Close timeout (#93)

* Add closeTimeout to websocket configuration * Add testCloseTimeout * Use swift-websocket 1.2.0
hummingbird-project · Dec 19, 2024 · 11b4b52 · 11b4b52
1 parent 8cb5af3
commit 11b4b52
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 1 deletion.
diff --git a/Package.swift b/Package.swift
@@ -16,7 +16,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/hummingbird-project/hummingbird.git", from: "2.5.0"),
-        .package(url: "https://github.com/hummingbird-project/swift-websocket.git", from: "1.0.0"),
+        .package(url: "https://github.com/hummingbird-project/swift-websocket.git", from: "1.2.0"),
         .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.22.0"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.5.0"),
     ],

diff --git a/Sources/HummingbirdWebSocket/WebSocketChannel.swift b/Sources/HummingbirdWebSocket/WebSocketChannel.swift
@@ -83,6 +83,7 @@ public struct HTTP1WebSocketUpgradeChannel: ServerChildChannel, HTTPChannelHandl
                                         configuration: .init(
                                             extensions: extensions,
                                             autoPing: configuration.autoPing,
+                                            closeTimeout: configuration.closeTimeout,
                                             validateUTF8: configuration.validateUTF8
                                         ),
                                         asyncChannel: asyncChannel,

diff --git a/Sources/HummingbirdWebSocket/WebSocketServerConfiguration.swift b/Sources/HummingbirdWebSocket/WebSocketServerConfiguration.swift
@@ -22,6 +22,8 @@ public struct WebSocketServerConfiguration: Sendable {
     public var extensions: [any WebSocketExtensionBuilder]
     /// Auto ping
     public var autoPing: AutoPingSetup
+    /// How long server should wait for close frame from client before timing out
+    public var closeTimeout: Duration
     /// Should we check text messages are valid UTF8
     public var validateUTF8: Bool
 
@@ -35,11 +37,13 @@ public struct WebSocketServerConfiguration: Sendable {
         maxFrameSize: Int = (1 << 14),
         extensions: [WebSocketExtensionFactory] = [],
         autoPing: AutoPingSetup = .enabled(timePeriod: .seconds(30)),
+        closeTimeout: Duration = .seconds(15),
         validateUTF8: Bool = false
     ) {
         self.maxFrameSize = maxFrameSize
         self.extensions = extensions.map { $0.build() }
         self.autoPing = autoPing
+        self.closeTimeout = closeTimeout
         self.validateUTF8 = validateUTF8
     }
 }
diff --git a/Tests/HummingbirdWebSocketTests/WebSocketTests.swift b/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
@@ -573,6 +573,25 @@ final class HummingbirdWebSocketTests: XCTestCase {
         }
     }
 
+    func testCloseTimeout() async throws {
+        let app = Application(
+            router: Router(),
+            server: .http1WebSocketUpgrade { _, _, _ in
+                .upgrade([:]) { _, _, _ in
+                    try await cancelWhenGracefulShutdown {
+                        try await Task.sleep(for: .seconds(15))
+                        XCTFail("Should not reach here")
+                    }
+                }
+            }
+        )
+        try await app.test(.live) { client in
+            _ = try await client.ws("/", configuration: .init(closeTimeout: .seconds(1))) { _, outbound, _ in
+                try await outbound.write(.text("Hello"))
+            }
+        }
+    }
+
     func testUnrecognisedOpcode() async throws {
         let app = Application(
             router: Router(),