Custom code on WebSocket close

Author: swhitty

FlyingFox/Sources/WebSocket/WSFrame.swift

@@ -91,16 +91,24 @@ public struct WSFrame: Sendable, Hashable {
 }
 
 public extension WSFrame {
-    static func close(message: String? = nil, mask: Mask? = nil) -> Self {
-        var payload = message == nil ? Data([0x03, 0xE8]) : Data([0x03, 0xEA])
-        if let data = message?.data(using: .utf8) {
+    static func close(message: String = "", mask: Mask? = nil) -> Self {
+        close(
+            code: message.isEmpty ? 1000 : 1002,
+            message: message,
+            mask: mask
+        )
+    }
+
+    static func close(code: UInt16, message: String, mask: Mask? = nil) -> Self {
+        var payload = Data([UInt8(code >> 8), UInt8(code & 0xFF)])
+        if let data = message.data(using: .utf8) {
             payload.append(contentsOf: data)
         }
         return WSFrame(
             fin: true,
             opcode: .close,
             mask: mask,
-            payload: Data(payload)
+            payload: payload
         )
     }
 }

FlyingFox/Tests/WebSocket/WSFrameTests.swift

@@ -61,7 +61,6 @@ struct WSFrameTests {
                 payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")])
             )
         )
-
         #expect(
             WSFrame.close(message: "Err", mask: .mock) == .make(
                 fin: true,
@@ -70,6 +69,22 @@ struct WSFrameTests {
                 payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")])
             )
         )
+        #expect(
+            WSFrame.close(code: 4999, message: "Err") == .make(
+                fin: true,
+                opcode: .close,
+                mask: nil,
+                payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
+            )
+        )
+        #expect(
+            WSFrame.close(code: 4999, message: "Err", mask: .mock) == .make(
+                fin: true,
+                opcode: .close,
+                mask: .mock,
+                payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
+            )
+        )
     }
 }
 

FlyingFox/XCTests/WebSocket/WSFrameTests.swift

@@ -57,14 +57,31 @@ final class WSFrameTests: XCTestCase {
                   mask: nil,
                   payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")]))
         )
-
         XCTAssertEqual(
             WSFrame.close(message: "Err", mask: .mock),
             .make(fin: true,
                   opcode: .close,
                   mask: .mock,
                   payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")]))
         )
+        XCTAssertEqual(
+            WSFrame.close(code: 4999, message: "Err"),
+            .make(
+                fin: true,
+                opcode: .close,
+                mask: nil,
+                payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
+            )
+        )
+        XCTAssertEqual(
+            WSFrame.close(code: 4999, message: "Err", mask: .mock),
+            .make(
+                fin: true,
+                opcode: .close,
+                mask: .mock,
+                payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
+            )
+        )
     }
 }