From 6ced1ac5eb57916ae20ebf2afa929112890fc278 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 9 Jul 2026 23:15:00 +0100 Subject: [PATCH] Bound remote MCP stdio lines (#31805) ## Why Remote MCP stdio buffers executor-provided stdout until a JSON-RPC line completes and stderr until a diagnostic line completes. A malicious or broken executor can keep sending chunks without a newline, making orchestrator memory grow without bound. Executor event boundaries are arbitrary, so a valid message before an oversized line must not be lost merely because both arrived in one event. ## What changed - Limit one stdout JSON-RPC line to 8 MiB. - Limit one stderr diagnostic line to 1 MiB. - Check each logical line before copying its bytes into the retained buffer. - Preserve complete stdout messages that precede an oversized line and deliver them before closing the stream. - Close the MCP transport on overflow; existing transport cleanup terminates the executor-managed process. - Preserve multiple lines per chunk, CRLF framing, and final unterminated messages at EOF. - Box the pending stdio transport so the added framing state does not inflate every `ClientState` value. - Add regression coverage for cross-chunk overflow, bounded lines in a larger chunk, and a valid prefix before an oversized line. ## Scope This bounds pre-newline buffering for executor-backed MCP stdio only. It does not change local stdio framing or address memory amplification while parsing a complete JSON message. --- .../src/executor_process_transport.rs | 101 ++++++++++++++++-- .../src/executor_process_transport_tests.rs | 72 ++++++++++++- codex-rs/rmcp-client/src/rmcp_client.rs | 8 +- 3 files changed, 165 insertions(+), 16 deletions(-) diff --git a/codex-rs/rmcp-client/src/executor_process_transport.rs b/codex-rs/rmcp-client/src/executor_process_transport.rs index cff4f517ed..23bd8c5e65 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport.rs @@ -46,18 +46,75 @@ use tracing::info; use tracing::warn; static PROCESS_COUNTER: AtomicUsize = AtomicUsize::new(1); +// Tool results can make valid MCP responses large, so keep the protocol +// ceiling well above ordinary messages while still bounding hostile input. +const MAX_MCP_STDOUT_LINE_BYTES: usize = 8 * 1024 * 1024; +// Stderr is diagnostic only and does not need the protocol stream's allowance. +const MAX_MCP_STDERR_LINE_BYTES: usize = 1024 * 1024; -#[derive(Default)] #[cfg_attr(test, derive(Debug, PartialEq, Eq))] struct LineBuffer { bytes: BytesMut, /// Prefix already scanned and known not to contain a newline. scanned_len: usize, + /// Bytes after the last buffered newline. + pending_line_bytes: usize, + max_line_bytes: usize, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct LineTooLong { + max_line_bytes: usize, +} + +impl Default for LineBuffer { + fn default() -> Self { + Self::new(MAX_MCP_STDOUT_LINE_BYTES) + } } impl LineBuffer { - fn extend_from_slice(&mut self, bytes: &[u8]) { - self.bytes.extend_from_slice(bytes); + fn new(max_line_bytes: usize) -> Self { + Self { + bytes: BytesMut::new(), + scanned_len: 0, + pending_line_bytes: 0, + max_line_bytes, + } + } + + fn extend_from_slice(&mut self, bytes: &[u8]) -> Result<(), LineTooLong> { + let mut remaining = bytes; + while let Some(newline_index) = memchr(b'\n', remaining) { + if newline_index > self.max_line_bytes.saturating_sub(self.pending_line_bytes) { + self.discard_pending_line(); + return Err(LineTooLong { + max_line_bytes: self.max_line_bytes, + }); + } + + let segment_len = newline_index + 1; + self.bytes.extend_from_slice(&remaining[..segment_len]); + self.pending_line_bytes = 0; + remaining = &remaining[segment_len..]; + } + if remaining.len() > self.max_line_bytes.saturating_sub(self.pending_line_bytes) { + self.discard_pending_line(); + return Err(LineTooLong { + max_line_bytes: self.max_line_bytes, + }); + } + + self.bytes.extend_from_slice(remaining); + self.pending_line_bytes += remaining.len(); + Ok(()) + } + + fn discard_pending_line(&mut self) { + let complete_line_bytes = self.bytes.len().saturating_sub(self.pending_line_bytes); + self.bytes.truncate(complete_line_bytes); + self.scanned_len = self.scanned_len.min(complete_line_bytes); + self.pending_line_bytes = 0; } fn take_line(&mut self) -> Option { @@ -79,8 +136,15 @@ impl LineBuffer { } self.scanned_len = 0; + self.pending_line_bytes = 0; Some(self.bytes.split()) } + + fn clear(&mut self) { + self.bytes = BytesMut::new(); + self.scanned_len = 0; + self.pending_line_bytes = 0; + } } // Remote public implementation. @@ -143,7 +207,7 @@ impl ExecutorProcessTransport { events, program_name, stdout: LineBuffer::default(), - stderr: LineBuffer::default(), + stderr: LineBuffer::new(MAX_MCP_STDERR_LINE_BYTES), closed: false, terminated: false, last_seq: 0, @@ -281,6 +345,9 @@ impl ExecutorProcessTransport { .map_err(io::Error::other)?; for chunk in response.chunks { self.push_process_output_if_new(chunk); + if self.closed { + return Ok(()); + } } self.last_seq = self.last_seq.max(response.next_seq.saturating_sub(1)); if let Some(message) = response.failure { @@ -309,16 +376,33 @@ impl ExecutorProcessTransport { // accepted defensively because the executor process API has a // unified stream enum, but remote MCP starts with `tty=false`. ExecOutputStream::Stdout | ExecOutputStream::Pty => { - self.stdout.extend_from_slice(&bytes); + if let Err(error) = self.stdout.extend_from_slice(&bytes) { + self.close_for_oversized_line("stdout", error); + } } // Stderr is intentionally out-of-band. It should help debug server // startup failures without entering rmcp framing. ExecOutputStream::Stderr => { - self.push_stderr(&bytes); + if let Err(error) = self.push_stderr(&bytes) { + self.stdout.clear(); + self.close_for_oversized_line("stderr", error); + } } } } + fn close_for_oversized_line(&mut self, stream_name: &str, error: LineTooLong) { + let max_line_bytes = error.max_line_bytes; + warn!( + "Remote MCP server {stream_name} line exceeds {max_line_bytes} bytes ({}); closing transport", + self.program_name + ); + self.stderr.clear(); + // Returning EOF makes rmcp drop the transport, whose Drop implementation + // terminates the executor-managed process. + self.closed = true; + } + fn take_stdout_message(&mut self, allow_partial: bool) -> Option> { // A normal MCP stdio server emits one JSON-RPC message per newline. // If the process has already closed, accept a final unterminated line @@ -343,10 +427,10 @@ impl ExecutorProcessTransport { } } - fn push_stderr(&mut self, bytes: &[u8]) { + fn push_stderr(&mut self, bytes: &[u8]) -> Result<(), LineTooLong> { // Keep stderr line-oriented in logs so a chatty MCP server does not // produce one log record per byte chunk. - self.stderr.extend_from_slice(bytes); + self.stderr.extend_from_slice(bytes)?; while let Some(line) = self.stderr.take_line() { let line = Self::trim_trailing_carriage_return(line); info!( @@ -355,6 +439,7 @@ impl ExecutorProcessTransport { String::from_utf8_lossy(&line) ); } + Ok(()) } fn flush_stderr(&mut self) { diff --git a/codex-rs/rmcp-client/src/executor_process_transport_tests.rs b/codex-rs/rmcp-client/src/executor_process_transport_tests.rs index aeef11a764..5a034f0272 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport_tests.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport_tests.rs @@ -2,32 +2,44 @@ use bytes::BytesMut; use pretty_assertions::assert_eq; use super::LineBuffer; +use super::LineTooLong; +use super::MAX_MCP_STDOUT_LINE_BYTES; #[test] fn searches_only_new_bytes_after_partial_line() { let mut buffer = LineBuffer::default(); - buffer.extend_from_slice(b"partial"); + buffer + .extend_from_slice(b"partial") + .expect("partial line should fit"); assert_eq!(buffer.take_line(), None); assert_eq!( buffer, LineBuffer { bytes: BytesMut::from(&b"partial"[..]), scanned_len: 7, + pending_line_bytes: 7, + max_line_bytes: MAX_MCP_STDOUT_LINE_BYTES, } ); - buffer.extend_from_slice(b" line"); + buffer + .extend_from_slice(b" line") + .expect("partial line should fit"); assert_eq!(buffer.take_line(), None); assert_eq!( buffer, LineBuffer { bytes: BytesMut::from(&b"partial line"[..]), scanned_len: 12, + pending_line_bytes: 12, + max_line_bytes: MAX_MCP_STDOUT_LINE_BYTES, } ); - buffer.extend_from_slice(b"\nnext"); + buffer + .extend_from_slice(b"\nnext") + .expect("completed line should fit"); assert_eq!( buffer.take_line(), Some(BytesMut::from(&b"partial line"[..])) @@ -37,6 +49,8 @@ fn searches_only_new_bytes_after_partial_line() { LineBuffer { bytes: BytesMut::from(&b"next"[..]), scanned_len: 0, + pending_line_bytes: 4, + max_line_bytes: MAX_MCP_STDOUT_LINE_BYTES, } ); } @@ -44,7 +58,9 @@ fn searches_only_new_bytes_after_partial_line() { #[test] fn splits_multiple_lines_and_retains_partial_tail() { let mut buffer = LineBuffer::default(); - buffer.extend_from_slice(b"first\nsecond\npartial"); + buffer + .extend_from_slice(b"first\nsecond\npartial") + .expect("lines should fit"); assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"first"[..]))); assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"second"[..]))); @@ -54,6 +70,8 @@ fn splits_multiple_lines_and_retains_partial_tail() { LineBuffer { bytes: BytesMut::from(&b"partial"[..]), scanned_len: 7, + pending_line_bytes: 7, + max_line_bytes: MAX_MCP_STDOUT_LINE_BYTES, } ); } @@ -61,7 +79,9 @@ fn splits_multiple_lines_and_retains_partial_tail() { #[test] fn takes_unterminated_remaining_bytes_at_eof() { let mut buffer = LineBuffer::default(); - buffer.extend_from_slice(b"remaining"); + buffer + .extend_from_slice(b"remaining") + .expect("remaining line should fit"); assert_eq!(buffer.take_line(), None); assert_eq!( @@ -70,3 +90,45 @@ fn takes_unterminated_remaining_bytes_at_eof() { ); assert_eq!(buffer, LineBuffer::default()); } + +#[test] +fn rejects_oversized_line_without_retaining_its_prefix() { + let mut buffer = LineBuffer::new(/*max_line_bytes*/ 5); + buffer + .extend_from_slice(b"12345") + .expect("line at the limit should fit"); + assert_eq!(buffer.take_line(), None); + + assert_eq!( + buffer.extend_from_slice(b"6"), + Err(LineTooLong { max_line_bytes: 5 }) + ); + assert_eq!(buffer, LineBuffer::new(/*max_line_bytes*/ 5)); +} + +#[test] +fn retains_complete_lines_before_an_oversized_line() { + let mut buffer = LineBuffer::new(/*max_line_bytes*/ 5); + + assert_eq!( + buffer.extend_from_slice(b"first\n123456"), + Err(LineTooLong { max_line_bytes: 5 }) + ); + + assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"first"[..]))); + assert_eq!(buffer.take_remaining(), None); +} + +#[test] +fn accepts_input_larger_than_limit_when_each_line_is_bounded() { + let mut buffer = LineBuffer::new(/*max_line_bytes*/ 5); + + buffer + .extend_from_slice(b"12345\nabcde\ntail") + .expect("each individual line should fit"); + + assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"12345"[..]))); + assert_eq!(buffer.take_line(), Some(BytesMut::from(&b"abcde"[..]))); + assert_eq!(buffer.take_line(), None); + assert_eq!(buffer.take_remaining(), Some(BytesMut::from(&b"tail"[..]))); +} diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 499c8db3bd..50251da13d 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -87,7 +87,7 @@ enum PendingTransport { transport: tokio::io::DuplexStream, }, Stdio { - transport: StdioServerTransport, + transport: Box, }, StreamableHttp { transport: StreamableHttpClientTransport, @@ -771,7 +771,9 @@ impl RmcpClient { } TransportRecipe::Stdio { command, launcher } => { let transport = launcher.launch(command.clone()).await?; - Ok(PendingTransport::Stdio { transport }) + Ok(PendingTransport::Stdio { + transport: Box::new(transport), + }) } TransportRecipe::StreamableHttp { server_name, @@ -890,7 +892,7 @@ impl RmcpClient { None, ), PendingTransport::Stdio { transport } => ( - service::serve_client(client_service, transport).boxed(), + service::serve_client(client_service, *transport).boxed(), None, ), PendingTransport::StreamableHttp { transport } => (