Forward left over bytes to websocket handler (#14)

hummingbird-project · Feb 21, 2023 · 8cb4bf5 · 8cb4bf5
1 parent 33c60b6
commit 8cb4bf5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/Sources/HummingbirdWSClient/WebSocketClient.swift b/Sources/HummingbirdWSClient/WebSocketClient.swift
@@ -104,7 +104,7 @@ public enum HBWebSocketClient {
         )
 
         // add HTTP handler with web socket upgrade
-        return channel.pipeline.addHTTPClientHandlers(withClientUpgrade: config).flatMap {
+        return channel.pipeline.addHTTPClientHandlers(leftOverBytesStrategy: .forwardBytes, withClientUpgrade: config).flatMap {
             channel.pipeline.addHandler(httpHandler)
         }
     }

diff --git a/Sources/HummingbirdWSCore/HTTPServer+WebSocket.swift b/Sources/HummingbirdWSCore/HTTPServer+WebSocket.swift
@@ -32,8 +32,9 @@ extension HBHTTPServer {
             },
             upgradePipelineHandler: { (channel: Channel, head: HTTPRequestHead) in
                 let webSocket = HBWebSocket(channel: channel, type: .server)
-                onUpgrade(webSocket, head)
-                return channel.pipeline.addHandler(WebSocketHandler(webSocket: webSocket))
+                return channel.pipeline.addHandler(WebSocketHandler(webSocket: webSocket)).map { _ in
+                    onUpgrade(webSocket, head)
+                }
             }
         )
         self.httpChannelInitializer.addProtocolUpgrader(upgrader)

diff --git a/Tests/HummingbirdWebSocketTests/WebSocketTests.swift b/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
@@ -23,7 +23,7 @@ final class HummingbirdWebSocketTests: XCTestCase {
     static var eventLoopGroup: EventLoopGroup!
 
     override class func setUp() {
-        Self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: System.coreCount)
+        Self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
     }
 
     override class func tearDown() {
@@ -124,6 +124,24 @@ final class HummingbirdWebSocketTests: XCTestCase {
         XCTAssertTrue(clientHello)
     }
 
+    func testServerImmediateWrite() throws {
+        let promise = TimeoutPromise(eventLoop: Self.eventLoopGroup.next(), timeout: .seconds(60))
+        let app = try self.setupClientAndServer(
+            onServer: { ws in
+                ws.write(.text("hello"), promise: nil)
+            },
+            onClient: { ws in
+                ws.onRead { data, _ in
+                    XCTAssertEqual(data, .text("hello"))
+                    promise.succeed()
+                }
+            }
+        )
+        defer { app.stop() }
+
+        try promise.wait()
+    }
+
     /* Commented out as ws://echo.websocket.org is not working anymore
       func testClient() throws {
          let eventLoop = Self.eventLoopGroup.next()