From fa1a69fb01ffa69f71f126954855fd4dbbcf7473 Mon Sep 17 00:00:00 2001 From: Bridger Maxwell Date: Sun, 23 Feb 2020 12:58:35 -0500 Subject: [PATCH] Allow WebSocket maxFrameSize to be configured. --- Sources/Vapor/Response/Response.swift | 2 +- Sources/Vapor/Routing/RoutesBuilder+WebSocket.swift | 3 ++- Sources/Vapor/Server/HTTPServerUpgradeHandler.swift | 5 +++-- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/Sources/Vapor/Response/Response.swift b/Sources/Vapor/Response/Response.swift index 727b308931..4753e194db 100644 --- a/Sources/Vapor/Response/Response.swift +++ b/Sources/Vapor/Response/Response.swift @@ -33,7 +33,7 @@ public final class Response: CustomStringConvertible { } internal enum Upgrader { - case webSocket(onUpgrade: (WebSocket) -> ()) + case webSocket(maxFrameSize: Int?, onUpgrade: (WebSocket) -> ()) } internal var upgrader: Upgrader? diff --git a/Sources/Vapor/Routing/RoutesBuilder+WebSocket.swift b/Sources/Vapor/Routing/RoutesBuilder+WebSocket.swift index d6eb8d2c4a..32fd68751f 100644 --- a/Sources/Vapor/Routing/RoutesBuilder+WebSocket.swift +++ b/Sources/Vapor/Routing/RoutesBuilder+WebSocket.swift @@ -2,11 +2,12 @@ extension RoutesBuilder { @discardableResult public func webSocket( _ path: PathComponent..., + maxFrameSize: Int? = nil, onUpgrade: @escaping (Request, WebSocket) -> () ) -> Route { return self.on(.GET, path) { request -> Response in let res = Response(status: .switchingProtocols) - res.upgrader = .webSocket(onUpgrade: { ws in + res.upgrader = .webSocket(maxFrameSize: maxFrameSize, onUpgrade: { ws in onUpgrade(request, ws) }) return res diff --git a/Sources/Vapor/Server/HTTPServerUpgradeHandler.swift b/Sources/Vapor/Server/HTTPServerUpgradeHandler.swift index 2ea7713933..d6f51607e0 100644 --- a/Sources/Vapor/Server/HTTPServerUpgradeHandler.swift +++ b/Sources/Vapor/Server/HTTPServerUpgradeHandler.swift @@ -51,8 +51,9 @@ final class HTTPServerUpgradeHandler: ChannelDuplexHandler, RemovableChannelHand self.upgradeState = .upgraded if res.status == .switchingProtocols, let upgrader = res.upgrader { switch upgrader { - case .webSocket(let onUpgrade): - let webSocketUpgrader = NIOWebSocketServerUpgrader(automaticErrorHandling: false, shouldUpgrade: { channel, _ in + case .webSocket(let maxFrameSize, let onUpgrade): + let maxFrameSize = maxFrameSize ?? 1 << 14 + let webSocketUpgrader = NIOWebSocketServerUpgrader(maxFrameSize: maxFrameSize, automaticErrorHandling: false, shouldUpgrade: { channel, _ in return channel.eventLoop.makeSucceededFuture([:]) }, upgradePipelineHandler: { channel, req in return WebSocket.server(on: channel, onUpgrade: onUpgrade)