diff --git a/codex-rs/exec-server-protocol/src/rpc.rs b/codex-rs/exec-server-protocol/src/rpc.rs index b004ade9d783..12bab5cddc99 100644 --- a/codex-rs/exec-server-protocol/src/rpc.rs +++ b/codex-rs/exec-server-protocol/src/rpc.rs @@ -7,10 +7,24 @@ use std::fmt; use codex_protocol::protocol::W3cTraceContext; use serde::Deserialize; +use serde::Deserializer; use serde::Serialize; +use serde::de; +use serde::de::DeserializeSeed; +use serde::de::MapAccess; +use serde::de::SeqAccess; +use serde::de::Visitor; +use serde_json::Map; +use serde_json::Number; +use serde_json::Value; pub const JSONRPC_VERSION: &str = "2.0"; +// A maximum-size fs/walk response has at most 50,000 entries and needs roughly +// 150,000 JSON values. Keep ample headroom for legitimate protocol messages +// while preventing compact arrays from expanding into millions of heap values. +const MAX_JSONRPC_VALUE_NODES: usize = 256 * 1024; + #[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Deserialize, Serialize, Hash, Eq)] #[serde(untagged)] pub enum RequestId { @@ -30,7 +44,7 @@ impl fmt::Display for RequestId { pub type Result = serde_json::Value; /// Any valid exec-server JSON-RPC object that can be decoded from or encoded onto the wire. -#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +#[derive(Debug, Clone, PartialEq, Serialize)] #[serde(untagged)] pub enum JSONRPCMessage { Request(JSONRPCRequest), @@ -39,6 +53,151 @@ pub enum JSONRPCMessage { Error(JSONRPCError), } +#[derive(Deserialize)] +#[serde(untagged)] +enum JSONRPCMessageRepr { + Request(JSONRPCRequest), + Notification(JSONRPCNotification), + Response(JSONRPCResponse), + Error(JSONRPCError), +} + +impl From for JSONRPCMessage { + fn from(value: JSONRPCMessageRepr) -> Self { + match value { + JSONRPCMessageRepr::Request(request) => Self::Request(request), + JSONRPCMessageRepr::Notification(notification) => Self::Notification(notification), + JSONRPCMessageRepr::Response(response) => Self::Response(response), + JSONRPCMessageRepr::Error(error) => Self::Error(error), + } + } +} + +impl<'de> Deserialize<'de> for JSONRPCMessage { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + let mut remaining = MAX_JSONRPC_VALUE_NODES; + let value = BoundedValueSeed { + remaining: &mut remaining, + } + .deserialize(deserializer)?; + JSONRPCMessageRepr::deserialize(value) + .map(Self::from) + .map_err(de::Error::custom) + } +} + +struct BoundedValueSeed<'a> { + remaining: &'a mut usize, +} + +impl<'de> DeserializeSeed<'de> for BoundedValueSeed<'_> { + type Value = Value; + + fn deserialize(self, deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + let Some(remaining) = self.remaining.checked_sub(1) else { + return Err(de::Error::custom(format!( + "JSON-RPC message exceeds the limit of {MAX_JSONRPC_VALUE_NODES} JSON values" + ))); + }; + *self.remaining = remaining; + deserializer.deserialize_any(BoundedValueVisitor { + remaining: self.remaining, + }) + } +} + +struct BoundedValueVisitor<'a> { + remaining: &'a mut usize, +} + +impl<'de> Visitor<'de> for BoundedValueVisitor<'_> { + type Value = Value; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a JSON value within the exec-server complexity limit") + } + + fn visit_bool(self, value: bool) -> std::result::Result { + Ok(Value::Bool(value)) + } + + fn visit_i64(self, value: i64) -> std::result::Result { + Ok(Value::Number(value.into())) + } + + fn visit_u64(self, value: u64) -> std::result::Result { + Ok(Value::Number(value.into())) + } + + fn visit_f64(self, value: f64) -> std::result::Result { + Ok(Number::from_f64(value).map_or(Value::Null, Value::Number)) + } + + fn visit_str(self, value: &str) -> std::result::Result { + Ok(Value::String(value.to_string())) + } + + fn visit_string(self, value: String) -> std::result::Result { + Ok(Value::String(value)) + } + + fn visit_none(self) -> std::result::Result { + Ok(Value::Null) + } + + fn visit_some(self, deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + BoundedValueSeed { + remaining: self.remaining, + } + .deserialize(deserializer) + } + + fn visit_unit(self) -> std::result::Result { + Ok(Value::Null) + } + + fn visit_seq(self, mut sequence: A) -> std::result::Result + where + A: SeqAccess<'de>, + { + let mut values = Vec::new(); + while let Some(value) = sequence.next_element_seed(BoundedValueSeed { + remaining: &mut *self.remaining, + })? { + values.push(value); + } + Ok(Value::Array(values)) + } + + fn visit_map(self, mut object: A) -> std::result::Result + where + A: MapAccess<'de>, + { + let mut values = Map::new(); + while let Some(key) = object.next_key::()? { + if values.contains_key(&key) { + return Err(de::Error::custom(format!( + "duplicate JSON object key `{key}`" + ))); + } + let value = object.next_value_seed(BoundedValueSeed { + remaining: &mut *self.remaining, + })?; + values.insert(key, value); + } + Ok(Value::Object(values)) + } +} + /// A request that expects a response. #[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] pub struct JSONRPCRequest { @@ -79,3 +238,7 @@ pub struct JSONRPCErrorError { pub data: Option, pub message: String, } + +#[cfg(test)] +#[path = "rpc_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server-protocol/src/rpc_tests.rs b/codex-rs/exec-server-protocol/src/rpc_tests.rs new file mode 100644 index 000000000000..d60f18ef6a31 --- /dev/null +++ b/codex-rs/exec-server-protocol/src/rpc_tests.rs @@ -0,0 +1,99 @@ +use pretty_assertions::assert_eq; +use serde_json::json; + +use super::JSONRPCError; +use super::JSONRPCErrorError; +use super::JSONRPCMessage; +use super::JSONRPCNotification; +use super::JSONRPCRequest; +use super::JSONRPCResponse; +use super::MAX_JSONRPC_VALUE_NODES; +use super::RequestId; + +#[test] +fn round_trips_every_jsonrpc_message_variant() -> serde_json::Result<()> { + let messages = [ + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "request".to_string(), + params: Some(json!({"items": [1, 2, 3]})), + trace: None, + }), + JSONRPCMessage::Notification(JSONRPCNotification { + method: "notification".to_string(), + params: Some(json!({"enabled": true})), + }), + JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::String("response".to_string()), + result: json!({"value": "ok"}), + }), + JSONRPCMessage::Error(JSONRPCError { + error: JSONRPCErrorError { + code: -32603, + data: Some(json!({"retryable": false})), + message: "failed".to_string(), + }, + id: RequestId::Integer(2), + }), + ]; + + for expected in messages { + let encoded = serde_json::to_string(&expected)?; + let actual = serde_json::from_str::(&encoded)?; + assert_eq!(actual, expected); + } + + Ok(()) +} + +#[test] +fn accepts_large_scalar_payload() -> serde_json::Result<()> { + let expected = JSONRPCMessage::Notification(JSONRPCNotification { + method: "large".to_string(), + params: Some(json!({"data": "x".repeat(MAX_JSONRPC_VALUE_NODES + 1)})), + }); + + let encoded = serde_json::to_string(&expected)?; + let actual = serde_json::from_str::(&encoded)?; + + assert_eq!(actual, expected); + Ok(()) +} + +#[test] +fn rejects_duplicate_object_keys() { + let error = serde_json::from_str::(r#"{"method":"safe","method":"dangerous"}"#) + .expect_err("duplicate JSON object keys should be rejected"); + + assert!( + error + .to_string() + .contains("duplicate JSON object key `method`"), + "unexpected error: {error}" + ); +} + +#[test] +fn rejects_compact_array_heap_amplification() { + const REPRO_VALUE_COUNT: usize = 2_097_137; + const REPRO_MESSAGE_BYTES: usize = 4_194_303; + + let mut encoded = String::with_capacity(REPRO_MESSAGE_BYTES); + encoded.push_str(r#"{"method":"probe","params":["#); + for index in 0..REPRO_VALUE_COUNT { + if index != 0 { + encoded.push(','); + } + encoded.push('0'); + } + encoded.push_str("]}"); + assert_eq!(encoded.len(), REPRO_MESSAGE_BYTES); + + let error = serde_json::from_str::(&encoded) + .expect_err("amplification payload should exceed the JSON value limit"); + let expected_error = format!("exceeds the limit of {MAX_JSONRPC_VALUE_NODES} JSON values"); + assert!( + error.to_string().contains(&expected_error), + "unexpected error: {error}" + ); +} diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 3c488f0ad7c3..a87c03a58584 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -63,6 +63,9 @@ use crate::telemetry::ExecServerTelemetry; use crate::telemetry::ProcessMetricGuard; const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024; +// Each process/read chunk needs four JSON values. Keep retained replay below the +// shared 256K-value JSON-RPC decoder budget even when output arrives in tiny chunks. +const RETAINED_OUTPUT_CHUNKS_PER_PROCESS: usize = 50_000; const NOTIFICATION_CHANNEL_CAPACITY: usize = 256; const PROCESS_EVENT_CHANNEL_CAPACITY: usize = 256; const RETAINED_STDIN_WRITE_IDS_PER_PROCESS: usize = 4096; @@ -783,7 +786,9 @@ async fn stream_output( stream, chunk: chunk.clone(), }); - while process.retained_bytes > RETAINED_OUTPUT_BYTES_PER_PROCESS { + while process.retained_bytes > RETAINED_OUTPUT_BYTES_PER_PROCESS + || process.output.len() > RETAINED_OUTPUT_CHUNKS_PER_PROCESS + { let Some(evicted) = process.output.pop_front() else { break; }; @@ -977,6 +982,9 @@ fn notification_sender(inner: &Inner) -> Option { #[cfg(test)] mod tests { use super::*; + use codex_exec_server_protocol::JSONRPCMessage; + use codex_exec_server_protocol::JSONRPCResponse; + use codex_exec_server_protocol::RequestId; use codex_otel::MetricsConfig; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_utils_path_uri::PathUri; @@ -1226,6 +1234,99 @@ mod tests { backend.shutdown().await; } + #[tokio::test] + async fn process_read_replay_is_bounded_by_chunk_count() { + let backend = LocalProcess::default(); + let process = spawn_test_process(&backend, "proc-chunk-count").await; + let retained_chunk_count = RETAINED_OUTPUT_CHUNKS_PER_PROCESS as u64; + + { + let mut processes = backend.inner.processes.lock().await; + let Some(ProcessEntry::Running(running)) = processes.get_mut(&process.process_id) + else { + panic!("process should be running"); + }; + running.output = (1..=retained_chunk_count) + .map(|seq| RetainedOutputChunk { + seq, + stream: ExecOutputStream::Stdout, + chunk: vec![b'x'], + }) + .collect(); + running.retained_bytes = RETAINED_OUTPUT_CHUNKS_PER_PROCESS; + running.next_seq = retained_chunk_count + 1; + } + + process + .stdout_tx + .send(vec![b'y']) + .await + .expect("send output beyond retained chunk limit"); + timeout(Duration::from_secs(1), async { + loop { + let output_recorded = { + let processes = backend.inner.processes.lock().await; + let Some(ProcessEntry::Running(running)) = processes.get(&process.process_id) + else { + panic!("process should be running"); + }; + running.next_seq == retained_chunk_count + 2 + }; + if output_recorded { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("output should be retained"); + + let response = backend + .exec_read(ReadParams { + process_id: process.process_id.clone(), + after_seq: None, + max_bytes: None, + wait_ms: Some(0), + }) + .await + .expect("read retained output"); + let mut expected_chunks = (2..=retained_chunk_count) + .map(|seq| ProcessOutputChunk { + seq, + stream: ExecOutputStream::Stdout, + chunk: vec![b'x'].into(), + }) + .collect::>(); + expected_chunks.push(ProcessOutputChunk { + seq: retained_chunk_count + 1, + stream: ExecOutputStream::Stdout, + chunk: vec![b'y'].into(), + }); + assert_eq!( + response, + ReadResponse { + chunks: expected_chunks, + next_seq: retained_chunk_count + 2, + exited: false, + exit_code: None, + closed: false, + failure: None, + sandbox_denied: false, + } + ); + + let message = JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result: serde_json::to_value(response).expect("serialize process/read response"), + }); + let encoded = serde_json::to_string(&message).expect("encode JSON-RPC response"); + let decoded = serde_json::from_str::(&encoded) + .expect("retained process/read response should fit the JSON value budget"); + assert_eq!(decoded, message); + + backend.shutdown().await; + } + #[tokio::test] async fn closed_process_is_evicted_after_retention() { let backend = LocalProcess::default(); diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 9671f25d3933..de5d05ecc4f5 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -607,10 +607,11 @@ where P: DeserializeOwned, { let params = params.unwrap_or(Value::Null); - match serde_json::from_value(params.clone()) { + let retry_as_null = matches!(¶ms, Value::Object(map) if map.is_empty()); + match serde_json::from_value(params) { Ok(params) => Ok(params), Err(err) => { - if matches!(params, Value::Object(ref map) if map.is_empty()) { + if retry_as_null { serde_json::from_value(Value::Null).map_err(|_| err) } else { Err(err) diff --git a/codex-rs/exec-server/src/server/file_system_handler.rs b/codex-rs/exec-server/src/server/file_system_handler.rs index cbf263dad939..21da65053a2c 100644 --- a/codex-rs/exec-server/src/server/file_system_handler.rs +++ b/codex-rs/exec-server/src/server/file_system_handler.rs @@ -11,6 +11,7 @@ use crate::ExecutorFileSystem; use crate::RemoveOptions; use crate::file_read::FileReadHandleManager; use crate::local_file_system::LocalFileSystem; +use crate::protocol::FS_READ_DIRECTORY_METHOD; use crate::protocol::FS_WRITE_FILE_METHOD; use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCanonicalizeResponse; @@ -42,6 +43,9 @@ use crate::rpc::invalid_request; use crate::rpc::not_found; const MAX_FILE_READ_HANDLE_ID_BYTES: usize = 32; +// Each read-directory entry needs four JSON values. Keep same-version +// producers comfortably below the shared 256K-value decoder budget. +const MAX_READ_DIRECTORY_ENTRIES: usize = 50_000; #[derive(Clone)] pub(crate) struct FileSystemHandler { @@ -189,7 +193,14 @@ impl FileSystemHandler { .file_system .read_directory(¶ms.path, params.sandbox.as_ref()) .await - .map_err(map_fs_error)? + .map_err(map_fs_error)?; + let entry_count = entries.len(); + if entry_count > MAX_READ_DIRECTORY_ENTRIES { + return Err(internal_error(format!( + "{FS_READ_DIRECTORY_METHOD} returned {entry_count} entries; limit is {MAX_READ_DIRECTORY_ENTRIES}" + ))); + } + let entries = entries .into_iter() .map(|entry| FsReadDirectoryEntry { file_name: entry.file_name,