clear pingData before writing new data (#11)

* clear pingData before writing new data * Add test
hummingbird-project · Jan 5, 2025 · b15960a · b15960a
1 parent 06e5e25
commit b15960a
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/Sources/WSCore/WebSocketStateMachine.swift b/Sources/WSCore/WebSocketStateMachine.swift
@@ -118,6 +118,7 @@ struct WebSocketStateMachine {
             }
             // creating random payload
             let random = (0..<Self.pingDataSize).map { _ in UInt8.random(in: 0...255) }
+            state.pingData.clear()
             state.pingData.writeBytes(random)
             state.lastPingTime = .now
             self.state = .open(state)

diff --git a/Tests/WebSocketTests/WebSocketStateMachineTests.swift b/Tests/WebSocketTests/WebSocketStateMachineTests.swift
@@ -96,4 +96,34 @@ final class WebSocketStateMachineTests: XCTestCase {
             return
         }
     }
+
+    // Verify ping buffer size doesnt grow
+    func testPingBufferSize() async throws {
+        var stateMachine = WebSocketStateMachine(autoPingSetup: .enabled(timePeriod: .milliseconds(1)))
+        var currentBuffer = ByteBuffer()
+        var count = 0
+        while true {
+            switch stateMachine.sendPing() {
+            case .sendPing(let buffer):
+                XCTAssertEqual(buffer.readableBytes, 16)
+                currentBuffer = buffer
+                count += 1
+                if count > 4 {
+                    return
+                }
+
+            case .wait(let time):
+                try await Task.sleep(for: time)
+                stateMachine.receivedPong(frameData: currentBuffer)
+
+            case .closeConnection:
+                XCTFail("Should not timeout")
+                return
+
+            case .stop:
+                XCTFail("Should not stop")
+                return
+            }
+        }
+    }
 }