Bound exec-server JSON-RPC decoding complexity (#33013)

## Why

Compact JSON arrays can expand into millions of heap values during decoding, and duplicate object keys make a message ambiguous.

## What changed

- Limit exec-server JSON-RPC messages to 256K JSON values and reject duplicate object keys.
- Cap `fs/read_directory` results and retained `process/read` output at 50,000 entries or chunks so locally produced responses remain within the decoder budget.

## Testing

Add coverage for all JSON-RPC variants, large scalar payloads, duplicate keys, compact array amplification, and retained process output at the chunk limit.

GitOrigin-RevId: e31d1f25ab0a7e2272015c98174fd2b7cdd669d7

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<JSONRPCMessageRepr> 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<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    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<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
+    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<E>(self, value: bool) -> std::result::Result<Self::Value, E> {
+        Ok(Value::Bool(value))
+    }
+
+    fn visit_i64<E>(self, value: i64) -> std::result::Result<Self::Value, E> {
+        Ok(Value::Number(value.into()))
+    }
+
+    fn visit_u64<E>(self, value: u64) -> std::result::Result<Self::Value, E> {
+        Ok(Value::Number(value.into()))
+    }
+
+    fn visit_f64<E>(self, value: f64) -> std::result::Result<Self::Value, E> {
+        Ok(Number::from_f64(value).map_or(Value::Null, Value::Number))
+    }
+
+    fn visit_str<E>(self, value: &str) -> std::result::Result<Self::Value, E> {
+        Ok(Value::String(value.to_string()))
+    }
+
+    fn visit_string<E>(self, value: String) -> std::result::Result<Self::Value, E> {
+        Ok(Value::String(value))
+    }
+
+    fn visit_none<E>(self) -> std::result::Result<Self::Value, E> {
+        Ok(Value::Null)
+    }
+
+    fn visit_some<D>(self, deserializer: D) -> std::result::Result<Self::Value, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        BoundedValueSeed {
+            remaining: self.remaining,
+        }
+        .deserialize(deserializer)
+    }
+
+    fn visit_unit<E>(self) -> std::result::Result<Self::Value, E> {
+        Ok(Value::Null)
+    }
+
+    fn visit_seq<A>(self, mut sequence: A) -> std::result::Result<Self::Value, A::Error>
+    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<A>(self, mut object: A) -> std::result::Result<Self::Value, A::Error>
+    where
+        A: MapAccess<'de>,
+    {
+        let mut values = Map::new();
+        while let Some(key) = object.next_key::<String>()? {
+            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<serde_json::Value>,
     pub message: String,
 }
+
+#[cfg(test)]
+#[path = "rpc_tests.rs"]
+mod tests;

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::<JSONRPCMessage>(&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::<JSONRPCMessage>(&encoded)?;
+
+    assert_eq!(actual, expected);
+    Ok(())
+}
+
+#[test]
+fn rejects_duplicate_object_keys() {
+    let error = serde_json::from_str::<JSONRPCMessage>(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::<JSONRPCMessage>(&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}"
+    );
+}

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<RpcNotificationSender> {
 #[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::<Vec<_>>();
+        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::<JSONRPCMessage>(&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();

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!(&params, 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)

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(&params.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,