From 4ba28150146e8ae8e949a07951639a929f5ddeae Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 10 Jul 2026 09:51:08 +0000 Subject: [PATCH] Bound streamed exec-server HTTP response bodies (#32112) ## Why Frame-count backpressure does not bound the amount of executor-controlled data retained in streamed HTTP response queues, and a single body delta could exceed the intended wire size. ## What changed - Limit each decoded `http/request/bodyDelta` payload to 1 MiB, split locally produced response chunks at that boundary, and reject oversized incoming deltas. - Apply a shared 16 MiB byte budget across queued HTTP response streams. Release capacity as deltas are consumed and fail a stream when the budget is exhausted. ## Testing Added coverage for rejecting an oversized delta and for failing a stream after its queued deltas exhaust the shared byte budget. GitOrigin-RevId: be9205fef44b0ec96b84b31a8e059b1cb9cd3f3b --- codex-rs/exec-server-protocol/src/protocol.rs | 2 + codex-rs/exec-server/src/client.rs | 8 +- .../src/client/http_response_body_stream.rs | 100 +++++++-- .../src/client/reqwest_http_client.rs | 31 +-- codex-rs/exec-server/tests/http_client.rs | 204 ++++++++++++++++++ 5 files changed, 315 insertions(+), 30 deletions(-) diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index a265aaed0c3a..caf172381442 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -40,6 +40,8 @@ pub const FS_COPY_METHOD: &str = "fs/copy"; pub const HTTP_REQUEST_METHOD: &str = "http/request"; /// JSON-RPC notification method for streamed executor HTTP response bodies. pub const HTTP_REQUEST_BODY_DELTA_METHOD: &str = "http/request/bodyDelta"; +/// Maximum decoded response-body bytes carried by one streamed HTTP notification. +pub const MAX_HTTP_BODY_DELTA_BYTES: usize = 1024 * 1024; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(transparent)] diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 73e49b2f54aa..6b8e22b7b241 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -15,6 +15,7 @@ use futures::future::BoxFuture; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::OnceCell; +use tokio::sync::Semaphore; use tokio::sync::mpsc; use tokio::sync::watch; use tokio_util::task::AbortOnDropHandle; @@ -24,6 +25,8 @@ use tracing::Instrument; use tracing::debug; use crate::ProcessId; +use crate::client::http_client::response_body_stream::MAX_QUEUED_HTTP_BODY_BYTES; +use crate::client::http_client::response_body_stream::QueuedHttpBodyDelta; use crate::client_api::ExecServerClientConnectOptions; use crate::client_api::ExecServerTransportParams; use crate::client_api::HttpClient; @@ -86,7 +89,6 @@ use crate::protocol::FsWalkResponse; use crate::protocol::FsWriteFileParams; use crate::protocol::FsWriteFileResponse; use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD; -use crate::protocol::HttpRequestBodyDeltaNotification; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; @@ -202,9 +204,10 @@ struct Inner { // because they share the same connection-global notification channel as // process output. Keep the routing table local to the client so higher // layers can consume body chunks like a normal byte stream. - http_body_streams: ArcSwap>>, + http_body_streams: ArcSwap>>, http_body_stream_failures: ArcSwap>, http_body_streams_write_lock: Mutex<()>, + http_body_stream_byte_budget: Arc, http_body_stream_next_id: AtomicU64, session_id: OnceLock, reconnect_strategy: Option, @@ -851,6 +854,7 @@ impl ExecServerClient { http_body_streams: ArcSwap::from_pointee(HashMap::new()), http_body_stream_failures: ArcSwap::from_pointee(HashMap::new()), http_body_streams_write_lock: Mutex::new(()), + http_body_stream_byte_budget: Arc::new(Semaphore::new(MAX_QUEUED_HTTP_BODY_BYTES)), http_body_stream_next_id: AtomicU64::new(1), session_id, reconnect_strategy, diff --git a/codex-rs/exec-server/src/client/http_response_body_stream.rs b/codex-rs/exec-server/src/client/http_response_body_stream.rs index 3435af7d59da..08a6655019fc 100644 --- a/codex-rs/exec-server/src/client/http_response_body_stream.rs +++ b/codex-rs/exec-server/src/client/http_response_body_stream.rs @@ -15,6 +15,7 @@ use reqwest::Response; use serde_json::Value; use serde_json::from_value; use tokio::runtime::Handle; +use tokio::sync::OwnedSemaphorePermit; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TrySendError; use tracing::debug; @@ -23,8 +24,29 @@ use crate::client::ExecServerError; use crate::client::Inner; use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD; use crate::protocol::HttpRequestBodyDeltaNotification; +use crate::protocol::MAX_HTTP_BODY_DELTA_BYTES; use crate::rpc::RpcNotificationSender; +pub(crate) const MAX_QUEUED_HTTP_BODY_BYTES: usize = 16 * 1024 * 1024; +const MAX_ENCODED_HTTP_BODY_DELTA_BYTES: usize = MAX_HTTP_BODY_DELTA_BYTES.div_ceil(3) * 4; + +pub(crate) struct QueuedHttpBodyDelta { + notification: HttpRequestBodyDeltaNotification, + _byte_permit: Option, +} + +impl QueuedHttpBodyDelta { + pub(crate) fn new( + notification: HttpRequestBodyDeltaNotification, + byte_permit: Option, + ) -> Self { + Self { + notification, + _byte_permit: byte_permit, + } + } +} + pub(super) struct HttpBodyStreamRegistration { inner: Arc, request_id: String, @@ -39,7 +61,7 @@ enum HttpResponseBodyStreamInner { inner: Arc, request_id: String, next_seq: u64, - rx: mpsc::Receiver, + rx: mpsc::Receiver, pending_eof: bool, closed: bool, }, @@ -66,7 +88,7 @@ impl HttpResponseBodyStream { pub(super) fn remote( inner: Arc, request_id: String, - rx: mpsc::Receiver, + rx: mpsc::Receiver, ) -> Self { Self { inner: HttpResponseBodyStreamInner::Remote { @@ -107,7 +129,11 @@ impl HttpResponseBodyStream { return Ok(None); } - let Some(delta) = rx.recv().await else { + let Some(QueuedHttpBodyDelta { + notification: delta, + .. + }) = rx.recv().await + else { finish_remote_stream(inner, request_id, closed).await; if let Some(error) = inner.take_http_body_stream_failure(request_id).await { return Err(ExecServerError::Protocol(format!( @@ -220,7 +246,22 @@ impl Inner { &self, params: Option, ) -> Result<(), ExecServerError> { - let params: HttpRequestBodyDeltaNotification = from_value(params.unwrap_or(Value::Null))?; + let params = params.unwrap_or(Value::Null); + if params + .get("deltaBase64") + .and_then(Value::as_str) + .is_some_and(|delta| delta.len() > MAX_ENCODED_HTTP_BODY_DELTA_BYTES) + { + return Err(ExecServerError::Protocol(format!( + "http response body delta exceeds {MAX_HTTP_BODY_DELTA_BYTES} bytes" + ))); + } + let params: HttpRequestBodyDeltaNotification = from_value(params)?; + if params.delta.0.len() > MAX_HTTP_BODY_DELTA_BYTES { + return Err(ExecServerError::Protocol(format!( + "http response body delta exceeds {MAX_HTTP_BODY_DELTA_BYTES} bytes" + ))); + } // Unknown request ids are ignored intentionally: a stream may have already // reached EOF and released its route. if let Some(tx) = self @@ -231,7 +272,33 @@ impl Inner { { let request_id = params.request_id.clone(); let terminal_delta = params.done || params.error.is_some(); - match tx.try_send(params) { + let queued_bytes = params + .delta + .0 + .len() + .saturating_add(params.error.as_deref().map_or(0, str::len)); + let byte_permit = if queued_bytes == 0 { + None + } else { + u32::try_from(queued_bytes).ok().and_then(|queued_bytes| { + Arc::clone(&self.http_body_stream_byte_budget) + .try_acquire_many_owned(queued_bytes) + .ok() + }) + }; + if queued_bytes > 0 && byte_permit.is_none() { + self.record_http_body_stream_failure( + &request_id, + format!("queued body deltas exceed {MAX_QUEUED_HTTP_BODY_BYTES} bytes"), + ) + .await; + self.remove_http_body_stream(&request_id).await; + debug!( + "closing http response stream `{request_id}` after exhausting the queued byte budget" + ); + return Ok(()); + } + match tx.try_send(QueuedHttpBodyDelta::new(params, byte_permit)) { Ok(()) => { if terminal_delta { self.remove_http_body_stream(&request_id).await; @@ -265,14 +332,19 @@ impl Inner { let streams = streams.as_ref().clone(); self.http_body_streams.store(Arc::new(HashMap::new())); for (request_id, tx) in streams { + // Failure notifications must wake every stream even when no + // byte-budget permits remain. if tx - .try_send(HttpRequestBodyDeltaNotification { - request_id: request_id.clone(), - seq: 1, - delta: Vec::new().into(), - done: true, - error: Some(message.clone()), - }) + .try_send(QueuedHttpBodyDelta::new( + HttpRequestBodyDeltaNotification { + request_id: request_id.clone(), + seq: 1, + delta: Vec::new().into(), + done: true, + error: Some(message.clone()), + }, + /*byte_permit*/ None, + )) .is_err() { let mut next_failures = self.http_body_stream_failures.load().as_ref().clone(); @@ -295,7 +367,7 @@ impl Inner { pub(super) async fn insert_http_body_stream( &self, request_id: String, - tx: mpsc::Sender, + tx: mpsc::Sender, ) -> Result<(), ExecServerError> { let _streams_write_guard = self.http_body_streams_write_lock.lock().await; let streams = self.http_body_streams.load(); @@ -321,7 +393,7 @@ impl Inner { pub(super) async fn remove_http_body_stream( &self, request_id: &str, - ) -> Option> { + ) -> Option> { let _streams_write_guard = self.http_body_streams_write_lock.lock().await; let streams = self.http_body_streams.load(); let stream = streams.get(request_id).cloned(); diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index d0cffa6a3247..47205874f801 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -30,6 +30,7 @@ use crate::protocol::HttpRedirectPolicy; use crate::protocol::HttpRequestBodyDeltaNotification; use crate::protocol::HttpRequestParams; use crate::protocol::HttpRequestResponse; +use crate::protocol::MAX_HTTP_BODY_DELTA_BYTES; use crate::rpc::RpcNotificationSender; use crate::rpc::internal_error; use crate::rpc::invalid_params; @@ -222,21 +223,23 @@ impl ReqwestHttpRequestRunner { while let Some(chunk) = body.next().await { match chunk { Ok(bytes) => { - if !send_body_delta( - ¬ifications, - HttpRequestBodyDeltaNotification { - request_id: request_id.clone(), - seq, - delta: bytes.to_vec().into(), - done: false, - error: None, - }, - ) - .await - { - return; + for chunk in bytes.chunks(MAX_HTTP_BODY_DELTA_BYTES) { + if !send_body_delta( + ¬ifications, + HttpRequestBodyDeltaNotification { + request_id: request_id.clone(), + seq, + delta: chunk.to_vec().into(), + done: false, + error: None, + }, + ) + .await + { + return; + } + seq += 1; } - seq += 1; } Err(error) => { let _ = send_body_delta( diff --git a/codex-rs/exec-server/tests/http_client.rs b/codex-rs/exec-server/tests/http_client.rs index b3f0182c63d1..6458e8dce58d 100644 --- a/codex-rs/exec-server/tests/http_client.rs +++ b/codex-rs/exec-server/tests/http_client.rs @@ -17,6 +17,7 @@ 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::MAX_HTTP_BODY_DELTA_BYTES; use codex_exec_server_protocol::RequestId; use futures::SinkExt; use futures::StreamExt; @@ -44,6 +45,7 @@ const INITIALIZE_METHOD: &str = "initialize"; const INITIALIZED_METHOD: &str = "initialized"; const TEST_TIMEOUT: Duration = Duration::from_secs(5); const HTTP_BODY_DELTA_CHANNEL_CAPACITY: u64 = 256; +const HTTP_BODY_DELTA_BYTE_BUDGET: usize = 16 * 1024 * 1024; const OVERFLOWING_BODY_DELTA_FRAMES: u64 = 1_024; /// What this tests: the buffered HTTP helper always sends a buffered @@ -764,6 +766,208 @@ async fn http_response_body_stream_fails_when_transport_disconnects() -> Result< Ok(()) } +/// What this tests: an executor cannot make the orchestrator decode and retain +/// a body frame larger than the response-stream wire contract allows. +#[tokio::test] +async fn http_response_body_stream_rejects_oversized_delta() -> Result<()> { + let (finish_tx, finish_rx) = oneshot::channel(); + let server = spawn_scripted_exec_server(|mut peer| async move { + let (_request_id, params) = peer.read_http_request().await?; + assert_eq!( + params, + HttpRequestParams { + method: "GET".to_string(), + url: "https://example.test/mcp/oversized-delta".to_string(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "http-1".to_string(), + stream_response: true, + } + ); + peer.write_body_delta(HttpRequestBodyDeltaNotification { + request_id: "http-1".to_string(), + seq: 1, + delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES + 1].into(), + done: false, + error: None, + }) + .await?; + finish_rx.await.expect("test should finish server task"); + Ok(()) + }) + .await?; + let client = server.connect_client().await?; + + let request = HttpRequestParams { + method: "GET".to_string(), + url: "https://example.test/mcp/oversized-delta".to_string(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "caller-stream-id".to_string(), + stream_response: false, + }; + let result = timeout(TEST_TIMEOUT, client.http_request_stream(request)) + .await + .context("oversized body delta should close the executor transport")?; + let error = match result { + Ok(_) => bail!("oversized body delta should fail the request"), + Err(error) => error, + }; + let error = error.to_string(); + assert_eq!(error, "exec-server transport disconnected"); + + finish_tx.send(()).expect("server task should stay active"); + drop(client); + server.finish().await?; + Ok(()) +} + +/// What this tests: frame-count backpressure cannot hide an unbounded amount +/// of executor-controlled body bytes across the orchestrator's stream queues. +#[tokio::test] +async fn http_response_body_stream_enforces_queued_byte_budget() -> Result<()> { + let (finish_tx, finish_rx) = oneshot::channel(); + let server = spawn_scripted_exec_server(|mut peer| async move { + let (request_id, params) = peer.read_http_request().await?; + assert_eq!( + params, + HttpRequestParams { + method: "GET".to_string(), + url: "https://example.test/mcp/byte-budget".to_string(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "http-1".to_string(), + stream_response: true, + } + ); + let frame_count = HTTP_BODY_DELTA_BYTE_BUDGET / MAX_HTTP_BODY_DELTA_BYTES + 1; + for seq in 1..=frame_count as u64 { + peer.write_body_delta(HttpRequestBodyDeltaNotification { + request_id: "http-1".to_string(), + seq, + delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES].into(), + done: false, + error: None, + }) + .await?; + } + peer.write_response( + request_id, + HttpRequestResponse { + status: 200, + headers: Vec::new(), + body: Vec::new().into(), + }, + ) + .await?; + + let (barrier_request_id, barrier_params) = peer.read_http_request().await?; + assert_eq!( + barrier_params, + HttpRequestParams { + method: "GET".to_string(), + url: "https://example.test/mcp/byte-budget-barrier".to_string(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "http-2".to_string(), + stream_response: true, + } + ); + peer.write_response( + barrier_request_id, + HttpRequestResponse { + status: 200, + headers: Vec::new(), + body: Vec::new().into(), + }, + ) + .await?; + peer.write_body_delta(HttpRequestBodyDeltaNotification { + request_id: "http-2".to_string(), + seq: 1, + delta: Vec::new().into(), + done: true, + error: None, + }) + .await?; + finish_rx.await.expect("test should finish server task"); + Ok(()) + }) + .await?; + let client = server.connect_client().await?; + + let (_response, mut body_stream) = timeout( + TEST_TIMEOUT, + client.http_request_stream(HttpRequestParams { + method: "GET".to_string(), + url: "https://example.test/mcp/byte-budget".to_string(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "caller-stream-id".to_string(), + stream_response: false, + }), + ) + .await + .context("streamed http/request should return headers")??; + + // Receiving this terminal notification proves the earlier byte-budget + // notifications have all passed through the ordered notification handler. + let (_response, mut barrier_stream) = timeout( + TEST_TIMEOUT, + client.http_request_stream(HttpRequestParams { + method: "GET".to_string(), + url: "https://example.test/mcp/byte-budget-barrier".to_string(), + headers: Vec::new(), + body: None, + timeout_ms: None, + redirect_policy: HttpRedirectPolicy::Follow, + request_id: "caller-barrier-id".to_string(), + stream_response: false, + }), + ) + .await + .context("barrier http/request should return headers")??; + assert_eq!( + timeout(TEST_TIMEOUT, barrier_stream.recv()) + .await + .context("barrier body stream should finish")??, + None + ); + + let mut delivered_bytes = 0; + let error = loop { + match timeout(TEST_TIMEOUT, body_stream.recv()) + .await + .context("queued body stream should finish")? + { + Ok(Some(chunk)) => delivered_bytes += chunk.len(), + Ok(None) => bail!("byte-budget exhaustion should not look like clean EOF"), + Err(error) => break error, + } + }; + assert_eq!(delivered_bytes, HTTP_BODY_DELTA_BYTE_BUDGET); + assert!( + error + .to_string() + .contains("queued body deltas exceed 16777216 bytes") + ); + + finish_tx.send(()).expect("server task should stay active"); + drop(client); + server.finish().await?; + Ok(()) +} + /// What this tests: transport disconnect still records a terminal stream /// failure even when the client-side body-delta queue is already full. #[tokio::test]