Merge pull request #146 from swhitty/WSCloseCode

swhitty · web-flow · commit c11f6b780641 · 2025-03-05T07:33:54.000+11:00
Include reason within WSCloseCode
diff --git a/FlyingFox/Sources/WebSocket/WSCloseCode.swift b/FlyingFox/Sources/WebSocket/WSCloseCode.swift
@@ -31,35 +31,37 @@
 
 import Foundation
 
-public struct WSCloseCode: RawRepresentable, Sendable, Hashable {
-    public var rawValue: UInt16
-
-    public init(rawValue: UInt16) {
-        self.rawValue = rawValue
-    }
+public struct WSCloseCode: Sendable, Hashable {
+    public var code: UInt16
+    public var reason: String
 
     public init(_ code: UInt16) {
-        self.rawValue = code
+        self.code = code
+        self.reason = ""
+    }
+    public init(_ code: UInt16, reason: String) {
+        self.code = code
+        self.reason = reason
     }
 }
 
 public extension WSCloseCode {
     // The following codes are based on:
     // https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
 
-    static let normalClosure = WSCloseCode(1000)
-    static let goingAway = WSCloseCode(1001)
-    static let protocolError = WSCloseCode(1002)
-    static let unsupportedData = WSCloseCode(1003)
-    static let noStatusReceived = WSCloseCode(1005)
-    static let abnormalClosure = WSCloseCode(1006)
-    static let invalidFramePayloadData = WSCloseCode(1007)
-    static let policyViolation = WSCloseCode(1008)
-    static let messageTooBig = WSCloseCode(1009)
-    static let mandatoryExtensionMissing = WSCloseCode(1010)
-    static let internalServerError = WSCloseCode(1011)
-    static let serviceRestart = WSCloseCode(1012)
-    static let tryAgainLater = WSCloseCode(1013)
-    static let badGateway = WSCloseCode(1014)
-    static let tlsHandshakeFailure = WSCloseCode(1015)
+    static let normalClosure                = WSCloseCode(1000)
+    static let goingAway                    = WSCloseCode(1001, reason: "Going Away")
+    static let protocolError                = WSCloseCode(1002, reason: "Protocol Error")
+    static let unsupportedData              = WSCloseCode(1003, reason: "Unsupported Data")
+    static let noStatusReceived             = WSCloseCode(1005, reason: "No Status Received")
+    static let abnormalClosure              = WSCloseCode(1006, reason: "Abnormal Closure")
+    static let invalidFramePayload          = WSCloseCode(1007, reason: "Invalid Frame Payload")
+    static let policyViolation              = WSCloseCode(1008, reason: "Policy Violation")
+    static let messageTooBig                = WSCloseCode(1009, reason: "Message Too Big")
+    static let mandatoryExtensionMissing    = WSCloseCode(1010, reason: "Mandatory Extension Missing")
+    static let internalServerError          = WSCloseCode(1011, reason: "Internal Server Error")
+    static let serviceRestart               = WSCloseCode(1012, reason: "Service Restart")
+    static let tryAgainLater                = WSCloseCode(1013, reason: "Try Again Later")
+    static let badGateway                   = WSCloseCode(1014, reason: "Bad Gateway")
+    static let tlsHandshakeFailure          = WSCloseCode(1015, reason: "TLS Handshake Failure")
 }
diff --git a/FlyingFox/Sources/WebSocket/WSFrame.swift b/FlyingFox/Sources/WebSocket/WSFrame.swift
@@ -93,15 +93,14 @@ public struct WSFrame: Sendable, Hashable {
 public extension WSFrame {
     static func close(message: String = "", mask: Mask? = nil) -> Self {
         close(
-            code: message.isEmpty ? .normalClosure : .protocolError,
-            message: message,
+            code: message.isEmpty ? .normalClosure : WSCloseCode(WSCloseCode.protocolError.code, reason: message),
             mask: mask
         )
     }
 
-    static func close(code: WSCloseCode, message: String, mask: Mask? = nil) -> Self {
-        var payload = Data([UInt8(code.rawValue >> 8), UInt8(code.rawValue & 0xFF)])
-        if let data = message.data(using: .utf8) {
+    static func close(code: WSCloseCode, mask: Mask? = nil) -> Self {
+        var payload = Data([UInt8(code.code >> 8), UInt8(code.code & 0xFF)])
+        if let data = code.reason.data(using: .utf8) {
             payload.append(contentsOf: data)
         }
         return WSFrame(
diff --git a/FlyingFox/Sources/WebSocket/WSHandler.swift b/FlyingFox/Sources/WebSocket/WSHandler.swift
@@ -119,10 +119,18 @@ public struct MessageFrameWSHandler: WSHandler {
                 }
             }
             group.addTask {
-                for await message in messagesOut {
-                    for frame in makeFrames(for: message) {
-                        framesOut.yield(frame)
+                do {
+                    for await message in messagesOut {
+                        for frame in makeFrames(for: message) {
+                            framesOut.yield(frame)
+                            if frame.opcode == .close {
+                                throw FrameError.closed(frame)
+                            }
+                        }
                     }
+                    framesOut.finish(throwing: nil)
+                } catch {
+                    framesOut.finish(throwing: nil)
                 }
             }
             await group.next()!
@@ -140,23 +148,22 @@ public struct MessageFrameWSHandler: WSHandler {
         case .binary:
             return .data(frame.payload)
         case .close:
-            let (code, reason) = try makeCloseCode(from: frame.payload)
-            return .close(code: code, reason: reason)
+            return try .close(makeCloseCode(from: frame.payload))
         default:
             return nil
         }
     }
 
-    func makeCloseCode(from payload: Data) throws -> (WSCloseCode, String) {
+    func makeCloseCode(from payload: Data) throws -> WSCloseCode {
         guard payload.count >= 2 else {
-            return (.noStatusReceived, "")
+            return .noStatusReceived
         }
 
         let statusCode = 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 (WSCloseCode(statusCode), reason)
+        return WSCloseCode(statusCode, reason: reason)
     }
 
     func makeResponseFrames(for frame: WSFrame) throws -> WSFrame? {
@@ -178,8 +185,8 @@ public struct MessageFrameWSHandler: WSHandler {
             return Self.makeFrames(opcode: .text, payload: string.data(using: .utf8)!, size: frameSize)
         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)]
+        case let .close(code):
+            return [WSFrame.close(code: code)]
         }
     }
 
diff --git a/FlyingFox/Sources/WebSocket/WSMessage.swift b/FlyingFox/Sources/WebSocket/WSMessage.swift
@@ -34,7 +34,7 @@ import Foundation
 public enum WSMessage: @unchecked Sendable, Hashable {
     case text(String)
     case data(Data)
-    case close(code: WSCloseCode = .normalClosure, reason: String = "")
+    case close(WSCloseCode = .normalClosure)
 }
 
 public protocol WSMessageHandler: Sendable {
diff --git a/FlyingFox/Tests/WebSocket/WSFrameTests.swift b/FlyingFox/Tests/WebSocket/WSFrameTests.swift
@@ -70,15 +70,15 @@ struct WSFrameTests {
             )
         )
         #expect(
-            WSFrame.close(code: WSCloseCode(4999), message: "Err") == .make(
+            WSFrame.close(code: WSCloseCode(4999, reason: "Err")) == .make(
                 fin: true,
                 opcode: .close,
                 mask: nil,
                 payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
             )
         )
         #expect(
-            WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock) == .make(
+            WSFrame.close(code: WSCloseCode(4999, reason: "Err"), mask: .mock) == .make(
                 fin: true,
                 opcode: .close,
                 mask: .mock,
diff --git a/FlyingFox/Tests/WebSocket/WSHandlerTests.swift b/FlyingFox/Tests/WebSocket/WSHandlerTests.swift
@@ -63,11 +63,11 @@ struct WSHandlerTests {
 
         #expect(
             try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)) ==
-                .close(code: WSCloseCode(4999), reason: "fish")
+                .close(WSCloseCode(4999, reason: "fish"))
         )
         #expect(
             try handler.makeMessage(for: .make(fin: true, opcode: .close)) ==
-                .close(code: .noStatusReceived, reason: "")
+                .close(.noStatusReceived)
         )
     }
 
diff --git a/FlyingFox/XCTests/WebSocket/WSFrameTests.swift b/FlyingFox/XCTests/WebSocket/WSFrameTests.swift
@@ -65,7 +65,7 @@ final class WSFrameTests: XCTestCase {
                   payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")]))
         )
         XCTAssertEqual(
-            WSFrame.close(code: WSCloseCode(4999), message: "Err"),
+            WSFrame.close(code: WSCloseCode(4999, reason: "Err")),
             .make(
                 fin: true,
                 opcode: .close,
@@ -74,7 +74,7 @@ final class WSFrameTests: XCTestCase {
             )
         )
         XCTAssertEqual(
-            WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock),
+            WSFrame.close(code: WSCloseCode(4999, reason: "Err"), mask: .mock),
             .make(
                 fin: true,
                 opcode: .close,
diff --git a/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift b/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift
@@ -65,11 +65,11 @@ final class WSHandlerTests: XCTestCase {
 
         XCTAssertEqual(
             try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)),
-            .close(code: WSCloseCode(4999), reason: "fish")
+            .close(WSCloseCode(4999, reason: "fish"))
         )
         XCTAssertEqual(
             try handler.makeMessage(for: .make(fin: true, opcode: .close)),
-            .close(code: .noStatusReceived, reason: "")
+            .close(.noStatusReceived)
         )
     }
 

Original file line number	Diff line number	Diff line change
`@@ -93,15 +93,14 @@ public struct WSFrame: Sendable, Hashable {`
`93`	`93`	`public extension WSFrame {`
`94`	`94`	`static func close(message: String = "", mask: Mask? = nil) -> Self {`
`95`	`95`	`close(`
`96`		`- code: message.isEmpty ? .normalClosure : .protocolError,`
`97`		`- message: message,`
	`96`	`+ code: message.isEmpty ? .normalClosure : WSCloseCode(WSCloseCode.protocolError.code, reason: message),`
`98`	`97`	`mask: mask`
`99`	`98`	`)`
`100`	`99`	`}`
`101`	`100`
`102`		`- static func close(code: WSCloseCode, message: String, mask: Mask? = nil) -> Self {`
`103`		`- var payload = Data([UInt8(code.rawValue >> 8), UInt8(code.rawValue & 0xFF)])`
`104`		`- if let data = message.data(using: .utf8) {`
	`101`	`+ static func close(code: WSCloseCode, mask: Mask? = nil) -> Self {`
	`102`	`+ var payload = Data([UInt8(code.code >> 8), UInt8(code.code & 0xFF)])`
	`103`	`+ if let data = code.reason.data(using: .utf8) {`
`105`	`104`	`payload.append(contentsOf: data)`
`106`	`105`	`}`
`107`	`106`	`return WSFrame(`
Original file line number	Diff line number	Diff line change
`@@ -119,10 +119,18 @@ public struct MessageFrameWSHandler: WSHandler {`
`119`	`119`	`}`
`120`	`120`	`}`
`121`	`121`	`group.addTask {`
`122`		`- for await message in messagesOut {`
`123`		`- for frame in makeFrames(for: message) {`
`124`		`- framesOut.yield(frame)`
	`122`	`+ do {`
	`123`	`+ for await message in messagesOut {`
	`124`	`+ for frame in makeFrames(for: message) {`
	`125`	`+ framesOut.yield(frame)`
	`126`	`+ if frame.opcode == .close {`
	`127`	`+ throw FrameError.closed(frame)`
	`128`	`+ }`
	`129`	`+ }`
`125`	`130`	`}`
	`131`	`+ framesOut.finish(throwing: nil)`
	`132`	`+ } catch {`
	`133`	`+ framesOut.finish(throwing: nil)`
`126`	`134`	`}`
`127`	`135`	`}`
`128`	`136`	`await group.next()!`
`@@ -140,23 +148,22 @@ public struct MessageFrameWSHandler: WSHandler {`
`140`	`148`	`case .binary:`
`141`	`149`	`return .data(frame.payload)`
`142`	`150`	`case .close:`
`143`		`- let (code, reason) = try makeCloseCode(from: frame.payload)`
`144`		`- return .close(code: code, reason: reason)`
	`151`	`+ return try .close(makeCloseCode(from: frame.payload))`
`145`	`152`	`default:`
`146`	`153`	`return nil`
`147`	`154`	`}`
`148`	`155`	`}`
`149`	`156`
`150`		`- func makeCloseCode(from payload: Data) throws -> (WSCloseCode, String) {`
	`157`	`+ func makeCloseCode(from payload: Data) throws -> WSCloseCode {`
`151`	`158`	`guard payload.count >= 2 else {`
`152`		`- return (.noStatusReceived, "")`
	`159`	`+ return .noStatusReceived`
`153`	`160`	`}`
`154`	`161`
`155`	`162`	`let statusCode = payload.withUnsafeBytes { $0.load(as: UInt16.self).bigEndian }`
`156`	`163`	`guard let reason = String(data: payload.dropFirst(2), encoding: .utf8) else {`
`157`	`164`	`throw FrameError.invalid("Invalid UTF8 Sequence")`
`158`	`165`	`}`
`159`		`- return (WSCloseCode(statusCode), reason)`
	`166`	`+ return WSCloseCode(statusCode, reason: reason)`
`160`	`167`	`}`
`161`	`168`
`162`	`169`	`func makeResponseFrames(for frame: WSFrame) throws -> WSFrame? {`
`@@ -178,8 +185,8 @@ public struct MessageFrameWSHandler: WSHandler {`
`178`	`185`	`return Self.makeFrames(opcode: .text, payload: string.data(using: .utf8)!, size: frameSize)`
`179`	`186`	`case let .data(data):`
`180`	`187`	`return Self.makeFrames(opcode: .binary, payload: data, size: frameSize)`
`181`		`- case let .close(code: code, reason: message):`
`182`		`- return [WSFrame.close(code: code, message: message)]`
	`188`	`+ case let .close(code):`
	`189`	`+ return [WSFrame.close(code: code)]`
`183`	`190`	`}`
`184`	`191`	`}`
`185`	`192`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ import Foundation`
`34`	`34`	`public enum WSMessage: @unchecked Sendable, Hashable {`
`35`	`35`	`case text(String)`
`36`	`36`	`case data(Data)`
`37`		`- case close(code: WSCloseCode = .normalClosure, reason: String = "")`
	`37`	`+ case close(WSCloseCode = .normalClosure)`
`38`	`38`	`}`
`39`	`39`
`40`	`40`	`public protocol WSMessageHandler: Sendable {`
Original file line number	Diff line number	Diff line change
`@@ -70,15 +70,15 @@ struct WSFrameTests {`
`70`	`70`	`)`
`71`	`71`	`)`
`72`	`72`	`#expect(`
`73`		`- WSFrame.close(code: WSCloseCode(4999), message: "Err") == .make(`
	`73`	`+ WSFrame.close(code: WSCloseCode(4999, reason: "Err")) == .make(`
`74`	`74`	`fin: true,`
`75`	`75`	`opcode: .close,`
`76`	`76`	`mask: nil,`
`77`	`77`	`payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])`
`78`	`78`	`)`
`79`	`79`	`)`
`80`	`80`	`#expect(`
`81`		`- WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock) == .make(`
	`81`	`+ WSFrame.close(code: WSCloseCode(4999, reason: "Err"), mask: .mock) == .make(`
`82`	`82`	`fin: true,`
`83`	`83`	`opcode: .close,`
`84`	`84`	`mask: .mock,`
Original file line number	Diff line number	Diff line change
`@@ -63,11 +63,11 @@ struct WSHandlerTests {`
`63`	`63`
`64`	`64`	`#expect(`
`65`	`65`	`try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)) ==`
`66`		`- .close(code: WSCloseCode(4999), reason: "fish")`
	`66`	`+ .close(WSCloseCode(4999, reason: "fish"))`
`67`	`67`	`)`
`68`	`68`	`#expect(`
`69`	`69`	`try handler.makeMessage(for: .make(fin: true, opcode: .close)) ==`
`70`		`- .close(code: .noStatusReceived, reason: "")`
	`70`	`+ .close(.noStatusReceived)`
`71`	`71`	`)`
`72`	`72`	`}`
`73`	`73`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ final class WSFrameTests: XCTestCase {`
`65`	`65`	`payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")]))`
`66`	`66`	`)`
`67`	`67`	`XCTAssertEqual(`
`68`		`- WSFrame.close(code: WSCloseCode(4999), message: "Err"),`
	`68`	`+ WSFrame.close(code: WSCloseCode(4999, reason: "Err")),`
`69`	`69`	`.make(`
`70`	`70`	`fin: true,`
`71`	`71`	`opcode: .close,`
`@@ -74,7 +74,7 @@ final class WSFrameTests: XCTestCase {`
`74`	`74`	`)`
`75`	`75`	`)`
`76`	`76`	`XCTAssertEqual(`
`77`		`- WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock),`
	`77`	`+ WSFrame.close(code: WSCloseCode(4999, reason: "Err"), mask: .mock),`
`78`	`78`	`.make(`
`79`	`79`	`fin: true,`
`80`	`80`	`opcode: .close,`
Original file line number	Diff line number	Diff line change
`@@ -65,11 +65,11 @@ final class WSHandlerTests: XCTestCase {`
`65`	`65`
`66`	`66`	`XCTAssertEqual(`
`67`	`67`	`try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)),`
`68`		`- .close(code: WSCloseCode(4999), reason: "fish")`
	`68`	`+ .close(WSCloseCode(4999, reason: "fish"))`
`69`	`69`	`)`
`70`	`70`	`XCTAssertEqual(`
`71`	`71`	`try handler.makeMessage(for: .make(fin: true, opcode: .close)),`
`72`		`- .close(code: .noStatusReceived, reason: "")`
	`72`	`+ .close(.noStatusReceived)`
`73`	`73`	`)`
`74`	`74`	`}`
`75`	`75`