diff --git a/Sources/JSONRPCStdio/ProcessTransport.swift b/Sources/JSONRPCStdio/ProcessTransport.swift index 4ae9dcc..c3be9d1 100644 --- a/Sources/JSONRPCStdio/ProcessTransport.swift +++ b/Sources/JSONRPCStdio/ProcessTransport.swift @@ -125,7 +125,7 @@ public final class ProcessTransport: JSONRPCMessageTran continuation.yield(message) } }, - onEOF: { continuation.finish() }) + onFinish: { continuation.finish(throwing: $0) }) continuation.onTermination = { [weak self] _ in self?.close() diff --git a/Sources/JSONRPCSubprocess/StdioMessageTransport.swift b/Sources/JSONRPCSubprocess/StdioMessageTransport.swift index 6ab134e..5d086b9 100644 --- a/Sources/JSONRPCSubprocess/StdioMessageTransport.swift +++ b/Sources/JSONRPCSubprocess/StdioMessageTransport.swift @@ -142,7 +142,7 @@ public final class StdioTransport: JSONRPCMessageTransp var decoder = framing // value copy → fresh buffer for try await buffer in execution.standardOutput { let bytes = buffer.withUnsafeBytes { Array($0) } - for body in decoder.push(Data(bytes)) { + try decoder.push(Data(bytes)) { body in for message in (try? JSONRPCMessage.decodeMessages(from: body)) ?? [] { inbound.yield(message) } @@ -188,7 +188,7 @@ public final class StdioTransport: JSONRPCMessageTransp inbound.yield(message) } }, - onEOF: { inbound.finish() }) + onFinish: { inbound.finish(throwing: $0) }) // Writer: a single task drains outbound to our stdout (no lock needed). return Task { diff --git a/Sources/JSONRPCTCP/TCPClientTransport.swift b/Sources/JSONRPCTCP/TCPClientTransport.swift index 26a0590..8543d97 100644 --- a/Sources/JSONRPCTCP/TCPClientTransport.swift +++ b/Sources/JSONRPCTCP/TCPClientTransport.swift @@ -162,7 +162,7 @@ public final class TCPClientTransport: JSONRPCMessageTr continuation.yield(message) } }, - onEOF: { continuation.finish() }) + onFinish: { continuation.finish(throwing: $0) }) continuation.onTermination = { [weak self] _ in self?.close() diff --git a/Sources/JSONRPCWire/FramedReaderThread.swift b/Sources/JSONRPCWire/FramedReaderThread.swift index a817518..17d7629 100644 --- a/Sources/JSONRPCWire/FramedReaderThread.swift +++ b/Sources/JSONRPCWire/FramedReaderThread.swift @@ -23,18 +23,23 @@ package func startFramedReaderThread( framing: some MessageFraming, readChunk: @escaping @Sendable () -> Data, onBody: @escaping @Sendable (Data) -> Void, - onEOF: @escaping @Sendable () -> Void + onFinish: @escaping @Sendable ((any Error)?) -> Void ) { let thread = Thread { var decoder = framing while true { let chunk = readChunk() if chunk.isEmpty { break } - for body in decoder.push(chunk) { - onBody(body) + do { + try decoder.push(chunk) { onBody($0) } + } catch { + // A framing failure leaves no boundary to resynchronise on, so the + // read ends here and the caller fails its stream rather than looping. + onFinish(error) + return } } - onEOF() + onFinish(nil) } thread.name = name thread.stackSize = 4 << 20 diff --git a/Sources/JSONRPCWire/MessageFraming.swift b/Sources/JSONRPCWire/MessageFraming.swift index b416ea7..afd7cdd 100644 --- a/Sources/JSONRPCWire/MessageFraming.swift +++ b/Sources/JSONRPCWire/MessageFraming.swift @@ -14,9 +14,45 @@ import Foundation public protocol MessageFraming: Sendable { /// Wrap one message body for the wire (prepend a header / append a terminator). func frame(_ body: Data) -> Data - /// Feed newly-read bytes; return every complete message body they now yield - /// (header/terminator stripped), buffering any partial remainder. - mutating func push(_ bytes: Data) -> [Data] + /// Feed newly-read bytes, handing each complete message body (header/terminator + /// stripped) to `emit` as it is decoded, and buffering any partial remainder. + /// + /// Throws ``FramingError`` when the bytes cannot yield further messages — a peer + /// sending more than the configured limit, say. Such a stream cannot be + /// resynchronised, so a transport answers by finishing its inbound stream with the + /// error rather than reading on. + /// + /// Delivery is a callback rather than a return value precisely because of that + /// throw: one read can carry a complete message *and* an oversized one, and the + /// complete message has already been emitted by the time the failure is reported. + mutating func push(_ bytes: Data, emit: (Data) -> Void) throws +} + +extension MessageFraming { + /// Collects into an array instead of emitting. A failure discards whatever the same + /// call had already decoded, so transports should prefer the emitting form; this is + /// for callers that treat any framing failure as fatal. + public mutating func push(_ bytes: Data) throws -> [Data] { + var messages: [Data] = [] + try push(bytes) { messages.append($0) } + return messages + } +} + +/// Why a framing could not turn the bytes it was given into messages. +public enum FramingError: Error, Equatable, Sendable, CustomStringConvertible { + /// One message exceeded the framing's `maxBytes`. `pending` is how many bytes had + /// accumulated when the limit was passed — for an unterminated flood that is what + /// arrived before the buffer was dropped, not the message's real size, which is + /// unknowable. + case messageTooLarge(limit: Int, pending: Int) + + public var description: String { + switch self { + case .messageTooLarge(let limit, let pending): + return "Message exceeded the \(limit)-byte framing limit (\(pending) bytes buffered)." + } + } } /// LSP base-protocol framing: `Content-Length: \r\n\r\n`. @@ -26,8 +62,15 @@ public protocol MessageFraming: Sendable { public struct ContentLengthFraming: MessageFraming { private var buffer = Data() private var expectedLength: Int? + /// Largest single message to accept, in bytes; `0` (the default) is unlimited. + /// + /// A declared `Content-Length` is checked *before* its body is buffered, so an + /// oversized message costs nothing but its header. + public let maxBytes: Int - public init() {} + public init(maxBytes: Int = 0) { + self.maxBytes = maxBytes + } public func frame(_ body: Data) -> Data { var out = Data("Content-Length: \(body.count)\r\n\r\n".utf8) @@ -35,14 +78,29 @@ public struct ContentLengthFraming: MessageFraming { return out } - public mutating func push(_ bytes: Data) -> [Data] { + public mutating func push(_ bytes: Data, emit: (Data) -> Void) throws { buffer.append(bytes) - var messages: [Data] = [] - while let message = next() { messages.append(message) } - return messages + while let message = try next() { emit(message) } + // Headers with no separator in sight would otherwise buffer without bound. This + // is *not* `maxBytes`: that limits a message body, while a read can split + // anywhere — including part-way through a header longer than a small body limit. + if expectedLength == nil, buffer.count > Self.maxHeaderBytes { + throw drop(pending: buffer.count) + } + } + + /// How much unterminated header to tolerate. Generous next to any real header, and + /// independent of `maxBytes` so a small body limit never rejects a legal header that + /// a read happened to split. + private static let maxHeaderBytes = 8 * 1024 + + private mutating func drop(pending: Int) -> FramingError { + buffer.removeAll(keepingCapacity: false) + expectedLength = nil + return FramingError.messageTooLarge(limit: maxBytes, pending: pending) } - private mutating func next() -> Data? { + private mutating func next() throws -> Data? { if let length = expectedLength { guard buffer.count >= length else { return nil } let body = Data(buffer.prefix(length)) @@ -54,9 +112,11 @@ public struct ContentLengthFraming: MessageFraming { let headerBytes = buffer[buffer.startIndex ..< separator.lowerBound] let length = Self.contentLength(in: headerBytes) buffer.removeSubrange(buffer.startIndex ..< separator.upperBound) - guard let length else { return next() } + guard let length else { return try next() } + // Checked before the body arrives: the header already says how big it is. + if maxBytes > 0, length > maxBytes { throw drop(pending: length) } expectedLength = length - return next() + return try next() } private static let headerSeparator = Data("\r\n\r\n".utf8) @@ -82,8 +142,17 @@ public struct ContentLengthFraming: MessageFraming { /// Newline-delimited JSON framing (ACP and MCP-over-stdio): `\n`. public struct LineFraming: MessageFraming { private var buffer = Data() + /// Largest single message to accept, in bytes; `0` (the default) is unlimited. + /// + /// Newline framing cannot know a message's size in advance, so the limit is + /// applied twice: to a completed line, and to an unterminated remainder that has + /// already passed it — which is what stops a peer that never sends a newline from + /// growing the buffer without bound. + public let maxBytes: Int - public init() {} + public init(maxBytes: Int = 0) { + self.maxBytes = maxBytes + } public func frame(_ body: Data) -> Data { var out = body @@ -91,14 +160,22 @@ public struct LineFraming: MessageFraming { return out } - public mutating func push(_ bytes: Data) -> [Data] { + public mutating func push(_ bytes: Data, emit: (Data) -> Void) throws { buffer.append(bytes) - var messages: [Data] = [] while let newline = buffer.firstIndex(of: 0x0A) { let line = buffer[buffer.startIndex ..< newline] + let size = line.count buffer.removeSubrange(buffer.startIndex ... newline) - if !line.isEmpty { messages.append(Data(line)) } + // Emitted before the check on the *next* line, so a message that arrived in + // the same read as an oversized one is still delivered. + if maxBytes > 0, size > maxBytes { throw drop(pending: size) } + if !line.isEmpty { emit(Data(line)) } } - return messages + if maxBytes > 0, buffer.count > maxBytes { throw drop(pending: buffer.count) } + } + + private mutating func drop(pending: Int) -> FramingError { + buffer.removeAll(keepingCapacity: false) + return FramingError.messageTooLarge(limit: maxBytes, pending: pending) } } diff --git a/Tests/JSONRPCWireTests/JSONRPCWireTests.swift b/Tests/JSONRPCWireTests/JSONRPCWireTests.swift index 02d94de..e512e93 100644 --- a/Tests/JSONRPCWireTests/JSONRPCWireTests.swift +++ b/Tests/JSONRPCWireTests/JSONRPCWireTests.swift @@ -7,40 +7,40 @@ private func text(_ data: Data) -> String? { String(data: data, encoding: .utf8) // MARK: - ContentLengthFraming -@Test func contentLengthRoundTripsOneMessage() { +@Test func contentLengthRoundTripsOneMessage() throws { var framing = ContentLengthFraming() - let out = framing.push(framing.frame(body(#"{"jsonrpc":"2.0","id":1}"#))) + let out = try framing.push(framing.frame(body(#"{"jsonrpc":"2.0","id":1}"#))) #expect(out.count == 1) #expect(text(out[0]) == #"{"jsonrpc":"2.0","id":1}"#) } -@Test func contentLengthSplitsTwoMessagesInOneChunk() { +@Test func contentLengthSplitsTwoMessagesInOneChunk() throws { var framing = ContentLengthFraming() let chunk = framing.frame(body(#"{"a":1}"#)) + framing.frame(body(#"{"b":2}"#)) - let out = framing.push(chunk) + let out = try framing.push(chunk) #expect(out.count == 2) #expect(text(out[1]) == #"{"b":2}"#) } -@Test func contentLengthReassemblesAcrossChunks() { +@Test func contentLengthReassemblesAcrossChunks() throws { var framing = ContentLengthFraming() var emitted: [Data] = [] - for byte in framing.frame(body(#"{"hello":"world"}"#)) { emitted += framing.push(Data([byte])) } + for byte in framing.frame(body(#"{"hello":"world"}"#)) { emitted += try framing.push(Data([byte])) } #expect(emitted.count == 1) #expect(text(emitted[0]) == #"{"hello":"world"}"#) } -@Test func contentLengthCountsBytesNotCharacters() { +@Test func contentLengthCountsBytesNotCharacters() throws { var framing = ContentLengthFraming() - let out = framing.push(framing.frame(body(#"{"v":"café"}"#))) // 5 UTF-8 bytes, 4 chars + let out = try framing.push(framing.frame(body(#"{"v":"café"}"#))) // 5 UTF-8 bytes, 4 chars #expect(out.count == 1) #expect(text(out[0]) == #"{"v":"café"}"#) } -@Test func contentLengthRejectsNegativeLength() { +@Test func contentLengthRejectsNegativeLength() throws { // A malformed `Content-Length: -1` must be dropped, not used as a frame size. var framing = ContentLengthFraming() - let out = framing.push(Data("Content-Length: -1\r\n\r\n{}".utf8)) + let out = try framing.push(Data("Content-Length: -1\r\n\r\n{}".utf8)) #expect(out.isEmpty) } @@ -50,74 +50,171 @@ private func text(_ data: Data) -> String? { String(data: data, encoding: .utf8) #expect(LineFraming().frame(body(#"{"id":1}"#)).last == 0x0A) } -@Test func lineFramingSplitsMultipleLines() { +@Test func lineFramingSplitsMultipleLines() throws { var framing = LineFraming() let chunk = framing.frame(body(#"{"a":1}"#)) + framing.frame(body(#"{"b":2}"#)) - let out = framing.push(chunk) + let out = try framing.push(chunk) #expect(out.count == 2) #expect(text(out[0]) == #"{"a":1}"#) } -@Test func lineFramingReassemblesAcrossChunks() { +@Test func lineFramingReassemblesAcrossChunks() throws { var framing = LineFraming() var emitted: [Data] = [] - for byte in framing.frame(body(#"{"x":42}"#)) { emitted += framing.push(Data([byte])) } + for byte in framing.frame(body(#"{"x":42}"#)) { emitted += try framing.push(Data([byte])) } #expect(emitted.count == 1) #expect(text(emitted[0]) == #"{"x":42}"#) } // MARK: - SSEEventDecoder -@Test func sseDecodesOneEvent() { +@Test func sseDecodesOneEvent() throws { var decoder = SSEEventDecoder() - let out = decoder.push(body("data: {\"id\":1}\n\n")) + let out = try decoder.push(body("data: {\"id\":1}\n\n")) #expect(out.count == 1) #expect(text(out[0]) == "{\"id\":1}") } -@Test func sseIgnoresCommentsAndNonDataFields() { +@Test func sseIgnoresCommentsAndNonDataFields() throws { var decoder = SSEEventDecoder() - let out = decoder.push(body(": keep-alive\nevent: message\nid: 7\ndata: {\"x\":1}\n\n")) + let out = try decoder.push(body(": keep-alive\nevent: message\nid: 7\ndata: {\"x\":1}\n\n")) #expect(out.count == 1) #expect(text(out[0]) == "{\"x\":1}") } -@Test func sseDecodesTwoEventsInOneChunk() { +@Test func sseDecodesTwoEventsInOneChunk() throws { var decoder = SSEEventDecoder() - let out = decoder.push(body("data: {\"a\":1}\n\ndata: {\"b\":2}\n\n")) + let out = try decoder.push(body("data: {\"a\":1}\n\ndata: {\"b\":2}\n\n")) #expect(out.count == 2) #expect(text(out[1]) == "{\"b\":2}") } -@Test func sseToleratesCRLF() { +@Test func sseToleratesCRLF() throws { var decoder = SSEEventDecoder() - let out = decoder.push(body("data: {\"id\":5}\r\n\r\n")) + let out = try decoder.push(body("data: {\"id\":5}\r\n\r\n")) #expect(out.count == 1) #expect(text(out[0]) == "{\"id\":5}") } -@Test func sseJoinsMultipleDataLines() { +@Test func sseJoinsMultipleDataLines() throws { // Doc-promised: multiple `data:` lines of one event join with "\n". var decoder = SSEEventDecoder() - let out = decoder.push(body("data: {\"a\":\ndata: 1}\n\n")) + let out = try decoder.push(body("data: {\"a\":\ndata: 1}\n\n")) #expect(out.count == 1) #expect(text(out[0]) == "{\"a\":\n1}") } -@Test func sseReassemblesAcrossChunks() { +@Test func sseReassemblesAcrossChunks() throws { // The SSE path is the one that actually sees arbitrary network chunking. var decoder = SSEEventDecoder() var emitted: [Data] = [] - for byte in body("data: {\"id\":9}\n\n") { emitted += decoder.push(Data([byte])) } + for byte in body("data: {\"id\":9}\n\n") { emitted += try decoder.push(Data([byte])) } #expect(emitted.count == 1) #expect(text(emitted[0]) == "{\"id\":9}") } -@Test func sseTreatsBareFieldNameAsEmptyValue() { +@Test func sseTreatsBareFieldNameAsEmptyValue() throws { // Doc-promised: a line with no colon is a field name with an empty value, so // a bare `data` line contributes an empty payload — still a dispatched event. var decoder = SSEEventDecoder() - let out = decoder.push(body("data\n\n")) + let out = try decoder.push(body("data\n\n")) #expect(out.count == 1) #expect(text(out[0]) == "") } + +// MARK: - Byte limits + +@Test func lineFramingIsUnlimitedByDefault() throws { + var framing = LineFraming() + let huge = body(String(repeating: "x", count: 1 << 20)) + #expect(try framing.push(framing.frame(huge)).count == 1) +} + +@Test func lineFramingRejectsAnOversizedMessage() throws { + var framing = LineFraming(maxBytes: 16) + #expect(throws: FramingError.messageTooLarge(limit: 16, pending: 32)) { + try framing.push(framing.frame(body(String(repeating: "x", count: 32)))) + } + // The buffer is dropped, so the framing does not keep answering with the failure. + #expect(try framing.push(framing.frame(body("{}"))).count == 1) +} + +/// The case a limit exists for: a peer that never terminates its line would otherwise +/// grow the buffer without bound. +@Test func lineFramingRejectsAnUnterminatedFlood() throws { + var framing = LineFraming(maxBytes: 8) + #expect(try framing.push(body("12345")).isEmpty) + #expect(throws: FramingError.messageTooLarge(limit: 8, pending: 10)) { + try framing.push(body("67890")) + } +} + +@Test func lineFramingAcceptsAMessageExactlyAtTheLimit() throws { + var framing = LineFraming(maxBytes: 4) + let out = try framing.push(framing.frame(body("abcd"))) + #expect(text(out.first ?? Data()) == "abcd") +} + +@Test func contentLengthRejectsByTheDeclaredLengthBeforeTheBody() throws { + var framing = ContentLengthFraming(maxBytes: 16) + // Only the header is fed: the length alone is enough to refuse it. + #expect(throws: FramingError.messageTooLarge(limit: 16, pending: 4096)) { + try framing.push(body("Content-Length: 4096\r\n\r\n")) + } +} + +/// Headers are capped on their own, generous limit — not on `maxBytes`, which governs +/// a body — so this takes far more than a small body limit to trip. +@Test func contentLengthRejectsHeadersThatNeverEnd() throws { + var framing = ContentLengthFraming(maxBytes: 8) + // Well under the header cap: still fine, however small the body limit is. + #expect(try framing.push(body(String(repeating: "X-Pad: 1\r\n", count: 4))).isEmpty) + // Past it: a peer that never sends a separator cannot buffer without bound. + #expect(throws: (any Error).self) { + try framing.push(body(String(repeating: "X-Pad: 1\r\n", count: 1024))) + } +} + +@Test func contentLengthIsUnlimitedByDefault() throws { + var framing = ContentLengthFraming() + let big = body(String(repeating: "y", count: 1 << 16)) + #expect(try framing.push(framing.frame(big)).count == 1) +} + +/// A read can carry a good message and an oversized one together. The good one is +/// already delivered when the failure is reported — losing it because of what followed +/// it in the same buffer would be a bug in the framing, not in the peer. +@Test func lineFramingEmitsCompletedMessagesBeforeFailing() { + var framing = LineFraming(maxBytes: 8) + var emitted: [String] = [] + let chunk = framing.frame(body("{\"a\":1}")) + framing.frame(body(String(repeating: "x", count: 32))) + + #expect(throws: FramingError.messageTooLarge(limit: 8, pending: 32)) { + try framing.push(chunk) { emitted.append(text($0) ?? "") } + } + #expect(emitted == ["{\"a\":1}"]) +} + +@Test func contentLengthEmitsCompletedMessagesBeforeFailing() { + var framing = ContentLengthFraming(maxBytes: 8) + var emitted: [String] = [] + let chunk = framing.frame(body("{\"a\":1}")) + framing.frame(body(String(repeating: "y", count: 64))) + + #expect(throws: FramingError.messageTooLarge(limit: 8, pending: 64)) { + try framing.push(chunk) { emitted.append(text($0) ?? "") } + } + #expect(emitted == ["{\"a\":1}"]) +} + +/// A transport read can split anywhere, including inside a header that is longer than a +/// small body limit. The body limit must not be applied to the header. +@Test func contentLengthAcceptsAHeaderSplitUnderASmallLimit() throws { + var framing = ContentLengthFraming(maxBytes: 16) + let frame = framing.frame(body("{}")) + let split = frame.count - 1 + + #expect(try framing.push(frame.prefix(split)).isEmpty) + let out = try framing.push(frame.suffix(from: split)) + #expect(out.count == 1) + #expect(text(out[0]) == "{}") +}