From e8870ee187f3d62fdb889d486c8ac5576d371f24 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 8 Jul 2026 14:48:23 +0100 Subject: [PATCH] Bound exec-server ingress --- codex-rs/exec-server/src/client_recovery.rs | 7 + codex-rs/exec-server/src/client_transport.rs | 29 ++- codex-rs/exec-server/src/connection.rs | 203 ++++++++++++++---- .../src/noise_relay/executor_stream.rs | 7 +- .../exec-server/src/noise_relay/harness.rs | 7 +- .../src/noise_relay/message_framing.rs | 17 +- .../src/noise_relay/message_framing_tests.rs | 6 +- codex-rs/exec-server/src/relay.rs | 32 ++- codex-rs/exec-server/src/rpc.rs | 104 +++++++-- codex-rs/exec-server/src/server/processor.rs | 2 +- 10 files changed, 323 insertions(+), 91 deletions(-) diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 252948cfc1..3f48913ba3 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -516,6 +516,13 @@ impl ExecServerClient { inner.request_recovery(rpc_client, disconnected_message(reason.as_deref())); return; } + RpcClientEvent::ProtocolError { reason } => { + rpc_client.close_transport().await; + inner + .fail(format!("exec-server protocol violation: {reason}")) + .await; + return; + } } } }); diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 1ccfd69152..432d8a6520 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -6,7 +6,6 @@ use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::process::Command; use tokio::time::timeout; -use tokio_tungstenite::connect_async; use tokio_tungstenite::connect_async_with_config; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tracing::debug; @@ -26,6 +25,7 @@ use crate::client_api::RemoteExecServerConnectArgs; use crate::client_api::StdioExecServerCommand; use crate::client_api::StdioExecServerConnectArgs; use crate::connection::JsonRpcConnection; +use crate::connection::jsonrpc_websocket_config; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_relay::NoiseHarnessConnectionArgs; use crate::noise_relay::noise_harness_connection_from_websocket; @@ -191,16 +191,23 @@ impl ExecServerClient { ensure_rustls_crypto_provider(); let websocket_url = args.websocket_url.clone(); let connect_timeout = args.connect_timeout; - let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) - .await - .map_err(|_| ExecServerError::WebSocketConnectTimeout { - url: websocket_url.clone(), - timeout: connect_timeout, - })? - .map_err(|source| ExecServerError::WebSocketConnect { - url: websocket_url.clone(), - source, - })?; + let (stream, _) = timeout( + connect_timeout, + connect_async_with_config( + websocket_url.as_str(), + Some(jsonrpc_websocket_config()), + /*disable_nagle*/ false, + ), + ) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: websocket_url.clone(), + timeout: connect_timeout, + })? + .map_err(|source| ExecServerError::WebSocketConnect { + url: websocket_url.clone(), + source, + })?; let connection_label = format!("exec-server websocket {websocket_url}"); let connection = if is_rendezvous_harness_url(&websocket_url) { diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 9bfe82c350..9ca748bf1d 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -24,11 +24,14 @@ 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; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; pub(crate) const CHANNEL_CAPACITY: usize = 128; +pub(crate) const MAX_JSONRPC_MESSAGE_BYTES: usize = 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); @@ -37,9 +40,16 @@ pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30 #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { - Message(JSONRPCMessage), - MalformedMessage { reason: String }, - Disconnected { reason: Option }, + Message { + message: JSONRPCMessage, + encoded_len: usize, + }, + MalformedMessage { + reason: String, + }, + Disconnected { + reason: Option, + }, } #[derive(Clone)] @@ -242,17 +252,51 @@ impl JsonRpcConnection { let incoming_tx_for_reader = incoming_tx.clone(); let disconnected_tx_for_reader = disconnected_tx.clone(); let reader_task = tokio::spawn(async move { - let mut lines = BufReader::new(reader).lines(); + let mut reader = BufReader::new(reader); loop { - match lines.next_line().await { - Ok(Some(line)) => { - if line.trim().is_empty() { + let mut encoded = Vec::new(); + match (&mut reader) + .take((MAX_JSONRPC_MESSAGE_BYTES + 2) as u64) + .read_until(b'\n', &mut encoded) + .await + { + Ok(0) => { + send_disconnected( + &incoming_tx_for_reader, + &disconnected_tx_for_reader, + /*reason*/ None, + ) + .await; + break; + } + Ok(_) => { + if encoded.last() == Some(&b'\n') { + encoded.pop(); + if encoded.last() == Some(&b'\r') { + encoded.pop(); + } + } + if encoded.len() > MAX_JSONRPC_MESSAGE_BYTES { + send_malformed_message( + &incoming_tx_for_reader, + Some(format!( + "JSON-RPC message from {reader_label} exceeds {MAX_JSONRPC_MESSAGE_BYTES} bytes" + )), + ) + .await; + break; + } + if encoded.iter().all(u8::is_ascii_whitespace) { continue; } - match serde_json::from_str::(&line) { + let encoded_len = encoded.len(); + match serde_json::from_slice::(&encoded) { Ok(message) => { if incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Message(message)) + .send(JsonRpcConnectionEvent::Message { + message, + encoded_len, + }) .await .is_err() { @@ -270,15 +314,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, @@ -392,36 +427,52 @@ impl JsonRpcConnection { } incoming_message = websocket.next() => { match incoming_message { - Some(Ok(message)) => match message.parse_jsonrpc_frame() { - Ok(JsonRpcWebSocketFrame::Message(message)) => { - if incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) - .await - .is_err() - { - break; - } - } - Ok(JsonRpcWebSocketFrame::Close) => { - send_disconnected( + Some(Ok(message)) => { + let encoded_len = message.jsonrpc_payload_len(); + if encoded_len.is_some_and(|len| len > MAX_JSONRPC_MESSAGE_BYTES) { + send_malformed_message( &incoming_tx, - &disconnected_tx, - /*reason*/ None, + Some(format!( + "websocket JSON-RPC message from {connection_label} exceeds {MAX_JSONRPC_MESSAGE_BYTES} bytes" + )), ) .await; break; } - Ok(JsonRpcWebSocketFrame::Ignore) => {} - Err(err) => { - send_malformed_message( - &incoming_tx, - Some(format!( - "failed to parse websocket JSON-RPC message from {connection_label}: {err}" - )), - ) - .await; + match message.parse_jsonrpc_frame() { + Ok(JsonRpcWebSocketFrame::Message(message)) => { + if incoming_tx + .send(JsonRpcConnectionEvent::Message { + message, + encoded_len: encoded_len.unwrap_or_default(), + }) + .await + .is_err() + { + break; + } + } + Ok(JsonRpcWebSocketFrame::Close) => { + send_disconnected( + &incoming_tx, + &disconnected_tx, + /*reason*/ None, + ) + .await; + break; + } + Ok(JsonRpcWebSocketFrame::Ignore) => {} + Err(err) => { + send_malformed_message( + &incoming_tx, + Some(format!( + "failed to parse websocket JSON-RPC message from {connection_label}: {err}" + )), + ) + .await; + } } - }, + } Some(Err(err)) => { send_disconnected( &incoming_tx, @@ -470,12 +521,21 @@ enum JsonRpcWebSocketFrame { } trait JsonRpcWebSocketMessage: Send + 'static { + fn jsonrpc_payload_len(&self) -> Option; fn parse_jsonrpc_frame(self) -> Result; fn from_text(text: String) -> Self; fn ping() -> Self; } impl JsonRpcWebSocketMessage for Message { + fn jsonrpc_payload_len(&self) -> Option { + match self { + Message::Text(text) => Some(text.len()), + Message::Binary(bytes) => Some(bytes.len()), + Message::Close(_) | Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => None, + } + } + fn parse_jsonrpc_frame(self) -> Result { match self { Message::Text(text) => { @@ -501,6 +561,16 @@ impl JsonRpcWebSocketMessage for Message { } impl JsonRpcWebSocketMessage for AxumWebSocketMessage { + fn jsonrpc_payload_len(&self) -> Option { + match self { + AxumWebSocketMessage::Text(text) => Some(text.len()), + AxumWebSocketMessage::Binary(bytes) => Some(bytes.len()), + AxumWebSocketMessage::Close(_) + | AxumWebSocketMessage::Ping(_) + | AxumWebSocketMessage::Pong(_) => None, + } + } + fn parse_jsonrpc_frame(self) -> Result { match self { AxumWebSocketMessage::Text(text) => { @@ -525,6 +595,12 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { } } +pub(crate) fn jsonrpc_websocket_config() -> WebSocketConfig { + WebSocketConfig::default() + .max_frame_size(Some(MAX_JSONRPC_MESSAGE_BYTES)) + .max_message_size(Some(MAX_JSONRPC_MESSAGE_BYTES)) +} + async fn send_disconnected( incoming_tx: &mpsc::Sender, disconnected_tx: &watch::Sender, @@ -675,13 +751,54 @@ mod tests { .await?; assert!(matches!( timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, - Some(JsonRpcConnectionEvent::Message(actual)) if actual == message + Some(JsonRpcConnectionEvent::Message { message: actual, .. }) if actual == message )); drop(connection); Ok(()) } + #[tokio::test] + async fn websocket_connection_rejects_oversized_jsonrpc_message() -> anyhow::Result<()> { + let (websocket, control, _outbound_rx) = + ControlledWebSocket::new(/*write_ready*/ true); + let mut connection = JsonRpcConnection::from_websocket_stream( + websocket, + "test".into(), + /*ping_interval*/ None, + ); + + control.send_inbound(Message::Text( + "x".repeat(MAX_JSONRPC_MESSAGE_BYTES + 1).into(), + ))?; + + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::MalformedMessage { reason }) + if reason.contains("exceeds") + )); + Ok(()) + } + + #[tokio::test] + async fn stdio_connection_rejects_oversized_jsonrpc_message() -> anyhow::Result<()> { + let (mut server_writer, client_reader) = tokio::io::duplex(MAX_JSONRPC_MESSAGE_BYTES + 2); + let (client_writer, _server_reader) = tokio::io::duplex(64); + let mut connection = + JsonRpcConnection::from_stdio(client_reader, client_writer, "test".into()); + + let mut oversized_message = vec![b'x'; MAX_JSONRPC_MESSAGE_BYTES + 1]; + oversized_message.push(b'\n'); + server_writer.write_all(&oversized_message).await?; + + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::MalformedMessage { reason }) + if reason.contains("exceeds") + )); + Ok(()) + } + #[tokio::test] async fn websocket_connection_keeps_outbound_message_while_send_is_backpressured() -> anyhow::Result<()> { diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 89f33d938a..6c89ab2e0a 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -74,9 +74,12 @@ impl NoiseVirtualStream { ExecServerError::Protocol(format!("Noise relay decryption failed: {error}")) })? }; - for message in self.inbound_decoder.push(&plaintext)? { + for (message, encoded_len) in self.inbound_decoder.push(&plaintext)? { self.incoming_tx - .try_send(JsonRpcConnectionEvent::Message(message)) + .try_send(JsonRpcConnectionEvent::Message { + message, + encoded_len, + }) .map_err(|_| { ExecServerError::Protocol( "Noise virtual stream inbound queue is full or closed".to_string(), diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 9969a3e130..8f0d77672f 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -582,10 +582,13 @@ async fn receive_data( // The authenticated byte stream can carry partial or multiple JSON-RPC // messages; emit only complete, successfully parsed messages. - for message in decoder.push(&plaintext)? { + for (message, encoded_len) in decoder.push(&plaintext)? { send_incoming_event( incoming_tx, - JsonRpcConnectionEvent::Message(message), + JsonRpcConnectionEvent::Message { + message, + encoded_len, + }, delivery_deadline, ) .await?; 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 402b68691a..7bdd6589c9 100644 --- a/codex-rs/exec-server/src/noise_relay/message_framing.rs +++ b/codex-rs/exec-server/src/noise_relay/message_framing.rs @@ -1,9 +1,9 @@ use codex_exec_server_protocol::JSONRPCMessage; use crate::ExecServerError; +use crate::connection::MAX_JSONRPC_MESSAGE_BYTES; const LENGTH_PREFIX_BYTES: usize = size_of::(); -const MAX_NOISE_JSONRPC_MESSAGE_LEN: usize = 64 * 1024 * 1024; pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024; /// Serialize one JSON-RPC message into the encrypted record byte stream. @@ -16,7 +16,7 @@ pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result, let mut framed = vec![0; LENGTH_PREFIX_BYTES]; serde_json::to_writer(&mut framed, message)?; let message_len = framed.len() - LENGTH_PREFIX_BYTES; - if message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN { + if message_len > MAX_JSONRPC_MESSAGE_BYTES { return Err(ExecServerError::Protocol( "Noise relay JSON-RPC message exceeds maximum length".to_string(), )); @@ -39,7 +39,7 @@ impl JsonRpcMessageDecoder { pub(crate) fn push( &mut self, plaintext_record: &[u8], - ) -> Result, ExecServerError> { + ) -> Result, ExecServerError> { if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN { return Err(ExecServerError::Protocol( "Noise relay plaintext record exceeds maximum length".to_string(), @@ -55,7 +55,7 @@ impl JsonRpcMessageDecoder { let message_len = u32::from_be_bytes([prefix[0], prefix[1], prefix[2], prefix[3]]) as usize; // Reject the authenticated length before waiting for its payload. - if message_len == 0 || message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN { + if message_len == 0 || message_len > MAX_JSONRPC_MESSAGE_BYTES { return Err(ExecServerError::Protocol( "Noise relay JSON-RPC message has invalid length".to_string(), )); @@ -64,14 +64,15 @@ impl JsonRpcMessageDecoder { if self.buffered.len() < framed_len { break; } - messages.push(serde_json::from_slice( - &self.buffered[LENGTH_PREFIX_BYTES..framed_len], - )?); + messages.push(( + serde_json::from_slice(&self.buffered[LENGTH_PREFIX_BYTES..framed_len])?, + message_len, + )); self.buffered.drain(..framed_len); } // Even before a message is complete, keep reassembly memory bounded. - if self.buffered.len() > LENGTH_PREFIX_BYTES + MAX_NOISE_JSONRPC_MESSAGE_LEN { + if self.buffered.len() > LENGTH_PREFIX_BYTES + MAX_JSONRPC_MESSAGE_BYTES { return Err(ExecServerError::Protocol( "Noise relay JSON-RPC reassembly buffer exceeds maximum length".to_string(), )); 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 3b9ac7939e..88154b49d8 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 @@ -3,10 +3,10 @@ use codex_exec_server_protocol::JSONRPCNotification; use pretty_assertions::assert_eq; use super::JsonRpcMessageDecoder; -use super::MAX_NOISE_JSONRPC_MESSAGE_LEN; use super::NOISE_RECORD_PLAINTEXT_LEN; use super::frame_jsonrpc_message; use crate::ExecServerError; +use crate::connection::MAX_JSONRPC_MESSAGE_BYTES; #[test] fn fragments_and_reassembles_large_jsonrpc_message() { @@ -25,13 +25,13 @@ fn fragments_and_reassembles_large_jsonrpc_message() { decoded.extend(decoder.push(record).unwrap()); } - assert_eq!(decoded, vec![message]); + assert_eq!(decoded, vec![(message, framed.len() - size_of::())]); } #[test] fn rejects_declared_message_length_above_limit_without_payload() { let mut decoder = JsonRpcMessageDecoder::default(); - let declared_len = (MAX_NOISE_JSONRPC_MESSAGE_LEN as u32 + 1).to_be_bytes(); + let declared_len = (MAX_JSONRPC_MESSAGE_BYTES as u32 + 1).to_be_bytes(); assert!(matches!( decoder.push(&declared_len), diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index b8d477f5db..b3d7bd1521 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -25,6 +25,7 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::MAX_JSONRPC_MESSAGE_BYTES; use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; @@ -194,9 +195,16 @@ impl RelayMessageFrame { } } - fn into_jsonrpc_message(self) -> Result { + fn into_jsonrpc_message(self) -> Result<(JSONRPCMessage, usize), ExecServerError> { let payload = self.into_data()?.payload; - serde_json::from_slice(&payload).map_err(ExecServerError::Json) + let encoded_len = payload.len(); + if encoded_len > MAX_JSONRPC_MESSAGE_BYTES { + return Err(ExecServerError::Protocol(format!( + "relay JSON-RPC message exceeds {MAX_JSONRPC_MESSAGE_BYTES} bytes" + ))); + } + let message = serde_json::from_slice(&payload).map_err(ExecServerError::Json)?; + Ok((message, encoded_len)) } pub(crate) fn into_handshake_payload(self) -> Result, ExecServerError> { @@ -365,12 +373,15 @@ where }; match kind { RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() { - Ok(message) => { + Ok((message, encoded_len)) => { match send_event_with_keepalive( &mut websocket, &mut keepalive, &incoming_tx, - JsonRpcConnectionEvent::Message(message), + JsonRpcConnectionEvent::Message { + message, + encoded_len, + }, ) .await { @@ -969,7 +980,7 @@ mod tests { .await?; assert!(matches!( timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, - Some(JsonRpcConnectionEvent::Message(actual)) if actual == message + Some(JsonRpcConnectionEvent::Message { message: actual, .. }) if actual == message )); drop(connection); @@ -1019,6 +1030,7 @@ mod tests { ControlledWebSocket::new(/*write_ready*/ true); let (incoming_tx, mut incoming_rx) = mpsc::channel(/*buffer*/ 1); let message = test_jsonrpc_message(); + let encoded_len = serde_json::to_vec(&message)?.len(); let expected_message = message.clone(); incoming_tx .send(JsonRpcConnectionEvent::MalformedMessage { @@ -1036,7 +1048,10 @@ mod tests { &mut websocket, &mut keepalive, &incoming_tx, - JsonRpcConnectionEvent::Message(message), + JsonRpcConnectionEvent::Message { + encoded_len, + message, + }, ) .await }); @@ -1055,7 +1070,8 @@ mod tests { )); assert!(matches!( incoming_rx.recv().await, - Some(JsonRpcConnectionEvent::Message(actual)) if actual == expected_message + Some(JsonRpcConnectionEvent::Message { message: actual, .. }) + if actual == expected_message )); Ok(()) } @@ -1129,7 +1145,7 @@ 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!(frame.into_jsonrpc_message()?.0, message); drop(connection); Ok(()) } diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index e762fe8b66..0e89b87649 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -6,6 +6,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::Instant; use codex_exec_server_protocol::JSONRPCError; use codex_exec_server_protocol::JSONRPCErrorError; @@ -29,6 +30,9 @@ use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010; +const INGRESS_RATE_WINDOW: Duration = Duration::from_secs(1); +const MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW: usize = 8 * 1024 * 1024; +const MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW: usize = 8 * 1024; #[derive(Debug)] pub(crate) enum RpcCallError { @@ -59,6 +63,54 @@ enum RpcCallTimeout { pub(crate) enum RpcClientEvent { Notification(JSONRPCNotification), Disconnected { reason: Option }, + ProtocolError { reason: String }, +} + +enum RpcReaderEnd { + Disconnected { reason: Option }, + ProtocolError { reason: String }, +} + +struct IngressRateWindow { + started_at: Instant, + bytes: usize, + messages: usize, +} + +impl IngressRateWindow { + fn new() -> Self { + Self { + started_at: Instant::now(), + bytes: 0, + messages: 0, + } + } + + fn record(&mut self, encoded_len: usize) -> Result<(), String> { + let now = Instant::now(); + if now.duration_since(self.started_at) >= INGRESS_RATE_WINDOW { + self.started_at = now; + self.bytes = 0; + self.messages = 0; + } + + let next_bytes = self.bytes.saturating_add(encoded_len); + let next_messages = self.messages.saturating_add(1); + if next_bytes > MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW { + return Err(format!( + "exec-server sent more than {MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW} JSON-RPC bytes in one second" + )); + } + if next_messages > MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW { + return Err(format!( + "exec-server sent more than {MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW} JSON-RPC messages in one second" + )); + } + + self.bytes = next_bytes; + self.messages = next_messages; + Ok(()) + } } #[derive(Debug, Clone, PartialEq)] @@ -263,36 +315,41 @@ impl RpcClient { let closed_for_reader = Arc::clone(&closed); let transport_for_reader = transport.clone(); let reader_task = tokio::spawn(async move { - let disconnect_reason = loop { + let mut ingress_rate = IngressRateWindow::new(); + let reader_end = loop { let Some(event) = incoming_rx.recv().await else { - break None; + break RpcReaderEnd::Disconnected { reason: None }; }; match event { - JsonRpcConnectionEvent::Message(message) => { + JsonRpcConnectionEvent::Message { + message, + encoded_len, + } => { + if let Err(reason) = ingress_rate.record(encoded_len) { + break RpcReaderEnd::ProtocolError { reason }; + } if let Err(err) = handle_server_message(&pending_for_reader, &event_tx, message).await { - let _ = err; - break None; + break RpcReaderEnd::ProtocolError { reason: err }; } } JsonRpcConnectionEvent::MalformedMessage { reason } => { - let _ = reason; - break None; + break RpcReaderEnd::ProtocolError { reason }; } JsonRpcConnectionEvent::Disconnected { reason } => { - break reason; + break RpcReaderEnd::Disconnected { reason }; } } }; closed_for_reader.store(true, Ordering::Release); drain_pending(&pending_for_reader).await; - let _ = event_tx - .send(RpcClientEvent::Disconnected { - reason: disconnect_reason, - }) - .await; + let event = match reader_end { + RpcReaderEnd::Disconnected { reason } => RpcClientEvent::Disconnected { reason }, + RpcReaderEnd::ProtocolError { reason } => RpcClientEvent::ProtocolError { reason }, + }; + let _ = event_tx.send(event).await; transport_for_reader.terminate(); }); @@ -626,6 +683,9 @@ mod tests { use tracing_subscriber::filter::filter_fn; use tracing_subscriber::prelude::*; + use super::IngressRateWindow; + use super::MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW; + use super::MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW; use super::RpcClient; use crate::connection::JsonRpcConnection; @@ -665,6 +725,24 @@ mod tests { } } + #[test] + fn ingress_rate_window_bounds_bytes() { + let mut window = IngressRateWindow::new(); + + assert!(window.record(MAX_INBOUND_JSONRPC_BYTES_PER_WINDOW).is_ok()); + assert!(window.record(1).is_err()); + } + + #[test] + fn ingress_rate_window_bounds_messages() { + let mut window = IngressRateWindow::new(); + + for _ in 0..MAX_INBOUND_JSONRPC_MESSAGES_PER_WINDOW { + assert!(window.record(0).is_ok()); + } + assert!(window.record(0).is_err()); + } + #[tokio::test] async fn rpc_client_matches_out_of_order_responses_by_request_id() { let (client_stdin, server_reader) = tokio::io::duplex(4096); diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 63f6bd45d9..2376c3e994 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -125,7 +125,7 @@ async fn run_connection( break; } } - JsonRpcConnectionEvent::Message(message) => match message { + JsonRpcConnectionEvent::Message { message, .. } => match message { codex_exec_server_protocol::JSONRPCMessage::Request(request) => { let request_started_at = Instant::now(); if let Some((method, route)) = router.request_route(request.method.as_str()) {