Skip to content

Commit 7d98b0c

Browse files
swhittyswhitty
authored
Merge pull request #145 from swhitty/WSMessage.close
WSMessage.close and WSCloseCode
2 parents 0285d3e + 92d43d4 commit 7d98b0c

File tree

8 files changed

+165
-20
lines changed

8 files changed

+165
-20
lines changed

‎FlyingFox/Sources/WebSocket/WSCloseCode.swift‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
//
2+
// WSCloseCode.swift
3+
// FlyingFox
4+
//
5+
// Created by Simon Whitty on 04/03/2025.
6+
// Copyright © 2025 Simon Whitty. All rights reserved.
7+
//
8+
// Distributed under the permissive MIT license
9+
// Get the latest version from here:
10+
//
11+
// https://github.com/swhitty/FlyingFox
12+
//
13+
// Permission is hereby granted, free of charge, to any person obtaining a copy
14+
// of this software and associated documentation files (the "Software"), to deal
15+
// in the Software without restriction, including without limitation the rights
16+
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
17+
// copies of the Software, and to permit persons to whom the Software is
18+
// furnished to do so, subject to the following conditions:
19+
//
20+
// The above copyright notice and this permission notice shall be included in all
21+
// copies or substantial portions of the Software.
22+
//
23+
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
24+
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
25+
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
26+
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
27+
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
28+
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
29+
// SOFTWARE.
30+
//
31+
32+
import Foundation
33+
34+
public struct WSCloseCode: RawRepresentable, Sendable, Hashable {
35+
public var rawValue: UInt16
36+
37+
public init(rawValue: UInt16) {
38+
self.rawValue = rawValue
39+
}
40+
41+
public init(_ code: UInt16) {
42+
self.rawValue = code
43+
}
44+
}
45+
46+
public extension WSCloseCode {
47+
// The following codes are based on:
48+
// https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
49+
50+
static let normalClosure = WSCloseCode(1000)
51+
static let goingAway = WSCloseCode(1001)
52+
static let protocolError = WSCloseCode(1002)
53+
static let unsupportedData = WSCloseCode(1003)
54+
static let noStatusReceived = WSCloseCode(1005)
55+
static let abnormalClosure = WSCloseCode(1006)
56+
static let invalidFramePayloadData = WSCloseCode(1007)
57+
static let policyViolation = WSCloseCode(1008)
58+
static let messageTooBig = WSCloseCode(1009)
59+
static let mandatoryExtensionMissing = WSCloseCode(1010)
60+
static let internalServerError = WSCloseCode(1011)
61+
static let serviceRestart = WSCloseCode(1012)
62+
static let tryAgainLater = WSCloseCode(1013)
63+
static let badGateway = WSCloseCode(1014)
64+
static let tlsHandshakeFailure = WSCloseCode(1015)
65+
}

‎FlyingFox/Sources/WebSocket/WSFrame.swift‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -91,16 +91,24 @@ public struct WSFrame: Sendable, Hashable {
9191
}
9292

9393
public extension WSFrame {
94-
static func close(message: String? = nil, mask: Mask? = nil) -> Self {
95-
var payload = message == nil ? Data([0x03, 0xE8]) : Data([0x03, 0xEA])
96-
if let data = message?.data(using: .utf8) {
94+
static func close(message: String = "", mask: Mask? = nil) -> Self {
95+
close(
96+
code: message.isEmpty ? .normalClosure : .protocolError,
97+
message: message,
98+
mask: mask
99+
)
100+
}
101+
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) {
97105
payload.append(contentsOf: data)
98106
}
99107
return WSFrame(
100108
fin: true,
101109
opcode: .close,
102110
mask: mask,
103-
payload: Data(payload)
111+
payload: payload
104112
)
105113
}
106114
}

