Add receivedPing to state machine

hummingbird-project · Nov 5, 2024 · b1a3bd3 · b1a3bd3
1 parent 299b341
commit b1a3bd3
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/Sources/HummingbirdWSCore/WebSocketHandler.swift b/Sources/HummingbirdWSCore/WebSocketHandler.swift
@@ -115,7 +115,7 @@ package actor WebSocketHandler {
                         if case .enabled = configuration.autoPing.value {
                             /// Add task sending ping frames every so often and verifying a pong frame was sent back
                             group.addTask {
-                                try await webSocketHandler.runPingLoop()
+                                try await webSocketHandler.runAutoPingLoop()
                                 return .init(closeCode: .goingAway, reason: "Ping timeout")
                             }
                         }
@@ -184,7 +184,7 @@ package actor WebSocketHandler {
         }
     }
 
-    func runPingLoop() async throws {
+    func runAutoPingLoop() async throws {
         let period = self.stateMachine.pingTimePeriod
         try await Task.sleep(for: period)
         loop: while true {
@@ -235,10 +235,16 @@ package actor WebSocketHandler {
 
     /// Respond to ping
     func onPing(_ frame: WebSocketFrame) async throws {
-        if frame.fin {
-            try await self.pong(data: frame.unmaskedData)
-        } else {
+        guard frame.fin else {
             try await self.close(code: .protocolError)
+            return
+        }
+        switch self.stateMachine.receivedPing(frameData: frame.unmaskedData) {
+        case .pong(let frameData):
+            try await self.pong(data: frameData)
+
+        case .doNothing:
+            break
         }
     }
 

diff --git a/Sources/HummingbirdWSCore/WebSocketStateMachine.swift b/Sources/HummingbirdWSCore/WebSocketStateMachine.swift
@@ -121,6 +121,24 @@ struct WebSocketStateMachine {
         }
     }
 
+    enum ReceivedPingResult {
+        case pong(ByteBuffer)
+        case doNothing
+    }
+
+    mutating func receivedPing(frameData: ByteBuffer) -> ReceivedPingResult {
+        switch self.state {
+        case .open:
+            return .pong(frameData)
+
+        case .closing:
+            return .pong(frameData)
+
+        case .closed:
+            return .doNothing
+        }
+    }
+
     mutating func receivedPong(frameData: ByteBuffer) {
         switch self.state {
         case .open(var state):

diff --git a/Tests/HummingbirdWebSocketTests/WebSocketTests.swift b/Tests/HummingbirdWebSocketTests/WebSocketTests.swift
@@ -524,11 +524,13 @@ final class HummingbirdWebSocketTests: XCTestCase {
             configuration: .init(address: .hostname("127.0.0.1", port: 0))
         )
         _ = try await application.test(.live) { client in
-            try await client.ws("/ws") { inbound, _, _ in
+            let frame = try await client.ws("/ws") { inbound, _, _ in
                 // don't handle any inbound data for a period much longer than the auto ping period
                 try await Task.sleep(for: .milliseconds(500))
                 for try await _ in inbound {}
             }
+            XCTAssertEqual(frame?.closeCode, .goingAway)
+            XCTAssertEqual(frame?.reason, "Ping timeout")
         }
     }