From 080059c1e9cb987a3b51ae0f6e15c8a513b042bd Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Thu, 9 Jul 2026 10:17:47 -0700 Subject: [PATCH] fix(exec-server): close remote policy decision gaps Co-authored-by: Codex noreply@openai.com --- codex-rs/core/src/exec.rs | 1 + codex-rs/core/src/sandboxing/mod.rs | 24 +++ codex-rs/core/src/tasks/user_shell.rs | 1 + .../tools/runtimes/shell/unix_escalation.rs | 2 + .../core/src/tools/runtimes/unified_exec.rs | 4 + .../core/src/unified_exec/process_manager.rs | 14 +- .../src/unified_exec/process_manager_tests.rs | 1 + codex-rs/exec-server-protocol/src/protocol.rs | 2 + codex-rs/exec-server/src/client.rs | 35 +++- codex-rs/exec-server/src/local_process.rs | 17 +- .../src/network_policy_decisions.rs | 79 +++++++-- .../src/network_policy_decisions_tests.rs | 132 +++++++++++++++ codex-rs/exec-server/src/process.rs | 8 +- codex-rs/exec-server/src/rpc.rs | 18 +++ codex-rs/exec-server/src/server/processor.rs | 153 ++++++++++++++++++ codex-rs/exec-server/src/server/registry.rs | 2 +- codex-rs/network-proxy/src/proxy.rs | 95 ++++++++++- 17 files changed, 549 insertions(+), 39 deletions(-) diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 47dd317de1..5f03f01f41 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -462,6 +462,7 @@ pub(crate) async fn execute_exec_request( exec_server_enforce_managed_network: _, exec_server_managed_network: _, exec_server_network_proxy: _, + network_policy_decider: _, } = exec_request; // TODO(anp): Keep PathUri through the local process launch boundary. diff --git a/codex-rs/core/src/sandboxing/mod.rs b/codex-rs/core/src/sandboxing/mod.rs index 9a59477709..9ff37bd01b 100644 --- a/codex-rs/core/src/sandboxing/mod.rs +++ b/codex-rs/core/src/sandboxing/mod.rs @@ -16,6 +16,7 @@ use crate::spawn::CODEX_SANDBOX_ENV_VAR; use crate::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR; use codex_file_system::FileSystemSandboxContext; use codex_network_proxy::ManagedNetworkSandboxContext; +use codex_network_proxy::NetworkPolicyDecider; use codex_network_proxy::NetworkProxy; use codex_network_proxy::RemoteNetworkProxyLaunchConfig; use codex_protocol::config_types::WindowsSandboxLevel; @@ -30,6 +31,7 @@ use codex_sandboxing::WindowsSandboxFilesystemOverrides; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use std::collections::HashMap; +use std::sync::Arc; #[derive(Debug)] pub(crate) struct ExecOptions { @@ -67,6 +69,26 @@ pub struct ExecRequest { pub(crate) exec_server_enforce_managed_network: bool, pub(crate) exec_server_managed_network: Option, pub(crate) exec_server_network_proxy: Option, + pub(crate) network_policy_decider: Option, +} + +#[derive(Clone)] +pub(crate) struct ExecNetworkPolicyDecider(Arc); + +impl ExecNetworkPolicyDecider { + pub(crate) fn new(decider: Arc) -> Self { + Self(decider) + } + + pub(crate) fn get(&self) -> Arc { + Arc::clone(&self.0) + } +} + +impl std::fmt::Debug for ExecNetworkPolicyDecider { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.debug_struct("ExecNetworkPolicyDecider").finish() + } } impl ExecRequest { @@ -113,6 +135,7 @@ impl ExecRequest { exec_server_enforce_managed_network: false, exec_server_managed_network: None, exec_server_network_proxy: None, + network_policy_decider: None, } } @@ -173,6 +196,7 @@ impl ExecRequest { exec_server_enforce_managed_network: false, exec_server_managed_network: None, exec_server_network_proxy: None, + network_policy_decider: None, } } } diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index ba40893bcc..0dc27272cc 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -232,6 +232,7 @@ pub(crate) async fn execute_user_shell_command( exec_server_enforce_managed_network: false, exec_server_managed_network: None, exec_server_network_proxy: None, + network_policy_decider: None, }; let stdout_stream = Some(StdoutStream { diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 42ffd3e9aa..b5003374c1 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -170,6 +170,7 @@ pub(super) async fn try_run_zsh_fork( exec_server_enforce_managed_network: _, exec_server_managed_network: _, exec_server_network_proxy: _, + network_policy_decider: _, } = sandbox_exec_request; let ParsedShellCommand { script, login, .. } = extract_shell_script(&command)?; let effective_timeout = Duration::from_millis( @@ -906,6 +907,7 @@ impl CoreShellCommandExecutor { exec_server_enforce_managed_network: false, exec_server_managed_network: None, exec_server_network_proxy: None, + network_policy_decider: None, }, /*stdout_stream*/ None, after_spawn, diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index e33ea42c2e..aa5d3363af 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -323,6 +323,9 @@ impl<'a> ToolRuntime for UnifiedExecRunt req.network.as_ref(), launch_sandbox_permissions, )); + let network_policy_decider = environment_is_remote + .then(|| managed_network.and_then(NetworkProxy::remote_policy_decider)) + .flatten(); let env = exec_env_for_sandbox_permissions(&req.env, launch_sandbox_permissions); let (env, managed_network_context, network_proxy_launch) = match managed_network { Some(network) if environment_is_remote => { @@ -484,6 +487,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt attempt, managed_network, network_proxy_launch, + network_policy_decider, /*environment_id*/ Some(&req.turn_environment.environment_id), req.exec_server_env_config.clone(), req.tty, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 6213d437a4..de693cc2d9 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -18,6 +18,7 @@ use crate::exec_env::CODEX_THREAD_ID_ENV_VAR; use crate::exec_env::create_env; use crate::exec_env::inject_permission_profile_env; use crate::exec_policy::ExecApprovalRequest; +use crate::sandboxing::ExecNetworkPolicyDecider; use crate::sandboxing::ExecOptions; use crate::sandboxing::ExecRequest; use crate::sandboxing::ExecServerEnvConfig; @@ -917,6 +918,7 @@ impl UnifiedExecProcessManager { attempt: &SandboxAttempt<'_>, network: Option<&NetworkProxy>, network_proxy_launch: Option, + network_policy_decider: Option>, environment_id: Option<&str>, exec_server_env_config: Option, tty: bool, @@ -931,9 +933,7 @@ impl UnifiedExecProcessManager { } .map_err(ToolError::Codex)?; request.exec_server_network_proxy = network_proxy_launch; - if remote { - request.network = network.cloned(); - } + request.network_policy_decider = network_policy_decider.map(ExecNetworkPolicyDecider::new); request.exec_server_env_config = exec_server_env_config; self.open_session_with_prepared_exec_env( process_id, @@ -1059,14 +1059,10 @@ impl UnifiedExecProcessManager { let backend = environment.get_exec_backend(); let params = exec_server_params_for_request(process_id, request, tty); - let started = match request - .network - .as_ref() - .and_then(NetworkProxy::remote_policy_decider) - { + let started = match request.network_policy_decider.as_ref() { Some(decider) => { backend - .start_with_network_policy_decider(params, decider) + .start_with_network_policy_decider(params, decider.get()) .await } None => backend.start(params).await, diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 92c4c70e21..c129c53efb 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -175,6 +175,7 @@ fn exec_server_params_use_path_uri_and_env_policy_overlay_contract() { exec_server_enforce_managed_network: true, exec_server_managed_network: Some(managed_network.clone()), exec_server_network_proxy: None, + network_policy_decider: None, }; let params = diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index f9bfe69208..64251a850b 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::time::Duration; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_file_system::FileSystemSandboxContext; @@ -28,6 +29,7 @@ pub const EXEC_EXITED_METHOD: &str = "process/exited"; pub const EXEC_CLOSED_METHOD: &str = "process/closed"; pub const NETWORK_POLICY_REQUEST_METHOD: &str = "network/policyRequest"; pub const NETWORK_POLICY_DECISION_METHOD: &str = "network/policyDecision"; +pub const NETWORK_POLICY_DECISION_TIMEOUT: Duration = Duration::from_secs(30); pub const ENVIRONMENT_INFO_METHOD: &str = "environment/info"; pub const FS_READ_FILE_METHOD: &str = "fs/readFile"; pub const FS_OPEN_METHOD: &str = "fs/open"; diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 6d78943d5f..6093dd94ee 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -19,6 +19,7 @@ use tokio::sync::Mutex; use tokio::sync::OnceCell; use tokio::sync::mpsc; use tokio::sync::watch; +use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; use tokio::time::timeout; @@ -94,6 +95,7 @@ use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; use crate::protocol::NETWORK_POLICY_DECISION_METHOD; +use crate::protocol::NETWORK_POLICY_DECISION_TIMEOUT; use crate::protocol::NETWORK_POLICY_REQUEST_METHOD; use crate::protocol::NetworkPolicyDecisionNotification; use crate::protocol::NetworkPolicyRequestNotification; @@ -169,6 +171,7 @@ pub(crate) struct SessionState { recoverable: AtomicBool, next_write_id: AtomicU64, network_policy_decider: Option>, + network_policy_decisions_cancelled: CancellationToken, } #[derive(Default)] @@ -880,6 +883,7 @@ impl SessionState { recoverable: AtomicBool::new(recoverable), next_write_id: AtomicU64::new(1), network_policy_decider, + network_policy_decisions_cancelled: CancellationToken::new(), } } @@ -936,6 +940,9 @@ impl SessionState { ordered_events.last_published_seq = next_seq; ordered_events.exit_published |= matches!(&event, ExecProcessEvent::Exited { .. }); let is_closed = matches!(&event, ExecProcessEvent::Closed { .. }); + if matches!(&event, ExecProcessEvent::Exited { .. }) || is_closed { + self.network_policy_decisions_cancelled.cancel(); + } ordered_events.closed_published |= is_closed; published_closed |= is_closed; self.events.publish(event); @@ -944,6 +951,7 @@ impl SessionState { } fn set_failure(&self, message: String) { + self.network_policy_decisions_cancelled.cancel(); let mut ordered_events = self .ordered_events .lock() @@ -1123,6 +1131,7 @@ impl Inner { let mut next_sessions = sessions.as_ref().clone(); next_sessions.remove(process_id); self.sessions.store(Arc::new(next_sessions)); + expected.network_policy_decisions_cancelled.cancel(); } fn take_all_sessions(&self) -> HashMap> { @@ -1233,15 +1242,29 @@ async fn handle_server_notification( NETWORK_POLICY_REQUEST_METHOD => { let params: NetworkPolicyRequestNotification = serde_json::from_value(notification.params.unwrap_or(Value::Null))?; - let decider = inner - .get_session(¶ms.process_id) - .and_then(|session| session.network_policy_decider.clone()); + let (decider, cancellation) = inner.get_session(¶ms.process_id).map_or_else( + || (None, CancellationToken::new()), + |session| { + ( + session.network_policy_decider.clone(), + session.network_policy_decisions_cancelled.clone(), + ) + }, + ); let inner = Arc::clone(inner); tokio::spawn(async move { - let decision = match decider { - Some(decider) => decider.decide(params.request).await, - None => NetworkDecision::deny("not_allowed"), + let decision = tokio::select! { + _ = cancellation.cancelled() => return, + decision = timeout(NETWORK_POLICY_DECISION_TIMEOUT, async { + match decider { + Some(decider) => decider.decide(params.request).await, + None => NetworkDecision::deny("not_allowed"), + } + }) => decision.unwrap_or_else(|_| NetworkDecision::deny("not_allowed")), }; + if cancellation.is_cancelled() { + return; + } let response = NetworkPolicyDecisionNotification { request_id: params.request_id, process_id: params.process_id, diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 383c4aec45..15e50019b9 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -257,10 +257,9 @@ impl LocalProcess { ) -> Result<(ExecResponse, watch::Sender, ExecProcessEventLog), JSONRPCErrorError> { let process_id = params.process_id.clone(); let network_policy_decider = params - .managed_network + .network_proxy .as_ref() - .and_then(|context| context.proxy_config.as_ref()) - .is_some_and(|config| config.request_policy_decisions) + .is_some_and(|launch| launch.proxy.request_policy_decisions) .then(|| { self.inner .network_policy_decisions @@ -662,6 +661,18 @@ impl ExecBackend for LocalProcess { fn start(&self, params: ExecParams) -> ExecBackendFuture<'_> { Box::pin(LocalProcess::start(self, params)) } + + fn start_with_network_policy_decider( + &self, + _params: ExecParams, + _decider: Arc, + ) -> ExecBackendFuture<'_> { + Box::pin(async { + Err(ExecServerError::Protocol( + "local exec backend does not accept a remote network policy decider".to_string(), + )) + }) + } } impl LocalExecProcess { diff --git a/codex-rs/exec-server/src/network_policy_decisions.rs b/codex-rs/exec-server/src/network_policy_decisions.rs index 1695cc5baf..ef3f395ceb 100644 --- a/codex-rs/exec-server/src/network_policy_decisions.rs +++ b/codex-rs/exec-server/src/network_policy_decisions.rs @@ -2,30 +2,66 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; use std::sync::RwLock; +use std::time::Duration; use codex_network_proxy::NetworkDecision; use codex_network_proxy::NetworkPolicyDecider; use codex_network_proxy::NetworkPolicyRequest; use tokio::sync::oneshot; +use tokio::time::timeout; use uuid::Uuid; use crate::ProcessId; +use crate::protocol::NETWORK_POLICY_DECISION_TIMEOUT; use crate::protocol::NETWORK_POLICY_REQUEST_METHOD; use crate::protocol::NetworkPolicyDecisionNotification; use crate::protocol::NetworkPolicyRequestNotification; use crate::rpc::RpcNotificationSender; -#[derive(Default)] pub(crate) struct NetworkPolicyDecisionRelay { pending: Mutex>, + decision_timeout: Duration, + max_pending: usize, } +const MAX_PENDING_NETWORK_POLICY_DECISIONS: usize = 256; + struct PendingNetworkPolicyDecision { process_id: ProcessId, response_tx: oneshot::Sender, } +struct PendingNetworkPolicyDecisionGuard<'a> { + relay: &'a NetworkPolicyDecisionRelay, + request_id: String, +} + +impl Drop for PendingNetworkPolicyDecisionGuard<'_> { + fn drop(&mut self) { + self.relay.fail(&self.request_id); + } +} + +impl Default for NetworkPolicyDecisionRelay { + fn default() -> Self { + Self { + pending: Mutex::new(HashMap::new()), + decision_timeout: NETWORK_POLICY_DECISION_TIMEOUT, + max_pending: MAX_PENDING_NETWORK_POLICY_DECISIONS, + } + } +} + impl NetworkPolicyDecisionRelay { + #[cfg(test)] + fn with_limits(decision_timeout: Duration, max_pending: usize) -> Self { + Self { + pending: Mutex::new(HashMap::new()), + decision_timeout, + max_pending, + } + } + pub(crate) fn decider( self: &Arc, process_id: ProcessId, @@ -78,16 +114,26 @@ impl NetworkPolicyDecisionRelay { ) -> NetworkDecision { let request_id = Uuid::new_v4().to_string(); let (response_tx, response_rx) = oneshot::channel(); - self.pending - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert( + { + let mut pending = self + .pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if pending.len() >= self.max_pending { + return NetworkDecision::deny("not_allowed"); + } + pending.insert( request_id.clone(), PendingNetworkPolicyDecision { process_id: process_id.clone(), response_tx, }, ); + } + let _pending = PendingNetworkPolicyDecisionGuard { + relay: self, + request_id: request_id.clone(), + }; let notification = NetworkPolicyRequestNotification { request_id: request_id.clone(), process_id, @@ -98,20 +144,19 @@ impl NetworkPolicyDecisionRelay { .unwrap_or_else(std::sync::PoisonError::into_inner) .clone() else { - self.fail(&request_id); return NetworkDecision::deny("not_allowed"); }; - if notifications - .notify(NETWORK_POLICY_REQUEST_METHOD, ¬ification) - .await - .is_err() - { - self.fail(&request_id); - return NetworkDecision::deny("not_allowed"); - } - response_rx - .await - .unwrap_or_else(|_| NetworkDecision::deny("not_allowed")) + timeout(self.decision_timeout, async { + notifications + .notify(NETWORK_POLICY_REQUEST_METHOD, ¬ification) + .await + .ok()?; + response_rx.await.ok() + }) + .await + .ok() + .flatten() + .unwrap_or_else(|| NetworkDecision::deny("not_allowed")) } fn fail(&self, request_id: &str) { diff --git a/codex-rs/exec-server/src/network_policy_decisions_tests.rs b/codex-rs/exec-server/src/network_policy_decisions_tests.rs index 8f51a79cb1..6b2e5ce99c 100644 --- a/codex-rs/exec-server/src/network_policy_decisions_tests.rs +++ b/codex-rs/exec-server/src/network_policy_decisions_tests.rs @@ -62,6 +62,25 @@ fn network_policy_relay() -> ( ) } +fn bounded_network_policy_relay( + decision_timeout: Duration, + max_pending: usize, +) -> ( + Arc, + Arc>>, + mpsc::Receiver, +) { + let (outgoing_tx, outgoing_rx) = mpsc::channel(NOTIFICATION_CHANNEL_CAPACITY); + ( + Arc::new(NetworkPolicyDecisionRelay::with_limits( + decision_timeout, + max_pending, + )), + Arc::new(RwLock::new(Some(RpcNotificationSender::new(outgoing_tx)))), + outgoing_rx, + ) +} + fn request_decision( relay: &Arc, notifications: &Arc>>, @@ -196,3 +215,116 @@ async fn process_exit_fails_only_its_decisions() { .expect("resolve remaining process decision"); assert_eq!(await_decision(decision_b).await, NetworkDecision::Allow); } + +#[tokio::test] +async fn pending_decisions_are_capacity_bounded() { + let (relay, notifications, mut outgoing_rx) = + bounded_network_policy_relay(Duration::from_secs(1), /*max_pending*/ 1); + let process_id = ProcessId::from("network-decision-capacity"); + let first = request_decision( + &relay, + ¬ifications, + process_id.clone(), + "first.example.com", + ); + let first_request = next_network_policy_request(&mut outgoing_rx).await; + + let second = request_decision( + &relay, + ¬ifications, + process_id.clone(), + "second.example.com", + ); + assert_eq!( + await_decision(second).await, + NetworkDecision::deny("not_allowed") + ); + assert!( + timeout(Duration::from_millis(20), outgoing_rx.recv()) + .await + .is_err(), + "a rejected request must not emit a notification" + ); + + relay + .resolve(NetworkPolicyDecisionNotification { + request_id: first_request.request_id, + process_id, + decision: NetworkDecision::Allow, + }) + .expect("resolve first network policy decision"); + assert_eq!(await_decision(first).await, NetworkDecision::Allow); +} + +#[tokio::test] +async fn pending_decisions_time_out_and_release_capacity() { + let (relay, notifications, mut outgoing_rx) = + bounded_network_policy_relay(Duration::from_millis(20), /*max_pending*/ 1); + let process_id = ProcessId::from("network-decision-timeout"); + let timed_out = request_decision( + &relay, + ¬ifications, + process_id.clone(), + "timeout.example.com", + ); + let timed_out_request = next_network_policy_request(&mut outgoing_rx).await; + assert_eq!( + await_decision(timed_out).await, + NetworkDecision::deny("not_allowed") + ); + relay + .resolve(NetworkPolicyDecisionNotification { + request_id: timed_out_request.request_id, + process_id: process_id.clone(), + decision: NetworkDecision::Allow, + }) + .expect("late decisions are ignored"); + + let next = request_decision( + &relay, + ¬ifications, + process_id.clone(), + "next.example.com", + ); + let next_request = next_network_policy_request(&mut outgoing_rx).await; + relay + .resolve(NetworkPolicyDecisionNotification { + request_id: next_request.request_id, + process_id, + decision: NetworkDecision::Allow, + }) + .expect("resolve next network policy decision"); + assert_eq!(await_decision(next).await, NetworkDecision::Allow); +} + +#[tokio::test] +async fn cancelling_a_pending_decision_releases_capacity() { + let (relay, notifications, mut outgoing_rx) = + bounded_network_policy_relay(Duration::from_secs(1), /*max_pending*/ 1); + let process_id = ProcessId::from("network-decision-cancelled"); + let cancelled = request_decision( + &relay, + ¬ifications, + process_id.clone(), + "cancelled.example.com", + ); + next_network_policy_request(&mut outgoing_rx).await; + cancelled.abort(); + let _ = cancelled.await; + + let next = request_decision( + &relay, + ¬ifications, + process_id.clone(), + "next.example.com", + ); + let next_request = next_network_policy_request(&mut outgoing_rx).await; + relay + .resolve(NetworkPolicyDecisionNotification { + request_id: next_request.request_id, + process_id, + decision: NetworkDecision::Allow, + }) + .expect("resolve next network policy decision"); + assert_eq!(await_decision(next).await, NetworkDecision::Allow); +} diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index d3bb47e133..cd612e7dd6 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/codex-rs/exec-server/src/process.rs @@ -205,10 +205,14 @@ pub trait ExecBackend: Send + Sync { fn start_with_network_policy_decider( &self, - params: ExecParams, + _params: ExecParams, _decider: Arc, ) -> ExecBackendFuture<'_> { - self.start(params) + Box::pin(async { + Err(ExecServerError::Protocol( + "exec backend does not support remote network policy decisions".to_string(), + )) + }) } } diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index e762fe8b66..6dd1645ceb 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::collections::HashSet; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -115,6 +116,7 @@ impl RpcNotificationSender { pub(crate) struct RpcRouter { request_routes: HashMap<&'static str, RequestRoute>, + concurrent_request_routes: HashSet<&'static str>, notification_routes: HashMap<&'static str, NotificationRoute>, } @@ -122,6 +124,7 @@ impl Default for RpcRouter { fn default() -> Self { Self { request_routes: HashMap::new(), + concurrent_request_routes: HashSet::new(), notification_routes: HashMap::new(), } } @@ -199,6 +202,17 @@ where ); } + pub(crate) fn concurrent_request(&mut self, method: &'static str, handler: F) + where + P: DeserializeOwned + Send + 'static, + R: Serialize + Send + 'static, + F: Fn(Arc, P) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.request(method, handler); + self.concurrent_request_routes.insert(method); + } + pub(crate) fn notification(&mut self, method: &'static str, handler: F) where P: DeserializeOwned + Send + 'static, @@ -227,6 +241,10 @@ where .map(|(&method, route)| (method, route)) } + pub(crate) fn request_runs_concurrently(&self, method: &str) -> bool { + self.concurrent_request_routes.contains(method) + } + pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { self.notification_routes.get(method) } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 2944c78041..e92ab31c1b 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,7 +1,10 @@ use std::sync::Arc; use std::time::Instant; +use tokio::sync::Semaphore; use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; use tracing::Instrument; use tracing::debug; use tracing::warn; @@ -21,6 +24,8 @@ use crate::server::session_registry::SessionRegistry; use crate::telemetry::ConnectionTransport; use crate::telemetry::ExecServerTelemetry; +const MAX_CONCURRENT_REQUESTS: usize = 256; + #[derive(Clone)] pub(crate) struct ConnectionProcessor { session_registry: Arc, @@ -89,6 +94,9 @@ async fn run_connection( notifications, runtime_paths, )); + let concurrent_request_shutdown = CancellationToken::new(); + let concurrent_request_tasks = TaskTracker::new(); + let concurrent_request_permits = Arc::new(Semaphore::new(MAX_CONCURRENT_REQUESTS)); let outbound_task = tokio::spawn(async move { while let Some(message) = outgoing_rx.recv().await { @@ -130,6 +138,79 @@ async fn run_connection( let request_started_at = Instant::now(); if let Some((method, route)) = router.request_route(request.method.as_str()) { let request_span = request_span(method, &request); + if router.request_runs_concurrently(method) { + let permit = match Arc::clone(&concurrent_request_permits) + .try_acquire_owned() + { + Ok(permit) => permit, + Err(_) => { + let result = "error"; + let message = RpcServerOutboundMessage::Error { + request_id: request.id, + error: invalid_request( + "too many concurrent exec-server requests".to_string(), + ), + }; + if outgoing_tx.send(message).await.is_err() { + request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); + break; + } + request_span.record("result", result); + telemetry.request_completed( + method, + result, + request_started_at.elapsed(), + ); + continue; + } + }; + let request = route(Arc::clone(&handler), request) + .instrument(request_span.clone()); + let outgoing_tx = outgoing_tx.clone(); + let mut disconnected_rx = disconnected_rx.clone(); + let shutdown = concurrent_request_shutdown.clone(); + let telemetry = telemetry.clone(); + concurrent_request_tasks.spawn(async move { + let _permit = permit; + let message = tokio::select! { + message = request => message, + _ = disconnected_rx.changed() => { + request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); + return; + } + _ = shutdown.cancelled() => return, + }; + let result = request_result(&message); + if let Some(message) = message + && outgoing_tx.send(message).await.is_err() + { + request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); + return; + } + request_span.record("result", result); + telemetry.request_completed( + method, + result, + request_started_at.elapsed(), + ); + }); + continue; + } let message = tokio::select! { message = route(Arc::clone(&handler), request).instrument(request_span.clone()) => message, _ = disconnected_rx.changed() => { @@ -231,6 +312,9 @@ async fn run_connection( } } + concurrent_request_shutdown.cancel(); + concurrent_request_tasks.close(); + concurrent_request_tasks.wait().await; handler.shutdown().await; drop(handler); drop(outgoing_tx); @@ -318,6 +402,8 @@ mod tests { use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; + use crate::protocol::NETWORK_POLICY_DECISION_METHOD; + use crate::protocol::NetworkPolicyDecisionNotification; use crate::protocol::ReadParams; use crate::protocol::TerminateParams; use crate::protocol::TerminateResponse; @@ -496,6 +582,73 @@ mod tests { .expect("second processor should join"); } + #[tokio::test] + async fn network_policy_decision_is_not_queued_behind_long_poll_read() { + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); + let (mut writer, mut lines, task) = + spawn_test_connection(registry, "network-policy-long-poll"); + + send_request( + &mut writer, + /*id*/ 1, + INITIALIZE_METHOD, + &InitializeParams { + client_name: "exec-server-test".to_string(), + resume_session_id: None, + }, + ) + .await; + let _: InitializeResponse = read_response(&mut lines, /*expected_id*/ 1).await; + send_notification(&mut writer, INITIALIZED_METHOD, &()).await; + + let process_id = ProcessId::from("proc-policy-long-poll"); + send_request( + &mut writer, + /*id*/ 2, + EXEC_METHOD, + &exec_params(process_id.clone()), + ) + .await; + let _: ExecResponse = read_response(&mut lines, /*expected_id*/ 2).await; + send_request( + &mut writer, + /*id*/ 3, + EXEC_READ_METHOD, + &ReadParams { + process_id: process_id.clone(), + after_seq: Some(0), + max_bytes: None, + wait_ms: Some(5_000), + }, + ) + .await; + send_notification( + &mut writer, + NETWORK_POLICY_DECISION_METHOD, + &NetworkPolicyDecisionNotification { + request_id: "already-completed".to_string(), + process_id, + decision: codex_network_proxy::NetworkDecision::Allow, + }, + ) + .await; + send_request(&mut writer, /*id*/ 4, ENVIRONMENT_INFO_METHOD, &()).await; + + let _: EnvironmentInfo = timeout( + Duration::from_millis(250), + read_response(&mut lines, /*expected_id*/ 4), + ) + .await + .expect("policy decision and following requests must bypass long-poll read"); + + drop(writer); + drop(lines); + timeout(Duration::from_secs(1), task) + .await + .expect("processor should exit") + .expect("processor should join"); + } + fn spawn_test_connection( registry: Arc, label: &str, diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index b288a335fe..aebda53682 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -79,7 +79,7 @@ pub(crate) fn build_router() -> RpcRouter { ENVIRONMENT_INFO_METHOD, |handler: Arc, _params: ()| async move { handler.environment_info() }, ); - router.request( + router.concurrent_request( EXEC_READ_METHOD, |handler: Arc, params: ReadParams| async move { handler.exec_read(params).await diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index 7860448fc5..4724321210 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -6,6 +6,8 @@ use crate::credential_broker::BROKERED_CREDENTIALS_ENV_KEY; use crate::credential_broker::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; use crate::http_proxy; use crate::network_policy::NetworkPolicyDecider; +use crate::runtime::BlockedRequest; +use crate::runtime::BlockedRequestArgs; use crate::runtime::BlockedRequestObserver; use crate::runtime::ConfigState; use crate::runtime::unix_socket_permissions_supported; @@ -237,6 +239,7 @@ impl NetworkProxyBuilder { )?)), reserved_listeners, policy_decider: self.policy_decider, + blocked_request_observer: self.blocked_request_observer, environment_proxies: Arc::new(Mutex::new(HashMap::new())), execution_scope: None, }) @@ -375,6 +378,7 @@ pub struct NetworkProxy { runtime_settings: Arc>, reserved_listeners: Option>, policy_decider: Option>, + blocked_request_observer: Option>, environment_proxies: Arc>>, execution_scope: Option>, } @@ -687,6 +691,7 @@ impl NetworkProxy { /// Returns the policy decider with this proxy's execution attribution attached. pub fn remote_policy_decider(&self) -> Option> { let decider = self.policy_decider.clone()?; + let blocked_request_observer = self.blocked_request_observer.clone(); let environment_id = self .execution_scope .as_ref() @@ -697,9 +702,18 @@ impl NetworkProxy { .map(|scope| scope.execution_id.clone()); Some(Arc::new(move |mut request: crate::NetworkPolicyRequest| { let decider = Arc::clone(&decider); + let blocked_request_observer = blocked_request_observer.clone(); request.environment_id = request.environment_id.or_else(|| environment_id.clone()); request.execution_id = request.execution_id.or_else(|| execution_id.clone()); - async move { decider.decide(request).await } + async move { + let decision = decider.decide(request.clone()).await; + if let Some(observer) = blocked_request_observer + && let Some(blocked) = remote_blocked_request(&request, &decision) + { + observer.on_blocked_request(blocked).await; + } + decision + } })) } @@ -1051,6 +1065,33 @@ impl NetworkProxy { } } +fn remote_blocked_request( + request: &crate::NetworkPolicyRequest, + decision: &crate::NetworkDecision, +) -> Option { + let crate::NetworkDecision::Deny { + reason, + source, + decision, + } = decision + else { + return None; + }; + let mut blocked = BlockedRequest::new(BlockedRequestArgs { + host: request.host.clone(), + reason: reason.clone(), + client: request.client_addr.clone(), + method: request.method.clone(), + mode: None, + protocol: request.protocol.as_policy_protocol().to_string(), + decision: Some(decision.as_str().to_string()), + source: Some(source.as_str().to_string()), + port: Some(request.port), + }); + blocked.execution_id = request.execution_id.clone(); + Some(blocked) +} + pub struct NetworkProxyHandle { http_task: Option>>, socks_task: Option>>, @@ -1195,6 +1236,58 @@ mod tests { ); } + #[tokio::test] + async fn remote_policy_decider_reports_denials_to_blocked_request_observer() { + let captured = Arc::new(Mutex::new(None)); + let captured_for_observer = Arc::clone(&captured); + let state = Arc::new(network_proxy_state_for_policy( + NetworkProxySettings::default(), + )); + let proxy = NetworkProxy::builder() + .state(state) + .policy_decider(|_request: crate::NetworkPolicyRequest| async { + NetworkDecision::deny("not_allowed") + }) + .blocked_request_observer(move |blocked: BlockedRequest| { + let captured = Arc::clone(&captured_for_observer); + async move { + *captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(blocked); + } + }) + .build() + .await + .expect("build proxy"); + let scoped = proxy + .for_execution("remote", "execution-1", "attribution-token".to_string()) + .expect("scope proxy to execution"); + let decider = scoped.remote_policy_decider().expect("remote decider"); + let request = crate::NetworkPolicyRequest::new(crate::NetworkPolicyRequestArgs { + protocol: crate::NetworkProtocol::HttpsConnect, + host: "blocked.example".to_string(), + port: 443, + environment_id: None, + client_addr: None, + method: Some("CONNECT".to_string()), + command: None, + exec_policy_hint: None, + }); + + assert_eq!( + decider.decide(request).await, + NetworkDecision::deny("not_allowed") + ); + let blocked = captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .expect("blocked request"); + assert_eq!(blocked.host, "blocked.example"); + assert_eq!(blocked.execution_id.as_deref(), Some("execution-1")); + assert_eq!(blocked.decision.as_deref(), Some("deny")); + } + #[tokio::test] async fn managed_proxy_builder_uses_loopback_ports() { let http_listener = StdTcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))).unwrap();