diff --git a/FlyingFox/Sources/HTTPConnection.swift b/FlyingFox/Sources/HTTPConnection.swift index 7f69d0e..62cd58b 100644 --- a/FlyingFox/Sources/HTTPConnection.swift +++ b/FlyingFox/Sources/HTTPConnection.swift @@ -84,15 +84,75 @@ struct HTTPConnection: Sendable { } func switchToWebSocket(with handler: some WSHandler, response: Data) async throws { + let (violations, violationsIn) = AsyncStream.makeStream() + // Reuse the connection-wide buffered stream so any bytes already // pulled past the upgrade request remain available to the WS framer. - let client = AsyncThrowingStream.decodingFrames(from: bytes) + let bytes = self.bytes + let client = AsyncThrowingStream { + do { + var frame = try await WSFrameEncoder.decodeFrame(from: bytes) + // RFC 6455 §5.1: "a client MUST mask all frames that it sends + // to the server. ... The server MUST close the connection upon + // receiving a frame that is not masked." + guard frame.mask != nil else { + violationsIn.yield(()) + return nil + } + // Handlers never see wire masks, so frames they echo back are + // safe to send unchanged. + frame.mask = nil + return frame + } catch SocketError.disconnected, is SequenceTerminationError { + return nil + } + } + let server = try await handler.makeFrames(for: client) try await socket.write(response) logger.logSwitchProtocol(self, to: "websocket") await requests.complete() - for await frame in server { - try await socket.write(WSFrameEncoder.encodeFrame(frame)) + try await withThrowingTaskGroup(of: Bool.self) { group in + group.addTask { + // Finishing `violations` on every exit — including a write + // error — guarantees the monitor task below always completes + // once output ends. + defer { violationsIn.finish() } + for await frame in server { + // RFC 6455 §5.1: "A server MUST NOT mask any frames that + // it sends to the client." + var frame = frame + frame.mask = nil + try await socket.write(WSFrameEncoder.encodeFrame(frame)) + } + return false + } + group.addTask { + for await _ in violations { + return true + } + return false + } + + var isViolation = try await group.next() ?? false + if isViolation { + // Stop and drain the output task before touching the socket + // so the close frame cannot interleave with another write. + group.cancelAll() + try? await group.waitForAll() + } else if let second = try await group.next() { + isViolation = second + } + if isViolation { + // RFC 6455 §5.1: a server "MAY send a Close frame with a + // status code of 1002 (protocol error)" and §7.1.7: "An + // endpoint SHOULD send a Close frame with an appropriate + // status code before closing the underlying connection." + // Throwing then fails the connection regardless of handler + // behaviour. + try? await socket.write(WSFrameEncoder.encodeFrame(.close(code: .protocolError))) + throw Error("Unmasked WebSocket frame received") + } } } diff --git a/FlyingFox/Sources/WebSocket/WSFrameEncoder.swift b/FlyingFox/Sources/WebSocket/WSFrameEncoder.swift index 07c574b..a5051da 100644 --- a/FlyingFox/Sources/WebSocket/WSFrameEncoder.swift +++ b/FlyingFox/Sources/WebSocket/WSFrameEncoder.swift @@ -70,6 +70,10 @@ struct WSFrameEncoder { static func decodeFrame(from bytes: some AsyncBufferedSequence) async throws -> WSFrame { var frame = try await decodeFrame(from: bytes.take()) let (length, mask) = try await decodeLengthMask(from: bytes) + // The payload is stored unmasked; the mask is preserved so servers can + // enforce RFC 6455 §5.1 — "a client MUST mask all frames that it sends + // to the server." + frame.mask = mask frame.payload = try await decodePayload(from: bytes, length: length, mask: mask) return frame } diff --git a/FlyingFox/Tests/HTTPConnectionTests.swift b/FlyingFox/Tests/HTTPConnectionTests.swift index 1b475ac..6b3a1b9 100644 --- a/FlyingFox/Tests/HTTPConnectionTests.swift +++ b/FlyingFox/Tests/HTTPConnectionTests.swift @@ -144,6 +144,175 @@ struct HTTPConnectionTests { HTTPConnection.makeIdentifier(from: .unix("/var/sock/fox")) == "/var/sock/fox" ) } + + @Test + func webSocket_UnmaskedClientFrame_FailsConnectionWithProtocolErrorClose() async throws { + // RFC 6455 §5.1: "The server MUST close the connection upon receiving + // a frame that is not masked. In this case, a server MAY send a Close + // frame with a status code of 1002 (protocol error)." + let (s1, s2) = try await AsyncSocket.makePair() + let connection = HTTPConnection(socket: s1) + + let response = Task { + try await connection.sendResponse(HTTPResponse(webSocket: MessageFrameWSHandler.make())) + } + + _ = try await s2.readResponse() + try await s2.writeFrame(.fish) + + #expect( + try await s2.readFrame() == .close(code: .protocolError) + ) + await #expect(throws: HTTPConnection.Error.self) { + try await response.value + } + + try s1.close() + try s2.close() + } + + @Test + func webSocket_UnmaskedClientFrame_FailsConnection_WhenHandlerSuppressesErrors() async throws { + // The connection owns RFC 6455 §5.1 termination: a handler that + // swallows input-stream failures and keeps its output open cannot + // keep the connection alive after an unmasked frame. + let (s1, s2) = try await AsyncSocket.makePair() + let connection = HTTPConnection(socket: s1) + + let response = Task { + try await connection.sendResponse(HTTPResponse(webSocket: ErrorSuppressingWSHandler())) + } + + _ = try await s2.readResponse() + try await s2.writeFrame(.fish) + + #expect( + try await s2.readFrame() == .close(code: .protocolError) + ) + await #expect(throws: HTTPConnection.Error.self) { + try await response.value + } + + try s1.close() + try s2.close() + } + + @Test + func webSocket_MaskedClientFrames_AreDeliveredToHandlerUnmasked() async throws { + // Wire masks are consumed at the connection boundary; handlers receive + // frames with `mask == nil` and the payload already unmasked. + let (s1, s2) = try await AsyncSocket.makePair() + let connection = HTTPConnection(socket: s1) + + let response = Task { + try await connection.sendResponse(HTTPResponse(webSocket: MaskReportingWSHandler())) + } + + _ = try await s2.readResponse() + try await s2.writeFrame(.fish.masked()) + + #expect( + try await s2.readFrame() == .make( + opcode: .binary, + payload: Data([1]) + "Fish".data(using: .utf8)! + ) + ) + + response.cancel() + try s1.close() + try s2.close() + } + + @Test + func webSocket_MaskedServerFrames_AreSentUnmasked() async throws { + // RFC 6455 §5.1: "A server MUST NOT mask any frames that it sends to + // the client." — even when a handler deliberately sets a mask. + let (s1, s2) = try await AsyncSocket.makePair() + let connection = HTTPConnection(socket: s1) + + let response = Task { + try await connection.sendResponse(HTTPResponse(webSocket: MaskedOutputWSHandler())) + } + + _ = try await s2.readResponse() + + #expect( + try await s2.readFrame() == .chips + ) + try await response.value + + try s1.close() + try s2.close() + } + + @Test + func webSocket_ClientDisconnect_EndsConnectionWithoutError() async throws { + // A peer that closes TCP without sending a Close frame ends the client + // stream (SocketError.disconnected → nil); the handler's output then + // finishes and the connection completes cleanly. + let (s1, s2) = try await AsyncSocket.makePair() + let connection = HTTPConnection(socket: s1) + + let response = Task { + try await connection.sendResponse(HTTPResponse(webSocket: MessageFrameWSHandler.make())) + } + + _ = try await s2.readResponse() + try s2.close() + + try await response.value + + try s1.close() + } +} + +private struct ErrorSuppressingWSHandler: WSHandler { + // Consumes client frames, swallows any input error, and never finishes + // its output stream. + func makeFrames(for client: AsyncThrowingStream) async throws -> AsyncStream { + AsyncStream { continuation in + let task = Task { + do { + for try await _ in client { } + } catch { } + // deliberately never calls continuation.finish() + } + continuation.onTermination = { _ in task.cancel() } + } + } +} + +private struct MaskReportingWSHandler: WSHandler { + // Echoes each frame as binary: first byte 1 when the received frame had + // no mask, followed by the received payload. + func makeFrames(for client: AsyncThrowingStream) async throws -> AsyncStream { + AsyncStream { continuation in + let task = Task { + do { + for try await frame in client { + continuation.yield( + WSFrame(fin: true, + opcode: .binary, + mask: nil, + payload: Data([frame.mask == nil ? 1 : 0]) + frame.payload) + ) + } + } catch { } + continuation.finish() + } + continuation.onTermination = { _ in task.cancel() } + } + } +} + +private struct MaskedOutputWSHandler: WSHandler { + // Ignores input and emits a single, deliberately masked frame. + func makeFrames(for client: AsyncThrowingStream) async throws -> AsyncStream { + AsyncStream { continuation in + continuation.yield(.chips.masked()) + continuation.finish() + } + } } private extension HTTPConnection { diff --git a/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift b/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift index 0045edc..f6f0787 100644 --- a/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift +++ b/FlyingFox/Tests/WebSocket/AsyncStream+WSFrameTests.swift @@ -46,6 +46,11 @@ struct WSFrameSequenceTests { #expect( try await AsyncThrowingStream.make([.close]).collectAll() == [.close] ) + // Decoding preserves the mask of a masked frame (RFC 6455 §5.1) while + // storing the payload unmasked. + #expect( + try await AsyncThrowingStream.make([.fish.masked()]).collectAll() == [.fish.masked()] + ) #expect( try await AsyncThrowingStream.make([]).collectAll() == [] ) diff --git a/FlyingFox/Tests/WebSocket/WSFrameTests.swift b/FlyingFox/Tests/WebSocket/WSFrameTests.swift index dfd2a9b..5316f1b 100644 --- a/FlyingFox/Tests/WebSocket/WSFrameTests.swift +++ b/FlyingFox/Tests/WebSocket/WSFrameTests.swift @@ -128,6 +128,14 @@ extension WSFrame { payload: text.data(using: .utf8)!) } + // Copy of the frame carrying a client masking key, as sent client → server. + // RFC 6455 §5.1: "a client MUST mask all frames that it sends to the server." + func masked(_ mask: Mask = .mock) -> Self { + var frame = self + frame.mask = mask + return frame + } + static func makeTextFrames(_ payload: String, maxCharacters: Int) -> [WSFrame] { var messages = payload.chunked(size: maxCharacters).enumerated().map { idx, substring in WSFrame.make(fin: false, isContinuation: idx != 0, text: String(substring))