WSMessage.close and WSCloseCode

FlyingFox/Sources/WebSocket/WSCloseCode.swift

@@ -0,0 +1,65 @@
+//
+//  WSCloseCode.swift
+//  FlyingFox
+//
+//  Created by Simon Whitty on 04/03/2025.
+//  Copyright © 2025 Simon Whitty. All rights reserved.
+//
+//  Distributed under the permissive MIT license
+//  Get the latest version from here:
+//
+//  https://github.com/swhitty/FlyingFox
+//
+//  Permission is hereby granted, free of charge, to any person obtaining a copy
+//  of this software and associated documentation files (the "Software"), to deal
+//  in the Software without restriction, including without limitation the rights
+//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+//  copies of the Software, and to permit persons to whom the Software is
+//  furnished to do so, subject to the following conditions:
+//
+//  The above copyright notice and this permission notice shall be included in all
+//  copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+//  SOFTWARE.
+//
+
+import Foundation
+
+public struct WSCloseCode: RawRepresentable, Sendable, Hashable {
+    public var rawValue: UInt16
+
+    public init(rawValue: UInt16) {
+        self.rawValue = rawValue
+    }
+
+    public init(_ code: UInt16) {
+        self.rawValue = code
+    }
+}
+
+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)
+}

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 ? .normalClosure : .protocolError,
+            message: 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) {
             payload.append(contentsOf: data)
         }
         return WSFrame(
             fin: true,
             opcode: .close,
             mask: mask,
-            payload: Data(payload)
+            payload: payload
         )
     }
 }

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 -> (WSCloseCode, String) {
+        guard payload.count >= 2 else {
+            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)
+    }
+
     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: WSCloseCode = .normalClosure, reason: String = "")
 }
 
 public protocol WSMessageHandler: Sendable {

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: WSCloseCode(4999), message: "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(
+                fin: true,
+                opcode: .close,
+                mask: .mock,
+                payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
+            )
+        )
     }
 }
 

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: WSCloseCode(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: .noStatusReceived, 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/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: WSCloseCode(4999), message: "Err"),
+            .make(
+                fin: true,
+                opcode: .close,
+                mask: nil,
+                payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
+            )
+        )
+        XCTAssertEqual(
+            WSFrame.close(code: WSCloseCode(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/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: WSCloseCode(4999), reason: "fish")
+        )
+        XCTAssertEqual(
+            try handler.makeMessage(for: .make(fin: true, opcode: .close)),
+            .close(code: .noStatusReceived, reason: "")
         )
     }
 
@@ -112,7 +123,7 @@ final class WSHandlerTests: XCTestCase {
 
         await AsyncAssertEqual(
             try await frames.collectAll(),
-            [.pong, .close(message: "Goodbye")]
+            [.pong, .close]
         )
     }