github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bound remote MCP stdio lines (#31805)

Browse Source 
## 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.
This commit is contained in:
jif
2026-07-09 23:15:00 +01:00
committed by GitHub
parent 8cf9a1b1f8
commit 6ced1ac5eb
3 changed files with 165 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/rmcp-client/src/executor_process_transport.rs
+93 -8
View File
@@ -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<BytesMut> {
@@ -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<RxJsonRpcMessage<RoleClient>> {
// 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) {
codex-rs/rmcp-client/src/executor_process_transport_tests.rs
+67 -5
View File
@@ -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"[..])));
}
codex-rs/rmcp-client/src/rmcp_client.rs
+5 -3
View File
@@ -87,7 +87,7 @@ enum PendingTransport {
transport: tokio::io::DuplexStream,
},
Stdio {
transport: StdioServerTransport,
transport: Box<StdioServerTransport>,
},
StreamableHttp {
transport: StreamableHttpClientTransport<StreamableHttpClientAdapter>,
@@ -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 } => (