‎FlyingFox/Sources/WebSocket/WSHandler.swift‎

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -106,10 +106,13 @@ public struct MessageFrameWSHandler: WSHandler {
106106
} else if let frame = try makeResponseFrames(for: frame) {
107107
framesOut.yield(frame)
108108
}
109+
if frame.opcode == .close {
110+
throw FrameError.closed(frame)
111+
}
109112
}
110113
framesOut.finish(throwing: nil)
111-
} catch FrameError.closed {
112-
framesOut.yield(.close(message: "Goodbye"))
114+
} catch FrameError.closed(let frame) {
115+
framesOut.yield(frame)
113116
framesOut.finish(throwing: nil)
114117
} catch {
115118
framesOut.finish(throwing: error)
@@ -136,11 +139,26 @@ public struct MessageFrameWSHandler: WSHandler {
136139
return .text(string)
137140
case .binary:
138141
return .data(frame.payload)
142+
case .close:
143+
let (code, reason) = try makeCloseCode(from: frame.payload)
144+
return .close(code: code, reason: reason)
139145
default:
140146
return nil
141147
}
142148
}
143149

150+
func makeCloseCode(from payload: Data) throws -> (WSCloseCode, String) {
151+
guard payload.count >= 2 else {
152+
return (.noStatusReceived, "")
153+
}
154+
155+
let statusCode = payload.withUnsafeBytes { $0.load(as: UInt16.self).bigEndian }
156+
guard let reason = String(data: payload.dropFirst(2), encoding: .utf8) else {
157+
throw FrameError.invalid("Invalid UTF8 Sequence")
158+
}
159+
return (WSCloseCode(statusCode), reason)
160+
}
161+
144162
func makeResponseFrames(for frame: WSFrame) throws -> WSFrame? {
145163
switch frame.opcode {
146164
case .ping:
@@ -149,19 +167,19 @@ public struct MessageFrameWSHandler: WSHandler {
149167
return response
150168
case .pong:
151169
return nil
152-
case .close:
153-
throw FrameError.closed
154170
default:
155171
throw FrameError.invalid("Unexpected Frame")
156172
}
157173
}
158174

159175
func makeFrames(for message: WSMessage) -> [WSFrame] {
160176
switch message {
161-
case .text(let string):
177+
case let .text(string):
162178
return Self.makeFrames(opcode: .text, payload: string.data(using: .utf8)!, size: frameSize)
163-
case .data(let data):
179+
case let .data(data):
164180
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)]
165183
}
166184
}
167185

@@ -179,7 +197,7 @@ public struct MessageFrameWSHandler: WSHandler {
179197
extension MessageFrameWSHandler {
180198

181199
enum FrameError: Error {
182-
case closed
200+
case closed(WSFrame)
183201
case invalid(String)
184202
}
185203
}

‎FlyingFox/Sources/WebSocket/WSMessage.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import Foundation
3434
public enum WSMessage: @unchecked Sendable, Hashable {
3535
case text(String)
3636
case data(Data)
37+
case close(code: WSCloseCode = .normalClosure, reason: String = "")
3738
}
3839

3940
public protocol WSMessageHandler: Sendable {

‎FlyingFox/Tests/WebSocket/WSFrameTests.swift‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,6 @@ struct WSFrameTests {
6161
payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")])
6262
)
6363
)
64-
6564
#expect(
6665
WSFrame.close(message: "Err", mask: .mock) == .make(
6766
fin: true,
@@ -70,6 +69,22 @@ struct WSFrameTests {
7069
payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")])
7170
)
7271
)
72+
#expect(
73+
WSFrame.close(code: WSCloseCode(4999), message: "Err") == .make(
74+
fin: true,
75+
opcode: .close,
76+
mask: nil,
77+
payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
78+
)
79+
)
80+
#expect(
81+
WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock) == .make(
82+
fin: true,
83+
opcode: .close,
84+
mask: .mock,
85+
payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
86+
)
87+
)
7388
}
7489
}
7590

