Merge pull request #126 from kateinoigakukun/katei/fix-ws-frame-length

Fix WebSocket frame length byte order for 16-bit length case
swhitty · Oct 20, 2024 · 211b297 · 211b297
2 parents ec371c0 + 63a52d8
commit 211b297
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/FlyingFox/Sources/WebSocket/WSFrameEncoder.swift b/FlyingFox/Sources/WebSocket/WSFrameEncoder.swift
@@ -138,8 +138,8 @@ struct WSFrameEncoder {
         case 0...125:
             return try await (Int(length0), hasMask ? decodeMask(from: bytes) : nil)
         case 126:
-            let length = try await UInt16(bytes.take()) |
-                                   UInt16(bytes.take()) << 8
+            let length = try await UInt16(bytes.take()) << 8 |
+                                   UInt16(bytes.take())
             return try await (Int(length), hasMask ? decodeMask(from: bytes) : nil)
         default:
             var length = try await UInt64(bytes.take())

diff --git a/FlyingFox/Tests/WebSocket/WSFrameEncoderTests.swift b/FlyingFox/Tests/WebSocket/WSFrameEncoderTests.swift
@@ -238,10 +238,10 @@ struct WSFrameEncoderTests {
             try await WSFrameEncoder.decodeLength(0x7D) == 125
         )
         #expect(
-            try await WSFrameEncoder.decodeLength(0x7E, 0x00, 0xFF) == 0xFF00
+            try await WSFrameEncoder.decodeLength(0x7E, 0x00, 0xFF) == 0x00FF
         )
         #expect(
-            try await WSFrameEncoder.decodeLength(0x7E, 0xFF, 0x00) == 0x00FF
+            try await WSFrameEncoder.decodeLength(0x7E, 0xFF, 0x00) == 0xFF00
         )
         #expect(
             try await WSFrameEncoder.decodeLength(0x7E, 0xFF, 0xFF) == 0xFFFF
@@ -371,6 +371,15 @@ struct WSFrameEncoderTests {
         frame = WSFrame.close(mask: .mock)
         try await socket.writeFrame(frame)
     }
+
+    @Test
+    func roundtripLength() async throws {
+        for length in 0..<256 {
+            let encoded = WSFrameEncoder.encodeLength(length, hasMask: false)
+            let decoded = try await WSFrameEncoder.decodeLengthMask(encoded)
+            #expect(length == decoded.length)
+        }
+    }
 }
 
 private extension WSFrameEncoder {

diff --git a/FlyingFox/XCTests/WebSocket/WSFrameEncoderTests.swift b/FlyingFox/XCTests/WebSocket/WSFrameEncoderTests.swift
@@ -209,11 +209,11 @@ final class WSFrameEncoderTests: XCTestCase {
             125
         )
         await AsyncAssertEqual(
-            try await WSFrameEncoder.decodeLength(0x7E, 0x00, 0xFF),
+            try await WSFrameEncoder.decodeLength(0x7E, 0xFF, 0x00),
             0xFF00
         )
         await AsyncAssertEqual(
-            try await WSFrameEncoder.decodeLength(0x7E, 0xFF, 0x00),
+            try await WSFrameEncoder.decodeLength(0x7E, 0x00, 0xFF),
             0x00FF
         )
         await AsyncAssertEqual(
@@ -346,6 +346,14 @@ final class WSFrameEncoderTests: XCTestCase {
         frame = WSFrame.close(mask: .mock)
         try await socket.writeFrame(frame)
     }
+
+    func testRoundtripLength() async throws {
+        for length in 0..<256 {
+            let encoded = WSFrameEncoder.encodeLength(length, hasMask: false)
+            let decoded = try await WSFrameEncoder.decodeLengthMask(encoded)
+            XCTAssertEqual(length, decoded.length)
+        }
+    }
 }
 
 private extension WSFrameEncoder {