diff --git a/codex-rs/exec-server/src/client_inbound_request_limit.rs b/codex-rs/exec-server/src/client_inbound_request_limit.rs new file mode 100644 index 0000000000..5a462c93c3 --- /dev/null +++ b/codex-rs/exec-server/src/client_inbound_request_limit.rs @@ -0,0 +1,40 @@ +use codex_exec_server_protocol::EXEC_CLOSED_METHOD; +use codex_exec_server_protocol::EXEC_EXITED_METHOD; +use codex_exec_server_protocol::EXEC_OUTPUT_DELTA_METHOD; +use codex_exec_server_protocol::HTTP_REQUEST_BODY_DELTA_METHOD; +use codex_exec_server_protocol::JSONRPCMessage; +use codex_exec_server_protocol::MAX_HTTP_BODY_DELTA_BYTES; + +// A transport may materialize one larger frame before its JSON-RPC kind is known. +pub(crate) const MAX_CLIENT_INBOUND_REQUEST_LEN: usize = 8 * 1024; +// Streamed HTTP bodies carry up to 1 MiB before base64 and JSON-RPC framing. +pub(crate) const MAX_CLIENT_INBOUND_NOTIFICATION_LEN: usize = 2 * MAX_HTTP_BODY_DELTA_BYTES; + +pub(crate) fn client_inbound_message_exceeded_limit( + message: Result<&JSONRPCMessage, &serde_json::Error>, + encoded_len: usize, + max_request_len: usize, +) -> Option { + let max_len = match message { + Ok(JSONRPCMessage::Notification(notification)) + if matches!( + notification.method.as_str(), + EXEC_OUTPUT_DELTA_METHOD + | EXEC_EXITED_METHOD + | EXEC_CLOSED_METHOD + | HTTP_REQUEST_BODY_DELTA_METHOD + ) => + { + MAX_CLIENT_INBOUND_NOTIFICATION_LEN + } + Ok(JSONRPCMessage::Request(_)) | Ok(JSONRPCMessage::Notification(_)) | Err(_) => { + max_request_len + } + Ok(JSONRPCMessage::Response(_)) | Ok(JSONRPCMessage::Error(_)) => return None, + }; + (encoded_len > max_len).then_some(max_len) +} + +#[cfg(test)] +#[path = "client_inbound_request_limit_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/client_inbound_request_limit_tests.rs b/codex-rs/exec-server/src/client_inbound_request_limit_tests.rs new file mode 100644 index 0000000000..bc7e537f5d --- /dev/null +++ b/codex-rs/exec-server/src/client_inbound_request_limit_tests.rs @@ -0,0 +1,115 @@ +use codex_exec_server_protocol::EXEC_CLOSED_METHOD; +use codex_exec_server_protocol::EXEC_EXITED_METHOD; +use codex_exec_server_protocol::EXEC_OUTPUT_DELTA_METHOD; +use codex_exec_server_protocol::HTTP_REQUEST_BODY_DELTA_METHOD; +use codex_exec_server_protocol::HttpRequestBodyDeltaNotification; +use codex_exec_server_protocol::JSONRPCMessage; +use codex_exec_server_protocol::MAX_HTTP_BODY_DELTA_BYTES; + +use super::MAX_CLIENT_INBOUND_NOTIFICATION_LEN; +use super::MAX_CLIENT_INBOUND_REQUEST_LEN; +use super::client_inbound_message_exceeded_limit; + +#[test] +fn rejects_oversized_request_hidden_in_raw_value_wrapper() { + let request = serde_json::json!({ + "method": "network/policyRequest", + "id": 1, + "params": { "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) }, + }); + let encoded = serde_json::to_vec(&serde_json::json!({ + "$serde_json::private::RawValue": request.to_string(), + })) + .expect("wrapped request should serialize"); + + let message = serde_json::from_slice::(&encoded); + assert!(matches!(&message, Ok(JSONRPCMessage::Request(_)))); + assert_eq!( + client_inbound_message_exceeded_limit( + message.as_ref(), + encoded.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ), + Some(MAX_CLIENT_INBOUND_REQUEST_LEN) + ); +} + +#[test] +fn bounds_notifications_without_rejecting_streamed_http_bodies() { + let streamed_body = serde_json::to_vec(&serde_json::json!({ + "method": HTTP_REQUEST_BODY_DELTA_METHOD, + "params": HttpRequestBodyDeltaNotification { + request_id: "request".to_string(), + seq: 1, + delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES].into(), + done: false, + error: None, + }, + })) + .expect("streamed body notification should serialize"); + let message = serde_json::from_slice::(&streamed_body); + assert_eq!( + client_inbound_message_exceeded_limit( + message.as_ref(), + streamed_body.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ), + None + ); + + for method in [ + EXEC_OUTPUT_DELTA_METHOD, + EXEC_EXITED_METHOD, + EXEC_CLOSED_METHOD, + ] { + let notification = serde_json::to_vec(&serde_json::json!({ + "method": method, + "params": { "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) }, + })) + .expect("process notification should serialize"); + let message = serde_json::from_slice::(¬ification); + assert_eq!( + client_inbound_message_exceeded_limit( + message.as_ref(), + notification.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ), + None + ); + } + + let oversized_notification = serde_json::to_vec(&serde_json::json!({ + "method": "hostile/notification", + "params": { "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) }, + })) + .expect("oversized notification should serialize"); + assert!(oversized_notification.len() < MAX_CLIENT_INBOUND_NOTIFICATION_LEN); + let message = serde_json::from_slice::(&oversized_notification); + assert_eq!( + client_inbound_message_exceeded_limit( + message.as_ref(), + oversized_notification.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ), + Some(MAX_CLIENT_INBOUND_REQUEST_LEN) + ); +} + +#[test] +fn preserves_oversized_responses() { + let response = serde_json::to_vec(&serde_json::json!({ + "id": 1, + "result": { "padding": "x".repeat(MAX_CLIENT_INBOUND_NOTIFICATION_LEN) }, + })) + .expect("response should serialize"); + let message = serde_json::from_slice::(&response); + + assert_eq!( + client_inbound_message_exceeded_limit( + message.as_ref(), + response.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ), + None + ); +} diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 0ec478adcf..a6867f2a27 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use std::time::Duration; use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncReadExt; use tokio::io::BufReader; use tokio::process::Command; use tokio::sync::OwnedSemaphorePermit; @@ -50,6 +51,7 @@ use crate::noise_relay::noise_relay_websocket_config; use crate::relay::harness_connection_from_websocket; use crate::trace_context::current_rendezvous_headers; +const MAX_STDIO_STDERR_LOG_LINE_LEN: u64 = 8 * 1024; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; const INITIAL_REGISTRY_MAX_RETRIES: u32 = 4; const INITIAL_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(6); @@ -773,11 +775,22 @@ impl ExecServerClient { })?; if let Some(stderr) = child.stderr.take() { tokio::spawn(async move { - let mut lines = BufReader::new(stderr).lines(); + let mut reader = BufReader::new(stderr); + let mut line = Vec::new(); loop { - match lines.next_line().await { - Ok(Some(line)) => debug!("exec-server stdio stderr: {line}"), - Ok(None) => break, + line.clear(); + match (&mut reader) + .take(MAX_STDIO_STDERR_LOG_LINE_LEN) + .read_until(b'\n', &mut line) + .await + { + Ok(0) => break, + Ok(_) => { + let line = line.strip_suffix(b"\n").unwrap_or(&line); + let line = line.strip_suffix(b"\r").unwrap_or(line); + let line = String::from_utf8_lossy(line); + debug!("exec-server stdio stderr: {line}"); + } Err(err) => { warn!("failed to read exec-server stdio stderr: {err}"); break; @@ -788,8 +801,12 @@ impl ExecServerClient { } Self::connect( - JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio command".to_string()) - .with_child_process(child), + JsonRpcConnection::client_from_stdio( + stdout, + stdin, + "exec-server stdio command".to_string(), + ) + .with_child_process(child), args.into(), ) .await diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 249b63330e..a0e4a18d86 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -32,6 +32,9 @@ use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::io::BufWriter; +use crate::client_inbound_request_limit::MAX_CLIENT_INBOUND_REQUEST_LEN; +use crate::client_inbound_request_limit::client_inbound_message_exceeded_limit; + 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. @@ -286,6 +289,21 @@ impl JsonRpcConnection { writer, connection_label, MAX_STDIO_JSONRPC_MESSAGE_LEN, + /*max_inbound_request_len*/ None, + ) + } + + pub(crate) fn client_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, + Some(MAX_CLIENT_INBOUND_REQUEST_LEN), ) } @@ -294,6 +312,7 @@ impl JsonRpcConnection { writer: W, connection_label: String, max_message_len: usize, + max_inbound_request_len: Option, ) -> Self where R: AsyncRead + Unpin + Send + 'static, @@ -366,7 +385,25 @@ impl JsonRpcConnection { if line.trim().is_empty() { continue; } - match serde_json::from_str::(&line) { + let message = serde_json::from_str::(&line); + if let Some(max_inbound_request_len) = max_inbound_request_len + && let Some(max_len) = client_inbound_message_exceeded_limit( + message.as_ref(), + line.len(), + max_inbound_request_len, + ) + { + send_disconnected( + &incoming_tx_for_reader, + &disconnected_tx_for_reader, + Some(format!( + "JSON-RPC message from {reader_label} exceeds maximum length of {max_len} bytes" + )), + ) + .await; + break; + } + match message { Ok(message) => { if incoming_tx_for_reader .send(JsonRpcConnectionEvent::message(message)) @@ -433,17 +470,28 @@ impl JsonRpcConnection { T: Sink + Stream> + Unpin + Send + 'static, E: std::fmt::Display + Send + 'static, { - Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) + Self::from_websocket_stream( + stream, + connection_label, + /*ping_interval*/ None, + Some(MAX_CLIENT_INBOUND_REQUEST_LEN), + ) } pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self { - Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL)) + Self::from_websocket_stream( + stream, + connection_label, + Some(WEBSOCKET_KEEPALIVE_INTERVAL), + /*max_inbound_request_len*/ None, + ) } fn from_websocket_stream( mut websocket: T, connection_label: String, ping_interval: Option, + max_inbound_request_len: Option, ) -> Self where T: Sink + Stream> + Unpin + Send + 'static, @@ -501,7 +549,7 @@ impl JsonRpcConnection { } incoming_message = websocket.next() => { match incoming_message { - Some(Ok(message)) => match message.parse_jsonrpc_frame() { + Some(Ok(message)) => match message.parse_jsonrpc_frame(max_inbound_request_len) { Ok(JsonRpcWebSocketFrame::Message(message)) => { if incoming_tx .send(JsonRpcConnectionEvent::message(message)) @@ -520,6 +568,17 @@ impl JsonRpcConnection { .await; break; } + Ok(JsonRpcWebSocketFrame::OverlongMessage { max_len }) => { + send_disconnected( + &incoming_tx, + &disconnected_tx, + Some(format!( + "JSON-RPC message from {connection_label} exceeds maximum length of {max_len} bytes" + )), + ) + .await; + break; + } Ok(JsonRpcWebSocketFrame::Ignore) => {} Err(err) => { send_malformed_message( @@ -575,24 +634,35 @@ impl JsonRpcConnection { enum JsonRpcWebSocketFrame { Message(JSONRPCMessage), Close, + OverlongMessage { max_len: usize }, Ignore, } trait JsonRpcWebSocketMessage: Send + 'static { - fn parse_jsonrpc_frame(self) -> Result; + fn parse_jsonrpc_frame( + self, + max_inbound_request_len: Option, + ) -> Result; fn from_text(text: String) -> Self; fn ping() -> Self; } impl JsonRpcWebSocketMessage for Message { - fn parse_jsonrpc_frame(self) -> Result { + fn parse_jsonrpc_frame( + self, + max_inbound_request_len: Option, + ) -> Result { match self { - Message::Text(text) => { - serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message) - } - Message::Binary(bytes) => { - serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message) - } + Message::Text(text) => parse_jsonrpc_message( + serde_json::from_str(text.as_ref()), + text.len(), + max_inbound_request_len, + ), + Message::Binary(bytes) => parse_jsonrpc_message( + serde_json::from_slice(bytes.as_ref()), + bytes.len(), + max_inbound_request_len, + ), Message::Close(_) => Ok(JsonRpcWebSocketFrame::Close), Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => { Ok(JsonRpcWebSocketFrame::Ignore) @@ -610,14 +680,21 @@ impl JsonRpcWebSocketMessage for Message { } impl JsonRpcWebSocketMessage for AxumWebSocketMessage { - fn parse_jsonrpc_frame(self) -> Result { + fn parse_jsonrpc_frame( + self, + max_inbound_request_len: Option, + ) -> Result { match self { - AxumWebSocketMessage::Text(text) => { - serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message) - } - AxumWebSocketMessage::Binary(bytes) => { - serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message) - } + AxumWebSocketMessage::Text(text) => parse_jsonrpc_message( + serde_json::from_str(text.as_ref()), + text.len(), + max_inbound_request_len, + ), + AxumWebSocketMessage::Binary(bytes) => parse_jsonrpc_message( + serde_json::from_slice(bytes.as_ref()), + bytes.len(), + max_inbound_request_len, + ), AxumWebSocketMessage::Close(_) => Ok(JsonRpcWebSocketFrame::Close), AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => { Ok(JsonRpcWebSocketFrame::Ignore) @@ -634,6 +711,23 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { } } +fn parse_jsonrpc_message( + message: Result, + encoded_len: usize, + max_inbound_request_len: Option, +) -> Result { + if let Some(max_inbound_request_len) = max_inbound_request_len + && let Some(max_len) = client_inbound_message_exceeded_limit( + message.as_ref(), + encoded_len, + max_inbound_request_len, + ) + { + return Ok(JsonRpcWebSocketFrame::OverlongMessage { max_len }); + } + message.map(JsonRpcWebSocketFrame::Message) +} + async fn send_disconnected( incoming_tx: &mpsc::Sender, disconnected_tx: &watch::Sender, @@ -732,6 +826,7 @@ mod tests { tokio::io::sink(), "test stdio peer".to_string(), max_message_len, + /*max_inbound_request_len*/ None, ); peer.write_all(encoded.as_bytes()).await?; @@ -762,6 +857,7 @@ mod tests { tokio::io::sink(), "hostile stdio peer".to_string(), max_message_len, + /*max_inbound_request_len*/ None, ); let overlong_message = vec![b'x'; max_message_len + 1]; @@ -785,6 +881,91 @@ mod tests { Ok(()) } + #[tokio::test] + async fn client_stdio_connection_rejects_overlong_request() -> anyhow::Result<()> { + let request = serde_json::json!({ + "params": { "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) }, + "method": "network/policyRequest", + "id": null, + }); + let encoded = serde_json::to_vec(&request)?; + let (reader, mut peer) = tokio::io::duplex(encoded.len().saturating_add(1)); + let mut connection = JsonRpcConnection::client_from_stdio( + reader, + tokio::io::sink(), + "hostile stdio peer".to_string(), + ); + + peer.write_all(&encoded).await?; + peer.write_all(b"\n").await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Disconnected { reason: Some(reason) }) + if reason == format!( + "JSON-RPC message from hostile stdio peer exceeds maximum length of {MAX_CLIENT_INBOUND_REQUEST_LEN} bytes" + ) + )); + + drop(peer); + drop(connection); + Ok(()) + } + + #[tokio::test] + async fn client_websocket_connection_rejects_overlong_request_and_accepts_large_response() + -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into()); + let response = serde_json::json!({ + "id": 1, + "result": { "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) }, + }); + + server_websocket + .send(Message::Text(response.to_string().into())) + .await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Message(JSONRPCMessage::Response(_))) + )); + + let request = serde_json::json!({ + "params": { "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) }, + "method": "network/policyRequest", + "id": 2, + }); + server_websocket + .send(Message::Binary(request.to_string().into_bytes().into())) + .await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Disconnected { reason: Some(reason) }) + if reason == format!( + "JSON-RPC message from test exceeds maximum length of {MAX_CLIENT_INBOUND_REQUEST_LEN} bytes" + ) + )); + + drop(connection); + Ok(()) + } + + #[test] + fn client_request_limit_fails_closed_for_malformed_oversized_frames() { + let malformed = format!( + "{{\"params\":{{\"padding\":\"{}\"}},\"method\":\"network/policyRequest\",\"id\":", + "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN) + ); + let message = serde_json::from_str::(&malformed); + assert_eq!( + client_inbound_message_exceeded_limit( + message.as_ref(), + malformed.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ), + Some(MAX_CLIENT_INBOUND_REQUEST_LEN) + ); + } + #[tokio::test] async fn websocket_connection_sends_configured_ping() -> anyhow::Result<()> { let (client_websocket, mut server_websocket) = websocket_pair().await?; @@ -792,6 +973,7 @@ mod tests { client_websocket, "test".into(), Some(WEBSOCKET_KEEPALIVE_INTERVAL), + /*max_inbound_request_len*/ None, ); let message = timeout(Duration::from_secs(1), server_websocket.next()) @@ -870,6 +1052,7 @@ mod tests { websocket, "test".into(), /*ping_interval*/ None, + /*max_inbound_request_len*/ None, ); let message = test_jsonrpc_message(); diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index c512d0d035..2671e83382 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -46,6 +46,9 @@ mod telemetry; mod trace_context; mod websocket_pong_watchdog; +// Shared limits for inbound executor messages across all transports. +mod client_inbound_request_limit; + use codex_exec_server_protocol as protocol; /// Process-local opt-in for tying a remote executor to its parent's stdin pipe. diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 793f7c1e12..8a682ea82f 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -32,6 +32,7 @@ use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; use crate::noise_channel::NoiseTransport; use crate::noise_channel::noise_channel_prologue; +use crate::noise_relay::message_framing::ClientJsonRpcMessageDecoder; use crate::noise_relay::message_framing::JsonRpcMessageDecoder; use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN; use crate::noise_relay::message_framing::frame_jsonrpc_message; @@ -274,7 +275,7 @@ where let mut websocket = websocket.peekable(); let mut next_outbound_seq = 0u32; let mut inbound_ciphertexts = OrderedCiphertextFrames::default(); - let mut inbound_decoder = JsonRpcMessageDecoder::default(); + let mut inbound_decoder = JsonRpcMessageDecoder::client(); let mut keepalive = tokio::time::interval_at( tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, WEBSOCKET_KEEPALIVE_INTERVAL, @@ -594,7 +595,7 @@ where async fn receive_data( inbound_ciphertexts: &mut OrderedCiphertextFrames, transport: &mut NoiseTransport, - decoder: &mut JsonRpcMessageDecoder, + decoder: &mut ClientJsonRpcMessageDecoder, data: RelayData, delivery_deadline: tokio::time::Instant, incoming_tx: &mpsc::Sender, diff --git a/codex-rs/exec-server/src/noise_relay/message_framing.rs b/codex-rs/exec-server/src/noise_relay/message_framing.rs index af9fbd4998..b6d30c2cee 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -4,6 +4,8 @@ use bytes::BytesMut; use codex_exec_server_protocol::JSONRPCMessage; use crate::ExecServerError; +use crate::client_inbound_request_limit::MAX_CLIENT_INBOUND_REQUEST_LEN; +use crate::client_inbound_request_limit::client_inbound_message_exceeded_limit; const LENGTH_PREFIX_BYTES: usize = size_of::(); pub(crate) const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024; @@ -109,6 +111,49 @@ impl MessageDecoder { } } +#[derive(Default)] +pub(crate) struct ClientJsonRpcMessageDecoder { + decoder: JsonRpcMessageDecoder, +} + +impl ClientJsonRpcMessageDecoder { + pub(crate) fn push( + &mut self, + plaintext_record: &[u8], + ) -> Result, ExecServerError> { + if self.decoder.decoder.buffered.is_empty() + && plaintext_record.len() <= MAX_CLIENT_INBOUND_REQUEST_LEN + { + return self.decoder.push(plaintext_record); + } + + self.decoder + .decoder + .push(plaintext_record)? + .into_iter() + .map(|encoded| { + let message = serde_json::from_slice(&encoded); + if let Some(max_len) = client_inbound_message_exceeded_limit( + message.as_ref(), + encoded.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ) { + return Err(ExecServerError::Protocol(format!( + "Noise relay JSON-RPC message exceeds maximum length of {max_len} bytes" + ))); + } + message.map_err(Into::into) + }) + .collect() + } +} + +impl JsonRpcMessageDecoder { + pub(crate) fn client() -> ClientJsonRpcMessageDecoder { + ClientJsonRpcMessageDecoder::default() + } +} + #[cfg(test)] #[path = "message_framing_tests.rs"] mod tests; diff --git a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs index 563e9e3bad..a2a713dafd 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing_tests.rs @@ -1,5 +1,9 @@ +use codex_exec_server_protocol::EXEC_OUTPUT_DELTA_METHOD; use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCNotification; +use codex_exec_server_protocol::JSONRPCRequest; +use codex_exec_server_protocol::JSONRPCResponse; +use codex_exec_server_protocol::RequestId; use pretty_assertions::assert_eq; use super::JsonRpcMessageDecoder; @@ -7,11 +11,12 @@ use super::MAX_NOISE_JSONRPC_MESSAGE_LEN; use super::NOISE_RECORD_PLAINTEXT_LEN; use super::frame_jsonrpc_message; use crate::ExecServerError; +use crate::client_inbound_request_limit::MAX_CLIENT_INBOUND_REQUEST_LEN; #[test] fn fragments_and_reassembles_large_jsonrpc_message() { let message = JSONRPCMessage::Notification(JSONRPCNotification { - method: "large/test".to_string(), + method: EXEC_OUTPUT_DELTA_METHOD.to_string(), params: Some(serde_json::json!({ "data": "x".repeat(128 * 1024), })), @@ -19,7 +24,7 @@ fn fragments_and_reassembles_large_jsonrpc_message() { let framed = frame_jsonrpc_message(&message).unwrap(); assert!(framed.len() > 128 * 1024); - let mut decoder = JsonRpcMessageDecoder::default(); + let mut decoder = JsonRpcMessageDecoder::client(); let mut decoded = Vec::new(); for record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) { decoded.extend(decoder.push(record).unwrap()); @@ -51,6 +56,63 @@ fn rejects_oversized_plaintext_record() { )); } +#[test] +fn client_decoder_rejects_oversized_request_and_preserves_large_response() { + let request = JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "network/policyRequest".to_string(), + params: Some(serde_json::json!({ + "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN), + })), + trace: None, + }); + let request = frame_jsonrpc_message(&request).unwrap(); + let mut client = JsonRpcMessageDecoder::client(); + assert!(matches!( + client.push(&request), + Err(ExecServerError::Protocol(message)) + if message == format!( + "Noise relay JSON-RPC message exceeds maximum length of {MAX_CLIENT_INBOUND_REQUEST_LEN} bytes" + ) + )); + + let mut client = JsonRpcMessageDecoder::client(); + let split = MAX_CLIENT_INBOUND_REQUEST_LEN / 2; + assert!(client.push(&request[..split]).unwrap().is_empty()); + assert!(matches!( + client.push(&request[split..]), + Err(ExecServerError::Protocol(message)) + if message == format!( + "Noise relay JSON-RPC message exceeds maximum length of {MAX_CLIENT_INBOUND_REQUEST_LEN} bytes" + ) + )); + + let response = JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result: serde_json::json!({ + "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN), + }), + }); + let framed = frame_jsonrpc_message(&response).unwrap(); + let mut client = JsonRpcMessageDecoder::client(); + assert_eq!(client.push(&framed).unwrap(), vec![response]); +} + +#[test] +fn executor_decoder_preserves_large_request() { + let request = JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "process/start".to_string(), + params: Some(serde_json::json!({ + "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN), + })), + trace: None, + }); + let framed = frame_jsonrpc_message(&request).unwrap(); + let mut executor = JsonRpcMessageDecoder::default(); + assert_eq!(executor.push(&framed).unwrap(), vec![request]); +} + #[test] fn reassembles_many_messages_from_one_record() { let message = JSONRPCMessage::Notification(JSONRPCNotification { diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 3fc1324eeb..bd1346953f 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -21,6 +21,8 @@ use tracing::warn; use uuid::Uuid; use crate::ExecServerError; +use crate::client_inbound_request_limit::MAX_CLIENT_INBOUND_REQUEST_LEN; +use crate::client_inbound_request_limit::client_inbound_message_exceeded_limit; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; @@ -204,11 +206,6 @@ impl RelayMessageFrame { } } - fn into_jsonrpc_message(self) -> Result { - let payload = self.into_data()?.payload; - serde_json::from_slice(&payload).map_err(ExecServerError::Json) - } - pub(crate) fn into_handshake_payload(self) -> Result, ExecServerError> { let kind = self.validate()?; if kind != RelayFrameBodyKind::Handshake { @@ -380,8 +377,35 @@ where } }; match kind { - RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() { - Ok(message) => { + RelayFrameBodyKind::Data => match frame.into_data() { + Ok(data) => { + let message = serde_json::from_slice(&data.payload); + if let Some(max_len) = client_inbound_message_exceeded_limit( + message.as_ref(), + data.payload.len(), + MAX_CLIENT_INBOUND_REQUEST_LEN, + ) { + let _ = disconnected_tx.send(true); + let _ = incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { + reason: Some(format!( + "relay JSON-RPC message from {reader_label} exceeds maximum length of {max_len} bytes" + )), + }) + .await; + break; + } + let message = match message { + Ok(message) => message, + Err(err) => { + let _ = incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: err.to_string(), + }) + .await; + continue; + } + }; match send_event_with_keepalive( &mut websocket, &mut keepalive, @@ -968,6 +992,7 @@ mod tests { use std::time::Duration; use codex_exec_server_protocol::JSONRPCRequest; + use codex_exec_server_protocol::JSONRPCResponse; use codex_exec_server_protocol::RequestId; use futures::Sink; use futures::Stream; @@ -1017,6 +1042,77 @@ mod tests { Ok(()) } + #[tokio::test] + async fn harness_connection_rejects_oversized_request_and_accepts_large_response() + -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = + harness_connection_from_websocket(client_websocket, "test".to_string()); + let stream_id = read_resume_stream_id(&mut server_websocket).await?; + let response = JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result: serde_json::json!({ + "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN), + }), + }); + + server_websocket + .send(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id: stream_id.clone(), + body: Some(relay_message_frame::Body::Data(RelayData { + seq: 0, + segment_index: 0, + segment_count: 1, + payload: jsonrpc_payload(&response)?, + })), + ..RelayMessageFrame::default() + }) + .into(), + )) + .await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Message(actual)) if actual == response + )); + + let request = JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(2), + method: "network/policyRequest".to_string(), + params: Some(serde_json::json!({ + "padding": "x".repeat(MAX_CLIENT_INBOUND_REQUEST_LEN), + })), + trace: None, + }); + server_websocket + .send(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + body: Some(relay_message_frame::Body::Data(RelayData { + seq: 1, + segment_index: 0, + segment_count: 1, + payload: jsonrpc_payload(&request)?, + })), + ..RelayMessageFrame::default() + }) + .into(), + )) + .await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Disconnected { reason: Some(reason) }) + if reason == format!( + "relay JSON-RPC message from test exceeds maximum length of {MAX_CLIENT_INBOUND_REQUEST_LEN} bytes" + ) + )); + + drop(connection); + Ok(()) + } + #[tokio::test] async fn multiplexed_environment_sends_keepalive() -> anyhow::Result<()> { let (client_websocket, mut server_websocket) = websocket_pair().await?; @@ -1170,7 +1266,10 @@ mod tests { }; let frame = decode_relay_message_frame(data_payload.as_ref())?; assert_eq!(frame.stream_id, stream_id); - assert_eq!(frame.into_jsonrpc_message()?, message); + assert_eq!( + serde_json::from_slice::(&frame.into_data()?.payload)?, + message + ); drop(connection); Ok(()) } diff --git a/codex-rs/exec-server/src/rpc_server_requests.rs b/codex-rs/exec-server/src/rpc_server_requests.rs index b7d158802b..54d7454628 100644 --- a/codex-rs/exec-server/src/rpc_server_requests.rs +++ b/codex-rs/exec-server/src/rpc_server_requests.rs @@ -20,6 +20,7 @@ use crate::rpc::RpcCallError; use crate::rpc::RpcServerOutboundMessage; pub(crate) const MAX_IN_FLIGHT_SERVER_CALLS: usize = 256; +const MAX_SERVER_REQUEST_TRACESTATE_LEN: usize = 512; type PendingRequest = oneshot::Sender>; @@ -98,11 +99,21 @@ impl RpcServerRequestSender { inner: Arc::clone(&self.inner), request_id: request_id.clone(), }; + let trace = codex_otel::current_span_w3c_trace_context().map(|mut trace| { + if trace + .tracestate + .as_ref() + .is_some_and(|tracestate| tracestate.len() > MAX_SERVER_REQUEST_TRACESTATE_LEN) + { + trace.tracestate = None; + } + trace + }); let request = RpcServerOutboundMessage::Request(JSONRPCRequest { id: request_id, method: method.to_string(), params: Some(params), - trace: codex_otel::current_span_w3c_trace_context(), + trace, }); let response = timeout(call_timeout, async { diff --git a/codex-rs/exec-server/src/rpc_server_requests_tests.rs b/codex-rs/exec-server/src/rpc_server_requests_tests.rs index 05692c32aa..e86789cb32 100644 --- a/codex-rs/exec-server/src/rpc_server_requests_tests.rs +++ b/codex-rs/exec-server/src/rpc_server_requests_tests.rs @@ -3,12 +3,25 @@ use std::time::Duration; use codex_exec_server_protocol::JSONRPCRequest; use codex_exec_server_protocol::RequestId; +use codex_protocol::protocol::W3cTraceContext; +use opentelemetry::Context; +use opentelemetry::trace::SpanContext; +use opentelemetry::trace::SpanId; +use opentelemetry::trace::TraceContextExt as _; +use opentelemetry::trace::TraceFlags; +use opentelemetry::trace::TraceId; +use opentelemetry::trace::TraceState; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; use tokio::sync::mpsc; use tokio::task::JoinSet; use tokio::time::timeout; +use tracing::Instrument; +use tracing_subscriber::prelude::*; use super::MAX_IN_FLIGHT_SERVER_CALLS; +use super::MAX_SERVER_REQUEST_TRACESTATE_LEN; use super::RpcServerRequestSender; use crate::rpc::RpcCallError; use crate::rpc::RpcServerOutboundMessage; @@ -36,6 +49,75 @@ impl RpcServerRequestSender { } } +#[tokio::test(flavor = "current_thread")] +async fn rpc_server_sender_bounds_tracestate_without_dropping_traceparent() { + let tracer_provider = SdkTracerProvider::builder().build(); + let tracer = tracer_provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _subscriber_guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); + + let ordinary_tracestate = TraceState::from_key_value([("vendor", "value")]) + .expect("ordinary tracestate should be valid"); + let oversized_tracestate = + TraceState::from_key_value([("first", "a".repeat(256)), ("second", "b".repeat(256))]) + .expect("oversized aggregate tracestate should have valid members"); + assert!(oversized_tracestate.header().len() > MAX_SERVER_REQUEST_TRACESTATE_LEN); + + for (tracestate, expected_tracestate) in [ + ( + ordinary_tracestate.clone(), + Some(ordinary_tracestate.header()), + ), + (oversized_tracestate, None), + ] { + let span_context = SpanContext::new( + TraceId::from_hex("00000000000000000000000000000001").expect("valid trace id"), + SpanId::from_hex("0000000000000002").expect("valid span id"), + TraceFlags::SAMPLED, + /*is_remote*/ true, + tracestate, + ); + let span = tracing::info_span!("executor-outbound-request"); + codex_otel::set_parent_from_context( + &span, + Context::new().with_remote_span_context(span_context), + ); + let expected_traceparent = codex_otel::span_w3c_trace_context(&span) + .expect("span should have trace context") + .traceparent; + + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*buffer*/ 1); + let requests = RpcServerRequestSender::new(outgoing_tx); + let params = serde_json::json!({}); + let call = requests + .call_with_timeout::<_, serde_json::Value>( + "network/policyRequest", + ¶ms, + Duration::from_secs(1), + ) + .instrument(span); + tokio::pin!(call); + assert!(futures::poll!(call.as_mut()).is_pending()); + + let request = receive_server_request(&mut outgoing_rx).await; + assert_eq!( + request.trace, + Some(W3cTraceContext { + traceparent: expected_traceparent, + tracestate: expected_tracestate, + }) + ); + + requests.complete(request.id, Ok(serde_json::json!({}))); + assert_eq!( + call.await.expect("traced request should complete"), + serde_json::json!({}) + ); + } +} + #[tokio::test] async fn rpc_server_sender_matches_out_of_order_responses_by_request_id() { let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*buffer*/ 8);