Reject unmasked client WebSocket frames (RFC 6455 §5.1)

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<Void>.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<WSFrame, any Swift.Error> {
+            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")
+            }
         }
     }
 

FlyingFox/Sources/WebSocket/WSFrameEncoder.swift

@@ -70,6 +70,10 @@ struct WSFrameEncoder {
     static func decodeFrame(from bytes: some AsyncBufferedSequence<UInt8>) 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
     }

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<WSFrame, any Error>) async throws -> AsyncStream<WSFrame> {
+        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<WSFrame, any Error>) async throws -> AsyncStream<WSFrame> {
+        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<WSFrame, any Error>) async throws -> AsyncStream<WSFrame> {
+        AsyncStream { continuation in
+            continuation.yield(.chips.masked())
+            continuation.finish()
+        }
+    }
 }
 
 private extension HTTPConnection {

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() == []
         )

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))