Allow WebSocket maxFrameSize to be configured.

vapor · Feb 28, 2020 · fa1a69f · fa1a69f
1 parent 638dc58
commit fa1a69f
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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)