‎FlyingFox/Tests/WebSocket/WSHandlerTests.swift‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -45,19 +45,29 @@ struct WSHandlerTests {
4545
#expect(throws: (any Error).self) {
4646
try handler.makeMessage(for: .make(fin: true, opcode: .text, payload: Data([0x03, 0xE8])))
4747
}
48-
4948
#expect(
5049
try handler.makeMessage(for: .make(fin: true, opcode: .binary, payload: Data([0x01, 0x02]))) == .data(Data([0x01, 0x02]))
5150
)
52-
5351
#expect(
5452
try handler.makeMessage(for: .make(fin: true, opcode: .ping)) == nil
5553
)
5654
#expect(
5755
try handler.makeMessage(for: .make(fin: true, opcode: .pong)) == nil
5856
)
57+
}
58+
59+
@Test
60+
func frames_CreatesCloseMessage() throws {
61+
let handler = MessageFrameWSHandler.make()
62+
let payload = Data([0x13, 0x87, .ascii("f"), .ascii("i"), .ascii("s"), .ascii("h")])
63+
64+
#expect(
65+
try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)) ==
66+
.close(code: WSCloseCode(4999), reason: "fish")
67+
)
5968
#expect(
60-
try handler.makeMessage(for: .make(fin: true, opcode: .close)) == nil
69+
try handler.makeMessage(for: .make(fin: true, opcode: .close)) ==
70+
.close(code: .noStatusReceived, reason: "")
6171
)
6272
}
6373

@@ -114,7 +124,7 @@ struct WSHandlerTests {
114124
)
115125

116126
#expect(
117-
try await frames.collectAll() == [.pong, .close(message: "Goodbye")]
127+
try await frames.collectAll() == [.pong, .close]
118128
)
119129
}
120130

‎FlyingFox/XCTests/WebSocket/WSFrameTests.swift‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,14 +57,31 @@ final class WSFrameTests: XCTestCase {
5757
mask: nil,
5858
payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")]))
5959
)
60-
6160
XCTAssertEqual(
6261
WSFrame.close(message: "Err", mask: .mock),
6362
.make(fin: true,
6463
opcode: .close,
6564
mask: .mock,
6665
payload: Data([0x03, 0xEA, .ascii("E"), .ascii("r"), .ascii("r")]))
6766
)
67+
XCTAssertEqual(
68+
WSFrame.close(code: WSCloseCode(4999), message: "Err"),
69+
.make(
70+
fin: true,
71+
opcode: .close,
72+
mask: nil,
73+
payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
74+
)
75+
)
76+
XCTAssertEqual(
77+
WSFrame.close(code: WSCloseCode(4999), message: "Err", mask: .mock),
78+
.make(
79+
fin: true,
80+
opcode: .close,
81+
mask: .mock,
82+
payload: Data([0x13, 0x87, .ascii("E"), .ascii("r"), .ascii("r")])
83+
)
84+
)
6885
}
6986
}
7087

‎FlyingFox/XCTests/WebSocket/WSHandlerTests.swift‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,19 @@ final class WSHandlerTests: XCTestCase {
5757
XCTAssertNil(
5858
try handler.makeMessage(for: .make(fin: true, opcode: .pong))
5959
)
60-
XCTAssertNil(
61-
try handler.makeMessage(for: .make(fin: true, opcode: .close))
60+
}
61+
62+
func testFrames_CreatesCloseMessage() throws {
63+
let handler = MessageFrameWSHandler.make()
64+
let payload = Data([0x13, 0x87, .ascii("f"), .ascii("i"), .ascii("s"), .ascii("h")])
65+
66+
XCTAssertEqual(
67+
try handler.makeMessage(for: .make(fin: true, opcode: .close, payload: payload)),
68+
.close(code: WSCloseCode(4999), reason: "fish")
69+
)
70+
XCTAssertEqual(
71+
try handler.makeMessage(for: .make(fin: true, opcode: .close)),
72+
.close(code: .noStatusReceived, reason: "")
6273
)
6374
}
6475

@@ -112,7 +123,7 @@ final class WSHandlerTests: XCTestCase {
112123

113124
await AsyncAssertEqual(
114125
try await frames.collectAll(),
115-
[.pong, .close(message: "Goodbye")]
126+
[.pong, .close]
116127
)
117128
}
118129

0 commit comments

Comments
 (0)