Bound exec-server stdio JSON-RPC messages

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<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,13 +666,74 @@ 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;
     use tokio_tungstenite::connect_async;
 
     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?;