WSMessage.close

Author: swhitty

FlyingFox/Sources/WebSocket/WSHandler.swift

@@ -106,10 +106,13 @@ public struct MessageFrameWSHandler: WSHandler {
                         } else if let frame = try makeResponseFrames(for: frame) {
                             framesOut.yield(frame)
                         }
+                        if frame.opcode == .close {
+                            throw FrameError.closed(frame)
+                        }
                     }
                     framesOut.finish(throwing: nil)
-                } catch FrameError.closed {
-                    framesOut.yield(.close(message: "Goodbye"))
+                } catch FrameError.closed(let frame) {
+                    framesOut.yield(frame)
                     framesOut.finish(throwing: nil)
                 } catch {
                     framesOut.finish(throwing: error)
@@ -136,11 +139,26 @@ public struct MessageFrameWSHandler: WSHandler {
             return .text(string)
         case .binary:
             return .data(frame.payload)
+        case .close:
+            let (code, reason) = try makeCloseCode(from: frame.payload)
+            return .close(code: code, reason: reason)
         default:
             return nil
         }
     }
 
+    func makeCloseCode(from payload: Data) throws -> (UInt16, String) {
+        guard payload.count >= 2 else {
+            return (1005, "")
+        }
+
+        let code = payload.withUnsafeBytes { $0.load(as: UInt16.self).bigEndian }
+        guard let reason = String(data: payload.dropFirst(2), encoding: .utf8) else {
+            throw FrameError.invalid("Invalid UTF8 Sequence")
+        }
+        return (code, reason)
+    }
+
     func makeResponseFrames(for frame: WSFrame) throws -> WSFrame? {
         switch frame.opcode {
         case .ping:
@@ -149,19 +167,19 @@ public struct MessageFrameWSHandler: WSHandler {
             return response
         case .pong:
             return nil
-        case .close:
-            throw FrameError.closed
         default:
             throw FrameError.invalid("Unexpected Frame")
         }
     }
 
     func makeFrames(for message: WSMessage) -> [WSFrame] {
         switch message {
-        case .text(let string):
+        case let .text(string):
             return Self.makeFrames(opcode: .text, payload: string.data(using: .utf8)!, size: frameSize)
-        case .data(let data):
+        case let .data(data):
             return Self.makeFrames(opcode: .binary, payload: data, size: frameSize)
+        case let .close(code: code, reason: message):
+            return [WSFrame.close(code: code, message: message)]
         }
     }
 
@@ -179,7 +197,7 @@ public struct MessageFrameWSHandler: WSHandler {
 extension MessageFrameWSHandler {
 
     enum FrameError: Error {
-        case closed
+        case closed(WSFrame)
         case invalid(String)
     }
 }

FlyingFox/Sources/WebSocket/WSMessage.swift

@@ -34,6 +34,7 @@ import Foundation
 public enum WSMessage: @unchecked Sendable, Hashable {
     case text(String)
     case data(Data)
+    case close(code: UInt16 = 1000, reason: String = "")
 }
 
 public protocol WSMessageHandler: Sendable {

FlyingFox/Tests/WebSocket/WSHandlerTests.swift

@@ -45,19 +45,29 @@ struct WSHandlerTests {
         #expect(throws: (any Error).self) {
             try handler.makeMessage(for: .make(fin: true, opcode: .text, payload: Data([0x03, 0xE8])))
         }
-
         #expect(
             try handler.makeMessage(for: .make(fin: true, opcode: .binary, payload: Data([0x01, 0x02]))) == .data(Data([0x01, 0x02]))
         )
-
         #expect(
             try handler.makeMessage(for: .make(fin: true, opcode: .ping)) == nil
         )
         #expect(
             try handler.makeMessage(for: .make(fin: true, opcode: .pong)) == nil
         )
+    }
+
+    @Test
+    func frames_CreatesCloseMessage() throws {
+        let handler = MessageFrameWSHandler.make()
+        let payload = Data([0x13, 0x87, .ascii("f"), .ascii("i"), .ascii("s"), .ascii("h")])
+
+        #expect(
+            try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)) ==
+                .close(code: 4999, reason: "fish")
+        )
         #expect(
-            try handler.makeMessage(for: .make(fin: true, opcode: .close)) == nil
+            try handler.makeMessage(for: .make(fin: true, opcode: .close)) ==
+                .close(code: 1005, reason: "")
         )
     }
 
@@ -114,7 +124,7 @@ struct WSHandlerTests {
         )
 
         #expect(
-            try await frames.collectAll() == [.pong, .close(message: "Goodbye")]
+            try await frames.collectAll() == [.pong, .close]
         )
     }
 

FlyingFox/XCTests/WebSocket/WSHandlerTests.swift

@@ -57,8 +57,19 @@ final class WSHandlerTests: XCTestCase {
         XCTAssertNil(
             try handler.makeMessage(for: .make(fin: true, opcode: .pong))
         )
-        XCTAssertNil(
-            try handler.makeMessage(for: .make(fin: true, opcode: .close))
+    }
+
+    func testFrames_CreatesCloseMessage() throws {
+        let handler = MessageFrameWSHandler.make()
+        let payload = Data([0x13, 0x87, .ascii("f"), .ascii("i"), .ascii("s"), .ascii("h")])
+
+        XCTAssertEqual(
+            try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)),
+            .close(code: 4999, reason: "fish")
+        )
+        XCTAssertEqual(
+            try handler.makeMessage(for: .make(fin: true, opcode: .close)),
+            .close(code: 1005, reason: "")
         )
     }
 
@@ -112,7 +123,7 @@ final class WSHandlerTests: XCTestCase {
 
         await AsyncAssertEqual(
             try await frames.collectAll(),
-            [.pong, .close(message: "Goodbye")]
+            [.pong, .close]
         )
     }