From e004af787287390c8f5c6490b352c19afcc9fce3 Mon Sep 17 00:00:00 2001 From: Simon Whitty Date: Tue, 4 Mar 2025 20:16:22 +1100 Subject: [PATCH 1/3] Custom code on WebSocket close --- FlyingFox/Sources/WebSocket/WSFrame.swift | 16 ++++++++++++---- FlyingFox/Tests/WebSocket/WSFrameTests.swift | 17 ++++++++++++++++- .../XCTests/WebSocket/WSFrameTests.swift | 19 ++++++++++++++++++- 3 files changed, 46 insertions(+), 6 deletions(-) diff --git a/FlyingFox/Sources/WebSocket/WSFrame.swift b/FlyingFox/Sources/WebSocket/WSFrame.swift index 8af4ca70..d610b016 100644 --- a/FlyingFox/Sources/WebSocket/WSFrame.swift +++ b/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 ) } } diff --git a/FlyingFox/Tests/WebSocket/WSFrameTests.swift b/FlyingFox/Tests/WebSocket/WSFrameTests.swift index 24130604..a3c50f8b 100644 --- a/FlyingFox/Tests/WebSocket/WSFrameTests.swift +++ b/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")]) + ) + ) } } diff --git a/FlyingFox/XCTests/WebSocket/WSFrameTests.swift b/FlyingFox/XCTests/WebSocket/WSFrameTests.swift index f79bc539..378ccdc2 100644 --- a/FlyingFox/XCTests/WebSocket/WSFrameTests.swift +++ b/FlyingFox/XCTests/WebSocket/WSFrameTests.swift @@ -57,7 +57,6 @@ 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, @@ -65,6 +64,24 @@ final class WSFrameTests: XCTestCase { 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")]) + ) + ) } } From 561c77060f80689180418517685eb98f2fcf438c Mon Sep 17 00:00:00 2001 From: Simon Whitty Date: Tue, 4 Mar 2025 21:07:21 +1100 Subject: [PATCH 2/3] WSMessage.close --- FlyingFox/Sources/WebSocket/WSHandler.swift | 32 +++++++++++++++---- FlyingFox/Sources/WebSocket/WSMessage.swift | 1 + .../Tests/WebSocket/WSHandlerTests.swift | 18 ++++++++--- .../XCTests/WebSocket/WSHandlerTests.swift | 17 ++++++++-- 4 files changed, 54 insertions(+), 14 deletions(-) diff --git a/FlyingFox/Sources/WebSocket/WSHandler.swift b/FlyingFox/Sources/WebSocket/WSHandler.swift index cdef1c96..edba4afe 100644 --- a/FlyingFox/Sources/WebSocket/WSHandler.swift +++ b/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,8 +167,6 @@ public struct MessageFrameWSHandler: WSHandler { return response case .pong: return nil - case .close: - throw FrameError.closed default: throw FrameError.invalid("Unexpected Frame") } @@ -158,10 +174,12 @@ public struct MessageFrameWSHandler: WSHandler { 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) } } diff --git a/FlyingFox/Sources/WebSocket/WSMessage.swift b/FlyingFox/Sources/WebSocket/WSMessage.swift index 54ab3f06..4393909c 100644 --- a/FlyingFox/Sources/WebSocket/WSMessage.swift +++ b/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 { diff --git a/FlyingFox/Tests/WebSocket/WSHandlerTests.swift b/FlyingFox/Tests/WebSocket/WSHandlerTests.swift index 2e13b4a5..d2b92d35 100644 --- a/FlyingFox/Tests/WebSocket/WSHandlerTests.swift +++ b/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] ) } diff --git a/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift b/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift index aa3876b6..2a20df06 100644 --- a/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift +++ b/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] ) } From 92d43d4bc462e29375f443b2a9fde7f3ee51157b Mon Sep 17 00:00:00 2001 From: Simon Whitty Date: Tue, 4 Mar 2025 21:32:06 +1100 Subject: [PATCH 3/3] WSCloseCode --- FlyingFox/Sources/WebSocket/WSCloseCode.swift | 65 +++++++++++++++++++ FlyingFox/Sources/WebSocket/WSFrame.swift | 6 +- FlyingFox/Sources/WebSocket/WSHandler.swift | 8 +-- FlyingFox/Sources/WebSocket/WSMessage.swift | 2 +- FlyingFox/Tests/WebSocket/WSFrameTests.swift | 4 +- .../Tests/WebSocket/WSHandlerTests.swift | 4 +- .../XCTests/WebSocket/WSFrameTests.swift | 4 +- .../XCTests/WebSocket/WSHandlerTests.swift | 4 +- 8 files changed, 81 insertions(+), 16 deletions(-) create mode 100644 FlyingFox/Sources/WebSocket/WSCloseCode.swift diff --git a/FlyingFox/Sources/WebSocket/WSCloseCode.swift b/FlyingFox/Sources/WebSocket/WSCloseCode.swift new file mode 100644 index 00000000..6ad3076d --- /dev/null +++ b/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) +} diff --git a/FlyingFox/Sources/WebSocket/WSFrame.swift b/FlyingFox/Sources/WebSocket/WSFrame.swift index d610b016..f9e43ca9 100644 --- a/FlyingFox/Sources/WebSocket/WSFrame.swift +++ b/FlyingFox/Sources/WebSocket/WSFrame.swift @@ -93,14 +93,14 @@ public struct WSFrame: Sendable, Hashable { public extension WSFrame { static func close(message: String = "", mask: Mask? = nil) -> Self { close( - code: message.isEmpty ? 1000 : 1002, + code: message.isEmpty ? .normalClosure : .protocolError, message: message, mask: mask ) } - static func close(code: UInt16, message: String, mask: Mask? = nil) -> Self { - var payload = Data([UInt8(code >> 8), UInt8(code & 0xFF)]) + 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) } diff --git a/FlyingFox/Sources/WebSocket/WSHandler.swift b/FlyingFox/Sources/WebSocket/WSHandler.swift index edba4afe..fd9ee6d0 100644 --- a/FlyingFox/Sources/WebSocket/WSHandler.swift +++ b/FlyingFox/Sources/WebSocket/WSHandler.swift @@ -147,16 +147,16 @@ public struct MessageFrameWSHandler: WSHandler { } } - func makeCloseCode(from payload: Data) throws -> (UInt16, String) { + func makeCloseCode(from payload: Data) throws -> (WSCloseCode, String) { guard payload.count >= 2 else { - return (1005, "") + return (.noStatusReceived, "") } - let code = payload.withUnsafeBytes { $0.load(as: UInt16.self).bigEndian } + 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 (code, reason) + return (WSCloseCode(statusCode), reason) } func makeResponseFrames(for frame: WSFrame) throws -> WSFrame? { diff --git a/FlyingFox/Sources/WebSocket/WSMessage.swift b/FlyingFox/Sources/WebSocket/WSMessage.swift index 4393909c..3c93baaf 100644 --- 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: UInt16 = 1000, reason: String = "") + case close(code: WSCloseCode = .normalClosure, reason: String = "") } public protocol WSMessageHandler: Sendable { diff --git a/FlyingFox/Tests/WebSocket/WSFrameTests.swift b/FlyingFox/Tests/WebSocket/WSFrameTests.swift index a3c50f8b..9057780c 100644 --- a/FlyingFox/Tests/WebSocket/WSFrameTests.swift +++ b/FlyingFox/Tests/WebSocket/WSFrameTests.swift @@ -70,7 +70,7 @@ struct WSFrameTests { ) ) #expect( - WSFrame.close(code: 4999, message: "Err") == .make( + WSFrame.close(code: WSCloseCode(4999), message: "Err") == .make( fin: true, opcode: .close, mask: nil, @@ -78,7 +78,7 @@ struct WSFrameTests { ) ) #expect( - WSFrame.close(code: 4999, message: "Err", mask: .mock) == .make( + WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock) == .make( fin: true, opcode: .close, mask: .mock, diff --git a/FlyingFox/Tests/WebSocket/WSHandlerTests.swift b/FlyingFox/Tests/WebSocket/WSHandlerTests.swift index d2b92d35..949bb5b3 100644 --- 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: 4999, reason: "fish") + .close(code: WSCloseCode(4999), reason: "fish") ) #expect( try handler.makeMessage(for: .make(fin: true, opcode: .close)) == - .close(code: 1005, reason: "") + .close(code: .noStatusReceived, reason: "") ) } diff --git a/FlyingFox/XCTests/WebSocket/WSFrameTests.swift b/FlyingFox/XCTests/WebSocket/WSFrameTests.swift index 378ccdc2..42475f12 100644 --- 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: 4999, message: "Err"), + WSFrame.close(code: WSCloseCode(4999), message: "Err"), .make( fin: true, opcode: .close, @@ -74,7 +74,7 @@ final class WSFrameTests: XCTestCase { ) ) XCTAssertEqual( - WSFrame.close(code: 4999, message: "Err", mask: .mock), + WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock), .make( fin: true, opcode: .close, diff --git a/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift b/FlyingFox/XCTests/WebSocket/WSHandlerTests.swift index 2a20df06..a91674f2 100644 --- 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: 4999, reason: "fish") + .close(code: WSCloseCode(4999), reason: "fish") ) XCTAssertEqual( try handler.makeMessage(for: .make(fin: true, opcode: .close)), - .close(code: 1005, reason: "") + .close(code: .noStatusReceived, reason: "") ) }