Add support for supplying additional headers

hummingbird-project · Apr 13, 2023 · 01537b4 · 01537b4
1 parent 5f48b71
commit 01537b4
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/Sources/HummingbirdWSClient/WebSocketClient.swift b/Sources/HummingbirdWSClient/WebSocketClient.swift
@@ -29,7 +29,7 @@ public enum HBWebSocketClient {
     ///   - configuration: Configuration of connection
     ///   - eventLoop: eventLoop to run connection on
     /// - Returns: EventLoopFuture which will be fulfilled with `HBWebSocket` once connection is made
-    public static func connect(url: HBURL, configuration: Configuration, on eventLoop: EventLoop) -> EventLoopFuture<HBWebSocket> {
+    public static func connect(url: HBURL, headers: HTTPHeaders = [:], configuration: Configuration, on eventLoop: EventLoop) -> EventLoopFuture<HBWebSocket> {
         let wsPromise = eventLoop.makePromise(of: HBWebSocket.self)
         do {
             let url = try SplitURL(url: url)
@@ -38,7 +38,7 @@ public enum HBWebSocketClient {
                 .channelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
                 .channelOption(ChannelOptions.socket(IPPROTO_TCP, TCP_NODELAY), value: 1)
                 .channelInitializer { channel in
-                    return Self.setupChannelForWebsockets(url: url, channel: channel, wsPromise: wsPromise, on: eventLoop)
+                    return Self.setupChannelForWebsockets(url: url, headers: headers, channel: channel, wsPromise: wsPromise, on: eventLoop)
                 }
                 .connect(host: url.host, port: url.port)
                 .cascadeFailure(to: wsPromise)
@@ -66,6 +66,7 @@ public enum HBWebSocketClient {
     /// setup for channel for websocket. Create initial HTTP request and include upgrade for when it is successful
     static func setupChannelForWebsockets(
         url: SplitURL,
+        headers: HTTPHeaders,
         channel: Channel,
         wsPromise: EventLoopPromise<HBWebSocket>,
         on eventLoop: EventLoop
@@ -78,6 +79,7 @@ public enum HBWebSocketClient {
         do {
             httpHandler = try WebSocketInitialRequestHandler(
                 url: url,
+                headers: headers,
                 upgradePromise: upgradePromise
             )
         } catch {

diff --git a/Tests/HummingbirdWebSocketTests/WebSocketTests.swift b/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
@@ -277,6 +277,32 @@ final class HummingbirdWebSocketTests: XCTestCase {
         let wsFuture = HBWebSocketClient.connect(url: "ws://localhost:8080/test?connect", configuration: .init(), on: eventLoop)
         _ = try wsFuture.wait()
     }
+
+    func testAdditionalHeaders() throws {
+        let app = HBApplication(configuration: .init(address: .hostname(port: 8080)))
+        // add HTTP to WebSocket upgrade
+        app.ws.addUpgrade()
+        // on websocket connect.
+        app.ws.on(
+            "/test",
+            shouldUpgrade: { request in
+                guard request.headers["Sec-WebSocket-Extensions"].first == "foo" else { return request.failure(HBHTTPError(.badRequest)) }
+                return request.success(nil)
+            },
+            onUpgrade: { _, _ in }
+        )
+        try app.start()
+        defer { app.stop() }
+
+        let eventLoop = app.eventLoopGroup.next()
+        let wsFuture = HBWebSocketClient.connect(
+            url: "ws://localhost:8080/test",
+            headers: ["Sec-WebSocket-Extensions": "foo"],
+            configuration: .init(),
+            on: eventLoop
+        )
+        _ = try wsFuture.wait()
+    }
 }
 
 @available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *)