diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 9bfe82c35051..29c553b5d23e 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -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(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( + 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?;