Bound exec-server stdio JSON-RPC messages (#32123)

## Why

Newline-delimited stdio input could buffer an unterminated JSON-RPC message without a per-message limit. Apply the same 64 MiB ceiling used by the other exec-server transports.

## What changed

- Read stdio messages with bounded lookahead and disconnect when a message exceeds the limit.
- Preserve LF and CRLF framing, including messages whose payload is exactly at the limit.

## Testing

Added tests for accepting a limit-sized CRLF message and rejecting an unterminated overlong message.

GitOrigin-RevId: 73eaf883324a7777960725037cc5b9c34720084f
This commit is contained in:
jif
2026-07-10 10:21:26 +00:00
committed by copyberry
parent c094a58c9f
commit c4c21b68a8
+138 -12
View File
@@ -24,11 +24,15 @@ use tracing::debug;
use tracing::warn;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::io::BufWriter;
pub(crate) const CHANNEL_CAPACITY: usize = 128;
// Match the existing serialized JSON-RPC message ceiling used by Noise and
// WebSocket transports so stdio has the same per-message bound.
const MAX_STDIO_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024;
const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2);
#[cfg(test)]
pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25);
@@ -230,6 +234,24 @@ pub(crate) struct JsonRpcConnection {
impl JsonRpcConnection {
pub(crate) fn from_stdio<R, W>(reader: R, writer: W, connection_label: String) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
Self::from_stdio_with_max_message_len(
reader,
writer,
connection_label,
MAX_STDIO_JSONRPC_MESSAGE_LEN,
)
}
fn from_stdio_with_max_message_len<R, W>(
reader: R,
writer: W,
connection_label: String,
max_message_len: usize,
) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
@@ -241,11 +263,63 @@ impl JsonRpcConnection {
let reader_label = connection_label.clone();
let incoming_tx_for_reader = incoming_tx.clone();
let disconnected_tx_for_reader = disconnected_tx.clone();
// Read one byte past the payload limit so an unterminated oversized
// message fails promptly. A trailing CR gets one more byte of lookahead
// because it may be the first half of a valid CRLF terminator.
let read_limit = u64::try_from(max_message_len.saturating_add(1)).unwrap_or(u64::MAX);
let reader_task = tokio::spawn(async move {
let mut lines = BufReader::new(reader).lines();
let mut reader = BufReader::new(reader);
let mut line = String::new();
loop {
match lines.next_line().await {
Ok(Some(line)) => {
line.clear();
let read_result = (&mut reader).take(read_limit).read_line(&mut line).await;
match read_result {
Ok(0) => {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
/*reason*/ None,
)
.await;
break;
}
Ok(_) => {
if line.ends_with('\n') {
line.pop();
if line.ends_with('\r') {
line.pop();
}
} else if line.len() > max_message_len && line.ends_with('\r') {
match reader.read_u8().await {
Ok(b'\n') => {
line.pop();
}
Ok(_) => {}
Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => {}
Err(err) => {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
Some(format!(
"failed to read JSON-RPC message from {reader_label}: {err}"
)),
)
.await;
break;
}
}
}
if line.len() > max_message_len {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
Some(format!(
"JSON-RPC message from {reader_label} exceeds maximum length of {max_message_len} bytes"
)),
)
.await;
break;
}
if line.trim().is_empty() {
continue;
}
@@ -270,15 +344,6 @@ impl JsonRpcConnection {
}
}
}
Ok(None) => {
send_disconnected(
&incoming_tx_for_reader,
&disconnected_tx_for_reader,
/*reason*/ None,
)
.await;
break;
}
Err(err) => {
send_disconnected(
&incoming_tx_for_reader,
@@ -601,6 +666,7 @@ mod tests {
use codex_exec_server_protocol::RequestId;
use futures::channel::mpsc as futures_mpsc;
use futures::task::AtomicWaker;
use pretty_assertions::assert_eq;
use tokio::net::TcpListener;
use tokio::time::timeout;
use tokio_tungstenite::accept_async;
@@ -608,6 +674,66 @@ mod tests {
use super::*;
#[tokio::test]
async fn stdio_connection_accepts_message_at_size_limit() -> anyhow::Result<()> {
let message = test_jsonrpc_message();
let encoded = serde_json::to_string(&message)?;
let max_message_len = encoded.len();
let (reader, mut peer) = tokio::io::duplex(max_message_len.saturating_add(1));
let mut connection = JsonRpcConnection::from_stdio_with_max_message_len(
reader,
tokio::io::sink(),
"test stdio peer".to_string(),
max_message_len,
);
peer.write_all(encoded.as_bytes()).await?;
peer.write_all(b"\r\n").await?;
let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv())
.await?
.expect("stdio connection should report the message");
match event {
JsonRpcConnectionEvent::Message(actual) => assert_eq!(actual, message),
event => anyhow::bail!("expected JSON-RPC message, got {event:?}"),
}
drop(peer);
drop(connection);
Ok(())
}
#[tokio::test]
async fn stdio_connection_rejects_overlong_unterminated_message() -> anyhow::Result<()> {
let max_message_len: usize = 32;
let (reader, mut peer) = tokio::io::duplex(max_message_len.saturating_add(1));
let mut connection = JsonRpcConnection::from_stdio_with_max_message_len(
reader,
tokio::io::sink(),
"hostile stdio peer".to_string(),
max_message_len,
);
let overlong_message = vec![b'x'; max_message_len + 1];
peer.write_all(&overlong_message).await?;
let event = timeout(Duration::from_secs(1), connection.incoming_rx.recv())
.await?
.expect("stdio connection should report the framing violation");
match event {
JsonRpcConnectionEvent::Disconnected { reason } => assert_eq!(
reason,
Some(
"JSON-RPC message from hostile stdio peer exceeds maximum length of 32 bytes"
.to_string()
)
),
event => anyhow::bail!("expected stdio disconnect, got {event:?}"),
}
drop(peer);
drop(connection);
Ok(())
}
#[tokio::test]
async fn websocket_connection_sends_configured_ping() -> anyhow::Result<()> {
let (client_websocket, mut server_websocket) = websocket_pair().await?;