github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 12:10:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bound inbound exec-server requests across client transports (#47362)

Browse Source 
## What changed

- Enforce an 8 KiB encoded-message limit for requests, unrecognized notifications, and malformed messages received by exec-server clients over stdio, WebSocket, relay, and Noise relay transports.
- Allow process output, exit, close, and HTTP request body notifications up to 2 MiB. Keep responses and errors subject to existing transport limits, and preserve executor-side request limits.
- Read stdio stderr in chunks of at most 8 KiB.
- Drop outgoing server-request `tracestate` values larger than 512 bytes while preserving `traceparent`.

## Testing

Add regression coverage for oversized requests across transports, malformed messages, requests wrapped in `RawValue`, fragmented Noise messages, large responses, streamed notifications, and trace-context preservation.

GitOrigin-RevId: 3fb87b6dc334296985ef40983d7276acfe9b9796
This commit is contained in:
viyatb-oai
2026-09-22 21:07:20 +00:00
committed by copyberry
parent c6ad6f33eb
commit 54487a5b61
11 changed files with 696 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/client_inbound_request_limit.rs
+40
View File
@@ -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<usize> {
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;
codex-rs/exec-server/src/client_inbound_request_limit_tests.rs
+115
View File
@@ -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::<JSONRPCMessage>(&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::<JSONRPCMessage>(&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::<JSONRPCMessage>(&notification);
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::<JSONRPCMessage>(&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::<JSONRPCMessage>(&response);
assert_eq!(
client_inbound_message_exceeded_limit(
message.as_ref(),
response.len(),
MAX_CLIENT_INBOUND_REQUEST_LEN,
),
None
);
}
codex-rs/exec-server/src/client_transport.rs
+23 -6
View File
@@ -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
codex-rs/exec-server/src/connection.rs
+202 -19
View File
@@ -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<R, W>(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<usize>,
) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
@@ -366,7 +385,25 @@ impl JsonRpcConnection {
if line.trim().is_empty() {
continue;
}
match serde_json::from_str::<JSONRPCMessage>(&line) {
let message = serde_json::from_str::<JSONRPCMessage>(&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<Message, Error = E> + Stream<Item = Result<Message, E>> + 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<T, M, E>(
mut websocket: T,
connection_label: String,
ping_interval: Option<Duration>,
max_inbound_request_len: Option<usize>,
) -> Self
where
T: Sink<M, Error = E> + Stream<Item = Result<M, E>> + 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<JsonRpcWebSocketFrame, serde_json::Error>;
fn parse_jsonrpc_frame(
self,
max_inbound_request_len: Option<usize>,
) -> Result<JsonRpcWebSocketFrame, serde_json::Error>;
fn from_text(text: String) -> Self;
fn ping() -> Self;
}
impl JsonRpcWebSocketMessage for Message {
fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
fn parse_jsonrpc_frame(
self,
max_inbound_request_len: Option<usize>,
) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
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<JsonRpcWebSocketFrame, serde_json::Error> {
fn parse_jsonrpc_frame(
self,
max_inbound_request_len: Option<usize>,
) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
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<JSONRPCMessage, serde_json::Error>,
encoded_len: usize,
max_inbound_request_len: Option<usize>,
) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
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<JsonRpcConnectionEvent>,
disconnected_tx: &watch::Sender<bool>,
@@ -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::<JSONRPCMessage>(&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();
codex-rs/exec-server/src/lib.rs
+3
View File
@@ -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.
codex-rs/exec-server/src/noise_relay/harness.rs
+3 -2
View File
@@ -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<JsonRpcConnectionEvent>,
codex-rs/exec-server/src/noise_relay/message_framing.rs
+45
View File
@@ -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::<u32>();
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<Vec<JSONRPCMessage>, 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;
codex-rs/exec-server/src/noise_relay/message_framing_tests.rs
+64 -2
View File
@@ -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 {
codex-rs/exec-server/src/relay.rs
+107 -8
View File
@@ -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<JSONRPCMessage, ExecServerError> {
let payload = self.into_data()?.payload;
serde_json::from_slice(&payload).map_err(ExecServerError::Json)
}
pub(crate) fn into_handshake_payload(self) -> Result<Vec<u8>, 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::<JSONRPCMessage>(&frame.into_data()?.payload)?,
message
);
drop(connection);
Ok(())
}
codex-rs/exec-server/src/rpc_server_requests.rs
+12 -1
View File
@@ -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<Result<Value, RpcCallError>>;
@@ -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 {
codex-rs/exec-server/src/rpc_server_requests_tests.rs
+82
View File
@@ -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",
&params,
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);