From 462ed19a05917b1ec88a015b2cf6342cfe8f7c85 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Mon, 27 Jul 2026 18:01:36 +0000 Subject: [PATCH] Enable network policy callbacks for remote exec (#35652) ## What changed - Forward remote managed-network policy requests to the controller-side decider when Guardian review is enabled. - Bound callbacks using the configured permission-hook and Guardian review timeouts, restore trusted execution attribution, and recheck live baseline policy before requesting a decision. - Cancel pending decisions when the process or execution scope ends, and propagate EOF across the sandbox proxy bridge with half-closes. ## Testing - Cover callback enablement, strict allowlist behavior, live policy updates, trusted attribution, and execution-scope cancellation. GitOrigin-RevId: 83b47657c8c4546dfe98eec8b272833e96a46b7e --- .../core/src/tools/runtimes/unified_exec.rs | 26 ++- codex-rs/core/src/unified_exec/mod_tests.rs | 4 + .../core/src/unified_exec/process_manager.rs | 34 ++-- .../suite/unified_exec_process_events.rs | 59 ++++++- codex-rs/exec-server/src/client.rs | 67 +++++-- .../src/client/tests/network_policy_tests.rs | 31 ++-- codex-rs/exec-server/src/process.rs | 14 ++ codex-rs/exec-server/src/remote_process.rs | 17 +- codex-rs/hooks/src/engine/mod.rs | 12 ++ codex-rs/hooks/src/registry.rs | 8 + codex-rs/linux-sandbox/src/proxy_routing.rs | 8 +- codex-rs/network-proxy/src/proxy.rs | 166 ++++++++++++++++++ .../src/proxy/execution_scope.rs | 4 + 13 files changed, 398 insertions(+), 52 deletions(-) diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 345a28a08d..a2781fb890 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -7,7 +7,9 @@ the process manager to spawn PTYs once an ExecRequest is prepared. use crate::command_canonicalization::canonicalize_command_for_approval; use crate::exec::ExecCapturePolicy; use crate::exec::ExecExpiration; +use crate::guardian::GUARDIAN_REVIEW_TIMEOUT; use crate::guardian::GuardianNetworkAccessTrigger; +use crate::guardian::routes_approval_to_guardian; use crate::sandboxing::ExecOptions; use crate::sandboxing::ExecServerEnvConfig; use crate::sandboxing::SandboxPermissions; @@ -54,9 +56,13 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use std::collections::HashMap; use std::io; +use std::time::Duration; use tokio_util::sync::CancellationToken; use tracing::error; +// Allow 5s for Guardian cleanup and 5s for controller processing after review. +const REMOTE_NETWORK_POLICY_DECISION_MARGIN: Duration = Duration::from_secs(10); + /// Request payload used by the unified-exec runtime after approvals and /// sandbox preferences have been resolved for the current turn. #[derive(Clone, Debug)] @@ -332,9 +338,26 @@ impl<'a> ToolRuntime for UnifiedExecRunt 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 => { - let launch = network.remote_launch_config().await.map_err(|err| { + let mut launch = network.remote_launch_config().await.map_err(|err| { ToolError::Codex(CodexErr::Io(io::Error::other(err.to_string()))) })?; + if routes_approval_to_guardian(&ctx.turn) + && network.remote_policy_decider().is_some() + { + let timeout = ctx + .session + .hooks() + .max_permission_request_timeout() + .saturating_add(GUARDIAN_REVIEW_TIMEOUT) + .saturating_add(REMOTE_NETWORK_POLICY_DECISION_MARGIN); + launch.policy_decision_timeout_ms = + Some(u64::try_from(timeout.as_millis()).map_err(|_| { + ToolError::Rejected( + "remote network policy decision timeout exceeds protocol limit" + .to_string(), + ) + })?); + } if !launch.proxy.enabled { (env, None, None) } else { @@ -466,6 +489,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt req.process_id, &prepared.exec_request, windows_sandbox_proxy_settings_mode, + /*network_policy_decider*/ None, req.tty, prepared.spawn_lifecycle, req.turn_environment.environment.as_ref(), diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index acf6d7d639..e0b58ea726 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -114,6 +114,7 @@ async fn exec_command_with_tty( process_id, &request, codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, + /*network_policy_decider*/ None, tty, Box::new(NoopSpawnLifecycle), turn.environments @@ -755,6 +756,7 @@ async fn completed_pipe_commands_preserve_exit_code() -> anyhow::Result<()> { /*process_id*/ 1234, &request, codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, + /*network_policy_decider*/ None, /*tty*/ false, Box::new(NoopSpawnLifecycle), &environment, @@ -796,6 +798,7 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul /*process_id*/ 1234, &request, codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, + /*network_policy_decider*/ None, /*tty*/ true, Box::new(NoopSpawnLifecycle), remote_test_env.environment(), @@ -855,6 +858,7 @@ async fn remote_exec_server_rejects_inherited_fd_launches() -> anyhow::Result<() /*process_id*/ 1234, &request, codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, + /*network_policy_decider*/ None, /*tty*/ true, Box::new(TestSpawnLifecycle { inherited_fds: vec![42], diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 098c70f026..dec5b3b9d0 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -59,6 +59,7 @@ use crate::unified_exec::process::OutputHandles; use crate::unified_exec::process::SpawnLifecycleHandle; use crate::unified_exec::process::UnifiedExecProcess; use codex_core_plugins::PluginCommandAttribution; +use codex_network_proxy::NetworkPolicyDecider; use codex_network_proxy::NetworkProxy; use codex_protocol::config_types::ShellEnvironmentPolicy; use codex_protocol::error::CodexErr; @@ -1019,12 +1020,17 @@ impl UnifiedExecProcessManager { attempt.env_for(command, options, network, environment_id) } .map_err(ToolError::Codex)?; + let network_policy_decider = network_proxy_launch + .as_ref() + .filter(|launch| launch.policy_decision_timeout_ms.is_some()) + .and_then(|_| network.and_then(NetworkProxy::remote_policy_decider)); request.exec_server_network_proxy = network_proxy_launch; request.exec_server_env_config = exec_server_env_config; self.open_session_with_prepared_exec_env( process_id, &request, windows_sandbox_proxy_settings_mode, + network_policy_decider, tty, spawn_lifecycle, environment, @@ -1041,11 +1047,13 @@ impl UnifiedExecProcessManager { }) } + #[allow(clippy::too_many_arguments)] pub(crate) async fn open_session_with_prepared_exec_env( &self, process_id: i32, request: &ExecRequest, windows_sandbox_proxy_settings_mode: codex_sandboxing::WindowsSandboxProxySettingsMode, + network_policy_decider: Option>, tty: bool, mut spawn_lifecycle: SpawnLifecycleHandle, environment: &codex_exec_server::Environment, @@ -1059,16 +1067,22 @@ impl UnifiedExecProcessManager { )); } - let started = environment - .get_exec_backend() - .start(exec_server_params_for_request( - process_id, - request, - windows_sandbox_proxy_settings_mode, - tty, - )) - .await - .map_err(|err| UnifiedExecError::create_process(err.to_string()))?; + let backend = environment.get_exec_backend(); + let params = exec_server_params_for_request( + process_id, + request, + windows_sandbox_proxy_settings_mode, + tty, + ); + let started = match network_policy_decider { + Some(decider) => { + backend + .start_with_network_policy_decider(params, decider) + .await + } + None => backend.start(params).await, + } + .map_err(|err| UnifiedExecError::create_process(err.to_string()))?; spawn_lifecycle.after_spawn(); return UnifiedExecProcess::from_exec_server_started(started).await; } diff --git a/codex-rs/core/tests/suite/unified_exec_process_events.rs b/codex-rs/core/tests/suite/unified_exec_process_events.rs index 3ce0b2728c..c8f9f24874 100644 --- a/codex-rs/core/tests/suite/unified_exec_process_events.rs +++ b/codex-rs/core/tests/suite/unified_exec_process_events.rs @@ -2,7 +2,9 @@ use anyhow::Context; use anyhow::Result; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_config::test_support::CloudConfigBundleFixture; use codex_features::Feature; +use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -388,15 +390,17 @@ async fn serve_exec_with_pushed_events( } } -#[test_case(PushedExecScenario::Complete, false ; "complete_event_stream")] -#[test_case(PushedExecScenario::DirectDenied, false ; "direct_sandbox_denial")] -#[test_case(PushedExecScenario::LegacyExit, false ; "legacy_exit_metadata")] -#[test_case(PushedExecScenario::ReplayGap, false ; "truncated_event_replay")] -#[test_case(PushedExecScenario::Complete, true ; "managed_network_uses_executor_proxy_launch")] +#[test_case(PushedExecScenario::Complete, false, false ; "complete_event_stream")] +#[test_case(PushedExecScenario::DirectDenied, false, false ; "direct_sandbox_denial")] +#[test_case(PushedExecScenario::LegacyExit, false, false ; "legacy_exit_metadata")] +#[test_case(PushedExecScenario::ReplayGap, false, false ; "truncated_event_replay")] +#[test_case(PushedExecScenario::Complete, true, true ; "managed_network_uses_executor_proxy_launch")] +#[test_case(PushedExecScenario::Complete, true, false ; "strict_managed_allowlist_omits_policy_callbacks")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn exec_command_consumes_pushed_remote_process_events( scenario: PushedExecScenario, managed_network: bool, + policy_callbacks: bool, ) -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let server = start_mock_server().await; @@ -428,8 +432,23 @@ async fn exec_command_consumes_pushed_remote_process_events( let exec_server = tokio::spawn(serve_exec_with_pushed_events(listener, scenario)); let mut builder = test_codex().with_exec_server_url(exec_server_url); if managed_network { + let cloud_config_bundle = if policy_callbacks { + managed_network_requirements_loader() + } else { + CloudConfigBundleFixture::loader_with_enterprise_requirement( + r#" +[experimental_network] +enabled = true +allow_local_binding = true +managed_allowed_domains_only = true + +[experimental_network.domains] +"allowed.example.com" = "allow" +"#, + ) + }; builder = builder - .with_cloud_config_bundle(managed_network_requirements_loader()) + .with_cloud_config_bundle(cloud_config_bundle) .with_pre_build_hook(|home| { fs::write( home.join("config.toml"), @@ -442,14 +461,30 @@ async fn exec_command_consumes_pushed_remote_process_events( enabled = true mode = "full" allow_local_binding = true + +[features] +hooks = true + +[hooks] + +[[hooks.PermissionRequest]] + +[[hooks.PermissionRequest.hooks]] +type = "command" +command = "unused" +timeout = 900 "#, ) .expect("write managed-network test config"); }); } - let mut builder = builder.with_config(|config| { + let mut builder = builder.with_config(move |config| { config.project_doc_max_bytes = 0; config.use_experimental_unified_exec_tool = true; + if managed_network { + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config.bypass_hook_trust = true; + } config .features .enable(Feature::UnifiedExec) @@ -476,7 +511,11 @@ allow_local_binding = true responsesapi_client_metadata: None, additional_context: Default::default(), thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { - approval_policy: Some(AskForApproval::Never), + approval_policy: Some(if managed_network { + AskForApproval::OnRequest + } else { + AskForApproval::Never + }), sandbox_policy: Some(sandbox_policy), permission_profile, collaboration_mode: Some(codex_protocol::config_types::CollaborationMode { @@ -522,6 +561,10 @@ allow_local_binding = true assert_eq!(params["env"]["HTTP_PROXY"], Value::Null); assert_eq!(params["networkProxy"]["proxy"]["enabled"], true); assert_eq!(params["networkProxy"]["proxy"]["mode"], "full"); + assert_eq!( + params["networkProxy"]["policyDecisionTimeoutMs"].as_u64(), + policy_callbacks.then_some(1_000_000) + ); assert_eq!(params["networkProxy"]["environmentId"], "remote"); assert!(params["networkProxy"]["executionId"].as_str().is_some()); timeout(Duration::from_secs(5), async { diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index e8a6be1bd9..3f825df39d 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -876,7 +876,31 @@ impl ExecServerClient { pub(crate) async fn start_process( &self, params: ExecParams, + network_policy_decider: Option>, ) -> Result { + let policy_decision_timeout_ms = params + .network_proxy + .as_ref() + .and_then(|launch| launch.policy_decision_timeout_ms); + let network_policy_controller = match ( + network_policy_decider.as_ref(), + policy_decision_timeout_ms, + ) { + (None, None) => None, + (Some(decider), Some(timeout_ms)) if timeout_ms > 0 => { + Some(NetworkPolicyDecisionController { + decider: Arc::clone(decider), + timeout: Duration::from_millis(timeout_ms), + }) + } + _ => { + return Err(ExecServerError::Protocol( + "network policy decision callback timeout must match the configured decider and be nonzero" + .to_string(), + )); + } + }; + loop { let rpc_client = self.rpc_client().await?; if !self.inner.begin_process_start(&rpc_client) { @@ -885,6 +909,11 @@ impl ExecServerClient { let process_id = params.process_id.clone(); let state = Arc::new(SessionState::new(/*recoverable*/ false)); + if let Some(controller) = network_policy_controller.as_ref() { + state + .network_policy_controller + .store(Some(Arc::new(controller.clone()))); + } if let Err(error) = self.inner.insert_session(&process_id, Arc::clone(&state)) { self.inner.finish_process_start(); return Err(error); @@ -1384,9 +1413,14 @@ impl Session { pub(crate) async fn terminate(&self) -> Result<(), ExecServerError> { self.client.terminate(&self.process_id).await?; + self.cancel_network_policy_decisions(); Ok(()) } + pub(crate) fn cancel_network_policy_decisions(&self) { + self.state.network_policy_cancelled.cancel(); + } + pub(crate) async fn unregister(&self) { self.client .inner @@ -1733,21 +1767,24 @@ mod tests { let process_id = ProcessId::from("trace-process"); let session = client - .start_process(ExecParams { - process_id: process_id.clone(), - argv: vec!["true".to_string()], - cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) - .expect("cwd URI"), - env_policy: None, - env: HashMap::new(), - tty: false, - pipe_stdin: false, - arg0: None, - sandbox: None, - enforce_managed_network: false, - managed_network: None, - network_proxy: None, - }) + .start_process( + ExecParams { + process_id: process_id.clone(), + argv: vec!["true".to_string()], + cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) + .expect("cwd URI"), + env_policy: None, + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }, + /*network_policy_decider*/ None, + ) .instrument(parent_span) .await .expect("process start should succeed"); diff --git a/codex-rs/exec-server/src/client/tests/network_policy_tests.rs b/codex-rs/exec-server/src/client/tests/network_policy_tests.rs index 4445c959ff..a3f934c925 100644 --- a/codex-rs/exec-server/src/client/tests/network_policy_tests.rs +++ b/codex-rs/exec-server/src/client/tests/network_policy_tests.rs @@ -127,22 +127,23 @@ async fn abandoned_process_start_unregisters_and_cleans_up() { let start_client = client.clone(); let start_process_id = process_id.clone(); let start = tokio::spawn(async move { + let params = ExecParams { + process_id: start_process_id, + argv: vec!["true".to_string()], + cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) + .expect("cwd URI"), + env_policy: None, + env: Default::default(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }; start_client - .start_process(ExecParams { - process_id: start_process_id, - argv: vec!["true".to_string()], - cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) - .expect("cwd URI"), - env_policy: None, - env: Default::default(), - tty: false, - pipe_stdin: false, - arg0: None, - sandbox: None, - enforce_managed_network: false, - managed_network: None, - network_proxy: None, - }) + .start_process(params, /*network_policy_decider*/ None) .await }); start_seen_rx.await.expect("start should be observed"); diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index afec44b818..0e9df9cd44 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/codex-rs/exec-server/src/process.rs @@ -4,6 +4,7 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use codex_network_proxy::NetworkPolicyDecider; use tokio::sync::broadcast; use tokio::sync::watch; @@ -202,6 +203,19 @@ pub type ExecProcessFuture<'a, T> = pub trait ExecBackend: Send + Sync { fn start(&self, params: ExecParams) -> ExecBackendFuture<'_>; + + /// Starts a process with an authoritative controller-side policy decider. + fn start_with_network_policy_decider( + &self, + _params: ExecParams, + _decider: Arc, + ) -> ExecBackendFuture<'_> { + Box::pin(async { + Err(ExecServerError::Protocol( + "exec backend does not support remote network policy decisions".to_string(), + )) + }) + } } pub type ExecBackendFuture<'a> = diff --git a/codex-rs/exec-server/src/remote_process.rs b/codex-rs/exec-server/src/remote_process.rs index 72f41ae70e..57abc6d176 100644 --- a/codex-rs/exec-server/src/remote_process.rs +++ b/codex-rs/exec-server/src/remote_process.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use codex_network_proxy::NetworkPolicyDecider; use tokio::sync::watch; use tracing::trace; @@ -34,9 +35,10 @@ impl RemoteProcess { async fn start( &self, params: ExecParams, + network_policy_decider: Option>, ) -> Result { let client = self.client.get().await?; - let session = client.start_process(params).await?; + let session = client.start_process(params, network_policy_decider).await?; Ok(StartedExecProcess { process: Arc::new(RemoteExecProcess { session }), @@ -46,7 +48,17 @@ impl RemoteProcess { impl ExecBackend for RemoteProcess { fn start(&self, params: ExecParams) -> ExecBackendFuture<'_> { - Box::pin(RemoteProcess::start(self, params)) + Box::pin(RemoteProcess::start( + self, params, /*network_policy_decider*/ None, + )) + } + + fn start_with_network_policy_decider( + &self, + params: ExecParams, + decider: Arc, + ) -> ExecBackendFuture<'_> { + Box::pin(RemoteProcess::start(self, params, Some(decider))) } } @@ -113,6 +125,7 @@ impl ExecProcess for RemoteExecProcess { impl Drop for RemoteExecProcess { fn drop(&mut self) { + self.session.cancel_network_policy_decisions(); let session = self.session.clone(); tokio::spawn(async move { session.unregister().await; diff --git a/codex-rs/hooks/src/engine/mod.rs b/codex-rs/hooks/src/engine/mod.rs index 0aaa4b0d33..d0c1aba1fe 100644 --- a/codex-rs/hooks/src/engine/mod.rs +++ b/codex-rs/hooks/src/engine/mod.rs @@ -34,6 +34,7 @@ use codex_protocol::protocol::HookSource; use codex_protocol::protocol::HookTrustStatus; use codex_utils_absolute_path::AbsolutePathBuf; use std::collections::HashMap; +use std::time::Duration; #[derive(Debug, Clone)] pub(crate) struct CommandShell { @@ -165,6 +166,17 @@ impl ClaudeHooksEngine { crate::events::permission_request::preview(&self.handlers, request) } + pub(crate) fn max_permission_request_timeout(&self) -> Duration { + Duration::from_secs( + self.handlers + .iter() + .filter(|handler| handler.event_name == HookEventName::PermissionRequest) + .map(|handler| handler.timeout_sec) + .max() + .unwrap_or_default(), + ) + } + pub(crate) fn preview_post_tool_use( &self, request: &PostToolUseRequest, diff --git a/codex-rs/hooks/src/registry.rs b/codex-rs/hooks/src/registry.rs index 8fdfaf9cf4..f729e7a0f2 100644 --- a/codex-rs/hooks/src/registry.rs +++ b/codex-rs/hooks/src/registry.rs @@ -1,5 +1,6 @@ use codex_config::ConfigLayerStack; use codex_plugin::PluginHookSource; +use std::time::Duration; use tokio::process::Command; use crate::engine::ClaudeHooksEngine; @@ -129,6 +130,13 @@ impl Hooks { self.engine.preview_permission_request(request) } + /// Maximum configured timeout among PermissionRequest hooks. + /// + /// Matching handlers run concurrently, so their aggregate timeout is bounded by this maximum. + pub fn max_permission_request_timeout(&self) -> Duration { + self.engine.max_permission_request_timeout() + } + pub fn preview_post_tool_use( &self, request: &PostToolUseRequest, diff --git a/codex-rs/linux-sandbox/src/proxy_routing.rs b/codex-rs/linux-sandbox/src/proxy_routing.rs index 57682aba81..245207c496 100644 --- a/codex-rs/linux-sandbox/src/proxy_routing.rs +++ b/codex-rs/linux-sandbox/src/proxy_routing.rs @@ -13,6 +13,7 @@ use std::io::Read; use std::io::Write; use std::net::IpAddr; use std::net::Ipv4Addr; +use std::net::Shutdown; use std::net::SocketAddr; use std::net::TcpListener; use std::net::TcpStream; @@ -685,8 +686,13 @@ fn harden_bridge_process() -> io::Result<()> { fn proxy_bidirectional(mut tcp_stream: TcpStream, mut unix_stream: UnixStream) -> io::Result<()> { let mut tcp_reader = tcp_stream.try_clone()?; let mut unix_writer = unix_stream.try_clone()?; - let tcp_to_unix = std::thread::spawn(move || std::io::copy(&mut tcp_reader, &mut unix_writer)); + let tcp_to_unix = std::thread::spawn(move || { + let result = std::io::copy(&mut tcp_reader, &mut unix_writer); + let _ = unix_writer.shutdown(Shutdown::Write); + result + }); let unix_to_tcp = std::io::copy(&mut unix_stream, &mut tcp_stream); + let _ = tcp_stream.shutdown(Shutdown::Write); let tcp_to_unix = tcp_to_unix .join() .map_err(|_| io::Error::other("bridge thread panicked"))?; diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index e1a8ce443b..2918e8e8c2 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -5,9 +5,13 @@ use crate::config; 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::NetworkDecision; +use crate::network_policy::NetworkDecisionSource; use crate::network_policy::NetworkPolicyDecider; use crate::runtime::BlockedRequestObserver; use crate::runtime::ConfigState; +use crate::runtime::HostBlockDecision; +use crate::runtime::HostBlockReason; use crate::runtime::unix_socket_permissions_supported; use crate::socks5; use crate::state::NetworkProxyState; @@ -835,6 +839,53 @@ impl NetworkProxy { }) } + /// Returns the policy decider with trusted execution attribution restored. + pub fn remote_policy_decider(&self) -> Option> { + let scope = self.execution_scope.as_ref()?; + self.state.for_execution_token(&scope.attribution_token)?; + let decider = Arc::clone(self.policy_decider.as_ref()?); + let state = Arc::clone(&self.state); + let environment_id = scope.environment_id.clone(); + let execution_id = scope.execution_id.clone(); + let execution_lifetime = scope.lifetime_tx.subscribe(); + Some(Arc::new(move |mut request: crate::NetworkPolicyRequest| { + let decider = Arc::clone(&decider); + let state = Arc::clone(&state); + let mut execution_lifetime = execution_lifetime.clone(); + request.environment_id = Some(environment_id.clone()); + request.execution_id = Some(execution_id.clone()); + async move { + tokio::select! { + biased; + _ = execution_lifetime.changed() => { + crate::NetworkDecision::deny(crate::reasons::REASON_NOT_ALLOWED) + } + decision = async { + match state.host_blocked(&request.host, request.port).await { + Ok(HostBlockDecision::Allowed) => NetworkDecision::Allow, + Ok(HostBlockDecision::Blocked(HostBlockReason::NotAllowed)) => { + decider.decide(request).await + } + Ok(HostBlockDecision::Blocked(reason)) => { + NetworkDecision::deny_with_source( + reason.as_str(), + NetworkDecisionSource::BaselinePolicy, + ) + } + Err(err) => { + warn!("failed to evaluate controller network policy: {err}"); + NetworkDecision::deny_with_source( + crate::reasons::REASON_NOT_ALLOWED, + NetworkDecisionSource::BaselinePolicy, + ) + } + } + } => decision, + } + } + })) + } + pub async fn add_allowed_domain(&self, host: &str) -> Result<()> { self.state.add_allowed_domain(host).await } @@ -1482,6 +1533,121 @@ mod tests { #[cfg(target_os = "windows")] static WINDOWS_INGRESS_TEST_LOCK: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1); + fn https_request(host: &str) -> crate::NetworkPolicyRequest { + crate::NetworkPolicyRequest::new(crate::NetworkPolicyRequestArgs { + protocol: crate::NetworkProtocol::HttpsConnect, + host: host.to_string(), + port: 443, + environment_id: None, + client_addr: None, + method: None, + command: None, + exec_policy_hint: None, + }) + } + + #[tokio::test] + async fn remote_policy_decider_rechecks_live_policy_and_restores_attribution() -> Result<()> { + let captured = Arc::new(Mutex::new(None)); + let captured_request = Arc::clone(&captured); + let config = NetworkProxyConfig { + allow_local_binding: true, + ..NetworkProxyConfig::default() + }; + let proxy = NetworkProxy::builder() + .state(Arc::new(network_proxy_state_for_policy(config))) + .policy_decider(move |request: crate::NetworkPolicyRequest| { + let captured = Arc::clone(&captured_request); + async move { + *captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some((request.host, request.environment_id, request.execution_id)); + crate::NetworkDecision::Allow + } + }) + .managed_by_codex(/*managed_by_codex*/ false) + .build() + .await?; + let scoped = proxy.for_execution("remote", "execution-1", "token-1".to_string())?; + let decider = scoped + .remote_policy_decider() + .expect("execution-scoped proxy should expose its policy decider"); + let mut request = https_request("allowed.example"); + request.environment_id = Some("forged-environment".to_string()); + request.execution_id = Some("forged-execution".to_string()); + + assert_eq!(decider.decide(request).await, crate::NetworkDecision::Allow); + assert_eq!( + *captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + Some(( + "allowed.example".to_string(), + Some("remote".to_string()), + Some("execution-1".to_string()) + )) + ); + + scoped.add_denied_domain("denied.example").await?; + assert_eq!( + decider.decide(https_request("denied.example")).await, + crate::NetworkDecision::deny_with_source( + crate::reasons::REASON_DENIED, + crate::NetworkDecisionSource::BaselinePolicy, + ) + ); + assert_eq!( + *captured + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + Some(( + "allowed.example".to_string(), + Some("remote".to_string()), + Some("execution-1".to_string()) + )) + ); + Ok(()) + } + + #[tokio::test] + async fn remote_policy_decider_stops_with_execution_scope() -> Result<()> { + let decision_started = Arc::new(tokio::sync::Notify::new()); + let decision_started_for_decider = Arc::clone(&decision_started); + let config = NetworkProxyConfig { + allow_local_binding: true, + ..NetworkProxyConfig::default() + }; + let proxy = NetworkProxy::builder() + .state(Arc::new(network_proxy_state_for_policy(config))) + .policy_decider(move |_request: crate::NetworkPolicyRequest| { + let decision_started = Arc::clone(&decision_started_for_decider); + async move { + decision_started.notify_one(); + std::future::pending::().await + } + }) + .managed_by_codex(/*managed_by_codex*/ false) + .build() + .await?; + let scoped = proxy.for_execution("remote", "execution-1", "token-1".to_string())?; + let decider = scoped + .remote_policy_decider() + .expect("execution-scoped proxy should expose its policy decider"); + let request = https_request("pending.example"); + let decision = tokio::spawn(async move { decider.decide(request).await }); + + decision_started.notified().await; + drop(scoped); + + let decision = tokio::time::timeout(std::time::Duration::from_secs(1), decision).await??; + assert_eq!( + decision, + crate::NetworkDecision::deny(crate::reasons::REASON_NOT_ALLOWED) + ); + Ok(()) + } + #[tokio::test] async fn managed_proxy_builder_uses_loopback_ports() { #[cfg(target_os = "windows")] diff --git a/codex-rs/network-proxy/src/proxy/execution_scope.rs b/codex-rs/network-proxy/src/proxy/execution_scope.rs index 621b047553..ac5fecbaee 100644 --- a/codex-rs/network-proxy/src/proxy/execution_scope.rs +++ b/codex-rs/network-proxy/src/proxy/execution_scope.rs @@ -4,6 +4,8 @@ pub(super) struct ExecutionScope { pub(super) environment_id: String, pub(super) execution_id: String, pub(super) attribution_token: String, + // Dropping the execution scope closes this channel and cancels remote reviews. + pub(super) lifetime_tx: tokio::sync::watch::Sender<()>, state: Arc, } @@ -28,11 +30,13 @@ impl NetworkProxy { self.state .register_execution(&attribution_token, environment_id, execution_id); + let (lifetime_tx, _) = tokio::sync::watch::channel(()); let mut proxy = self.clone(); proxy.execution_scope = Some(Arc::new(ExecutionScope { environment_id: environment_id.to_string(), execution_id: execution_id.to_string(), attribution_token, + lifetime_tx, state: Arc::clone(&self.state), })); Ok(proxy)