diff --git a/codex-rs/exec-server-protocol/Cargo.toml b/codex-rs/exec-server-protocol/Cargo.toml index 5a8628705f2..9dac711e1a0 100644 --- a/codex-rs/exec-server-protocol/Cargo.toml +++ b/codex-rs/exec-server-protocol/Cargo.toml @@ -20,7 +20,7 @@ codex-protocol = { workspace = true } codex-shell-command = { workspace = true } codex-utils-path-uri = { workspace = true } serde = { workspace = true, features = ["derive"] } -serde_json = { workspace = true } +serde_json = { workspace = true, features = ["arbitrary_precision", "raw_value"] } [dev-dependencies] pretty_assertions = { workspace = true } diff --git a/codex-rs/exec-server-protocol/src/rpc.rs b/codex-rs/exec-server-protocol/src/rpc.rs index 12bab5cddc9..1d79b2b6f34 100644 --- a/codex-rs/exec-server-protocol/src/rpc.rs +++ b/codex-rs/exec-server-protocol/src/rpc.rs @@ -25,6 +25,11 @@ pub const JSONRPC_VERSION: &str = "2.0"; // while preventing compact arrays from expanding into millions of heap values. const MAX_JSONRPC_VALUE_NODES: usize = 256 * 1024; +// With `arbitrary_precision` enabled, serde_json presents decimal, exponent, +// and out-of-range integer values to visitors as a synthetic one-entry map. +const SERDE_JSON_NUMBER_TOKEN: &str = "$serde_json::private::Number"; +const SERDE_JSON_RAW_VALUE_TOKEN: &str = "$serde_json::private::RawValue"; + #[derive(Debug, Clone, PartialEq, PartialOrd, Ord, Deserialize, Serialize, Hash, Eq)] #[serde(untagged)] pub enum RequestId { @@ -53,26 +58,6 @@ 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 @@ -83,9 +68,29 @@ impl<'de> Deserialize<'de> for JSONRPCMessage { remaining: &mut remaining, } .deserialize(deserializer)?; - JSONRPCMessageRepr::deserialize(value) - .map(Self::from) - .map_err(de::Error::custom) + let object = value + .as_object() + .ok_or_else(|| de::Error::custom("expected a JSON-RPC object"))?; + + if object.contains_key("method") { + if object.contains_key("id") { + JSONRPCRequest::deserialize(value) + .map(Self::Request) + .map_err(de::Error::custom) + } else { + JSONRPCNotification::deserialize(value) + .map(Self::Notification) + .map_err(de::Error::custom) + } + } else if object.contains_key("result") { + JSONRPCResponse::deserialize(value) + .map(Self::Response) + .map_err(de::Error::custom) + } else { + JSONRPCError::deserialize(value) + .map(Self::Error) + .map_err(de::Error::custom) + } } } @@ -182,7 +187,36 @@ impl<'de> Visitor<'de> for BoundedValueVisitor<'_> { where A: MapAccess<'de>, { + let Some(first_key) = object.next_key::()? else { + return Ok(Value::Object(Map::new())); + }; + + if first_key == SERDE_JSON_NUMBER_TOKEN { + let encoded = object.next_value::()?; + let number = encoded.parse::().map_err(de::Error::custom)?; + return Ok(Value::Number(number)); + } + + if first_key == SERDE_JSON_RAW_VALUE_TOKEN { + let encoded = object.next_value::()?; + let mut deserializer = serde_json::Deserializer::from_str(&encoded); + // The raw wrapper already consumed one value from the budget. Reuse + // that slot for the decoded root while charging all of its children. + let value = (&mut deserializer) + .deserialize_any(BoundedValueVisitor { + remaining: &mut *self.remaining, + }) + .map_err(de::Error::custom)?; + deserializer.end().map_err(de::Error::custom)?; + return Ok(value); + } + let mut values = Map::new(); + let first_value = object.next_value_seed(BoundedValueSeed { + remaining: &mut *self.remaining, + })?; + values.insert(first_key, first_value); + while let Some(key) = object.next_key::()? { if values.contains_key(&key) { return Err(de::Error::custom(format!( diff --git a/codex-rs/exec-server-protocol/src/rpc_tests.rs b/codex-rs/exec-server-protocol/src/rpc_tests.rs index d60f18ef6a3..4ace85be94a 100644 --- a/codex-rs/exec-server-protocol/src/rpc_tests.rs +++ b/codex-rs/exec-server-protocol/src/rpc_tests.rs @@ -1,4 +1,5 @@ use pretty_assertions::assert_eq; +use serde_json::Value; use serde_json::json; use super::JSONRPCError; @@ -9,6 +10,7 @@ use super::JSONRPCRequest; use super::JSONRPCResponse; use super::MAX_JSONRPC_VALUE_NODES; use super::RequestId; +use super::SERDE_JSON_RAW_VALUE_TOKEN; #[test] fn round_trips_every_jsonrpc_message_variant() -> serde_json::Result<()> { @@ -46,6 +48,51 @@ fn round_trips_every_jsonrpc_message_variant() -> serde_json::Result<()> { Ok(()) } +#[test] +fn round_trips_arbitrary_precision_numbers() -> serde_json::Result<()> { + let encoded = r#"{"method":"numbers","params":{"decimal":1.5,"exponent":1e100,"largeInteger":18446744073709551616}}"#; + let expected = serde_json::from_str::(encoded)?; + + let message = serde_json::from_str::(encoded)?; + let actual = serde_json::to_value(message)?; + + assert_eq!(actual, expected); + Ok(()) +} + +#[test] +fn applies_value_limit_to_raw_value_wrapper() -> serde_json::Result<()> { + let encoded = + format!(r#"{{"method":"raw","params":{{"{SERDE_JSON_RAW_VALUE_TOKEN}":"[0,1]"}}}}"#); + let actual = serde_json::from_str::(&encoded)?; + let expected = JSONRPCMessage::Notification(JSONRPCNotification { + method: "raw".to_string(), + params: Some(json!([0, 1])), + }); + assert_eq!(actual, expected); + + let mut wrapped = String::with_capacity(2 * MAX_JSONRPC_VALUE_NODES + 1); + wrapped.push('['); + for index in 0..MAX_JSONRPC_VALUE_NODES { + if index != 0 { + wrapped.push(','); + } + wrapped.push('0'); + } + wrapped.push(']'); + let encoded = + format!(r#"{{"method":"raw","params":{{"{SERDE_JSON_RAW_VALUE_TOKEN}":"{wrapped}"}}}}"#); + + let error = serde_json::from_str::(&encoded) + .expect_err("raw value wrapper should not bypass 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}" + ); + Ok(()) +} + #[test] fn accepts_large_scalar_payload() -> serde_json::Result<()> { let expected = JSONRPCMessage::Notification(JSONRPCNotification {