Allow to specify the max websockets frame size when using the upgrader.

Sources/NIOWebSocket/WebSocketFrameDecoder.swift

@@ -221,7 +221,7 @@ public final class WebSocketFrameDecoder: ByteToMessageDecoder {
     public var cumulationBuffer: ByteBuffer? = nil
 
     /// The maximum frame size the decoder is willing to tolerate from the remote peer.
-    private let maxFrameSize: Int
+    /* private but tests */ let maxFrameSize: Int
 
     /// Our parser state.
     private var parser = WSParser()

Sources/NIOWebSocket/WebSocketUpgrader.swift

@@ -59,6 +59,7 @@ public final class WebSocketUpgrader: HTTPProtocolUpgrader {
 
     private let shouldUpgrade: (HTTPRequestHead) -> HTTPHeaders?
     private let upgradePipelineHandler: (Channel, HTTPRequestHead) -> EventLoopFuture<Void>
+    private let maxFrameSize: Int
 
     /// Create a new `WebSocketUpgrader`.
     ///
@@ -73,10 +74,44 @@ public final class WebSocketUpgrader: HTTPProtocolUpgrader {
     ///         websocket protocol. This only needs to add the user handlers: the
     ///         `WebSocketFrameEncoder` and `WebSocketFrameDecoder` will have been added to the
     ///         pipeline automatically.
-    public init(shouldUpgrade: @escaping (HTTPRequestHead) -> HTTPHeaders?,
+    ///     - maxFrameSize: The maximum frame size the decoder is willing to tolerate from the
+    ///         remote peer. WebSockets in principle allows frame sizes up to `2**64` bytes, but
+    ///         this is an objectively unreasonable maximum value (on AMD64 systems it is not
+    ///         possible to even allocate a buffer large enough to handle this size), so we
+    ///         set a lower one. The default value is the same as the default HTTP/2 max frame
+    ///         size, `2**14` bytes. Users may override this to any value up to `UInt32.max`.
+    ///         Users are strongly encouraged not to increase this value unless they absolutely
+    ///         must, as the decoder will not produce partial frames, meaning that it will hold
+    ///         on to data until the *entire* body is received.
+    public convenience init(shouldUpgrade: @escaping (HTTPRequestHead) -> HTTPHeaders?,
                 upgradePipelineHandler: @escaping (Channel, HTTPRequestHead) -> EventLoopFuture<Void>) {
+        self.init(maxFrameSize: 1 << 14, shouldUpgrade: shouldUpgrade, upgradePipelineHandler: upgradePipelineHandler)
+    }
+
+
+    /// Create a new `WebSocketUpgrader`.
+    ///
+    /// - parameters:
+    ///     - maxFrameSize: The maximum frame size the decoder is willing to tolerate from the
+    ///         remote peer. WebSockets in principle allows frame sizes up to `2**64` bytes, but
+    ///         this is an objectively unreasonable maximum value (on AMD64 systems it is not
+    ///         possible to even. Users may set this to any value up to `UInt32.max`.
+    ///     - shouldUpgrade: A callback that determines whether the websocket request should be
+    ///         upgraded. This callback is responsible for creating a `HTTPHeaders` object with
+    ///         any headers that it needs on the response *except for* the `Upgrade`, `Connection`,
+    ///         and `Sec-WebSocket-Accept` headers, which this upgrader will handle. Should return
+    ///         `nil` if the upgrade should be refused.
+    ///     - upgradePipelineHandler: A function that will be called once the upgrade response is
+    ///         flushed, and that is expected to mutate the `Channel` appropriately to handle the
+    ///         websocket protocol. This only needs to add the user handlers: the
+    ///         `WebSocketFrameEncoder` and `WebSocketFrameDecoder` will have been added to the
+    ///         pipeline automatically.
+    public init(maxFrameSize: Int, shouldUpgrade: @escaping (HTTPRequestHead) -> HTTPHeaders?,
+                upgradePipelineHandler: @escaping (Channel, HTTPRequestHead) -> EventLoopFuture<Void>) {
+        precondition(maxFrameSize <= UInt32.max, "invalid overlarge max frame size")
         self.shouldUpgrade = shouldUpgrade
         self.upgradePipelineHandler = upgradePipelineHandler
+        self.maxFrameSize = maxFrameSize
     }
 
     public func buildUpgradeResponse(upgradeRequest: HTTPRequestHead, initialResponseHeaders: HTTPHeaders) throws -> HTTPHeaders {
@@ -111,7 +146,7 @@ public final class WebSocketUpgrader: HTTPProtocolUpgrader {
 
     public func upgrade(ctx: ChannelHandlerContext, upgradeRequest: HTTPRequestHead) -> EventLoopFuture<Void> {
         return ctx.pipeline.add(handler: WebSocketFrameEncoder()).then {
-            ctx.pipeline.add(handler: WebSocketFrameDecoder())
+            ctx.pipeline.add(handler: WebSocketFrameDecoder(maxFrameSize: self.maxFrameSize))
         }.then {
             self.upgradePipelineHandler(ctx.channel, upgradeRequest)
         }

Tests/NIOWebSocketTests/EndToEndTests+XCTest.swift

@@ -34,6 +34,7 @@ extension EndToEndTests {
                 ("testUpgradeMayAddCustomHeaders", testUpgradeMayAddCustomHeaders),
                 ("testMayRegisterMultipleWebSocketEndpoints", testMayRegisterMultipleWebSocketEndpoints),
                 ("testSendAFewFrames", testSendAFewFrames),
+                ("testMaxFrameSize", testMaxFrameSize),
            ]
    }
 }

Tests/NIOWebSocketTests/EndToEndTests.swift

@@ -15,7 +15,7 @@
 import XCTest
 import NIO
 import NIOHTTP1
-import NIOWebSocket
+@testable import NIOWebSocket
 
 extension EmbeddedChannel {
     func readAllInboundBuffers() -> ByteBuffer {
@@ -340,4 +340,29 @@ class EndToEndTests: XCTestCase {
 
         XCTAssertEqual(recorder.frames, [dataFrame, pingFrame])
     }
+
+    func testMaxFrameSize() throws {
+        let basicUpgrader = WebSocketUpgrader(maxFrameSize: 16, shouldUpgrade: { head in HTTPHeaders() },
+                                              upgradePipelineHandler: { (channel, req) in
+            return channel.eventLoop.newSucceededFuture(result: ())
+        })
+        let (loop, server, client) = createTestFixtures(upgraders: [basicUpgrader])
+        defer {
+            XCTAssertNoThrow(try client.finish())
+            XCTAssertNoThrow(try server.finish())
+            XCTAssertNoThrow(try loop.syncShutdownGracefully())
+        }
+
+        let upgradeRequest = self.upgradeRequest(extraHeaders: ["Sec-WebSocket-Version": "13", "Sec-WebSocket-Key": "AQIDBAUGBwgJCgsMDQ4PEC=="])
+        XCTAssertNoThrow(try client.writeString(upgradeRequest).wait())
+        interactInMemory(client, server)
+
+        let receivedResponse = client.readAllInboundBuffers().allAsString()
+        assertResponseIs(response: receivedResponse,
+                         expectedResponseLine: "HTTP/1.1 101 Switching Protocols",
+                         expectedResponseHeaders: ["Upgrade: websocket", "Sec-WebSocket-Accept: OfS0wDaT5NoxF2gqm7Zj2YtetzM=", "Connection: upgrade"])
+
+        let decoder = (try server.pipeline.context(handlerType: WebSocketFrameDecoder.self).wait()).handler as! WebSocketFrameDecoder
+        XCTAssertEqual(16, decoder.maxFrameSize)
+    }
 }