diff --git a/codex-rs/core/src/network_policy_decision_tests.rs b/codex-rs/core/src/network_policy_decision_tests.rs index 0640041847..a52a521324 100644 --- a/codex-rs/core/src/network_policy_decision_tests.rs +++ b/codex-rs/core/src/network_policy_decision_tests.rs @@ -164,6 +164,7 @@ fn denied_network_policy_message_requires_deny_decision() { decision: Some("ask".to_string()), source: Some("decider".to_string()), port: Some(80), + request_origin: None, timestamp: 0, }; assert_eq!(denied_network_policy_message(&blocked), None); @@ -181,6 +182,7 @@ fn denied_network_policy_message_for_denylist_block_is_explicit() { decision: Some("deny".to_string()), source: Some("baseline_policy".to_string()), port: Some(80), + request_origin: None, timestamp: 0, }; assert_eq!( diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index 2b63164b4d..f5acdd3395 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -30,6 +30,7 @@ use codex_protocol::protocol::WarningEvent; use indexmap::IndexMap; use std::collections::HashMap; use std::collections::HashSet; +use std::io; use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::Notify; @@ -89,6 +90,7 @@ pub(crate) struct ActiveNetworkApproval { registration_id: Option, mode: NetworkApprovalMode, cancellation_token: CancellationToken, + network: NetworkProxy, } impl ActiveNetworkApproval { @@ -100,11 +102,16 @@ impl ActiveNetworkApproval { self.cancellation_token.clone() } + pub(crate) fn network(&self) -> &NetworkProxy { + &self.network + } + pub(crate) fn into_deferred(self) -> Option { let ActiveNetworkApproval { registration_id, mode, cancellation_token, + network: _, } = self; match (mode, registration_id) { (NetworkApprovalMode::Deferred, Some(registration_id)) => { @@ -245,6 +252,8 @@ enum ActiveNetworkApprovalAttribution { struct NetworkApprovalCallState { active_calls: IndexMap>, call_outcomes: HashMap, + // Keep each listener alive for the same lifetime as its approval registration. + request_scopes: HashMap, } pub(crate) struct NetworkApprovalService { @@ -305,9 +314,8 @@ impl NetworkApprovalService { async fn resolve_single_active_call(&self) -> Option> { let calls = self.calls.lock().await; - // Blocked proxy requests are not attributed to a specific tool call. Only pick an owner + // Legacy/unscoped proxy requests can still arrive without an origin. Only pick an owner // when there is exactly one candidate; with concurrent calls, canceling one would be a guess. - // TODO: Carry blocked-request attribution so concurrent active calls can be handled safely. if calls.active_calls.len() == 1 { return calls.active_calls.values().next().cloned(); } @@ -377,6 +385,7 @@ impl NetworkApprovalService { async fn remove_call(&self, registration_id: &str) -> Option { let mut calls = self.calls.lock().await; calls.active_calls.shift_remove(registration_id); + calls.request_scopes.remove(registration_id); calls.call_outcomes.remove(registration_id) } @@ -393,8 +402,12 @@ impl NetworkApprovalService { return; }; - self.record_outcome_for_single_active_call(NetworkApprovalOutcome::DeniedByPolicy(message)) - .await; + let outcome = NetworkApprovalOutcome::DeniedByPolicy(message); + if let Some(request_origin) = blocked.request_origin.as_deref() { + self.record_call_outcome(request_origin, outcome).await; + } else { + self.record_outcome_for_single_active_call(outcome).await; + } } async fn active_turn_context( @@ -431,28 +444,56 @@ impl NetworkApprovalService { NetworkProtocol::Socks5Tcp => NetworkApprovalProtocol::Socks5Tcp, NetworkProtocol::Socks5Udp => NetworkApprovalProtocol::Socks5Udp, }; - let (owner_call, active_environment_id) = - if let Some(environment_id) = request.environment_id.clone() { - let owner_call = match self.resolve_active_call_attribution().await { - ActiveNetworkApprovalAttribution::Single(call) => { - (call.environment_id == environment_id).then_some(call) - } - ActiveNetworkApprovalAttribution::None - | ActiveNetworkApprovalAttribution::Ambiguous => None, + let attributed_call = match request.request_origin.as_deref() { + Some(request_origin) => { + let Some(call) = self + .calls + .lock() + .await + .active_calls + .get(request_origin) + .cloned() + else { + return NetworkDecision::deny(REASON_NOT_ALLOWED); }; - (owner_call, Some(environment_id)) - } else { - match self.resolve_active_call_attribution().await { - ActiveNetworkApprovalAttribution::None => (None, None), - ActiveNetworkApprovalAttribution::Single(call) => { - let environment_id = call.environment_id.clone(); - (Some(call), Some(environment_id)) - } - ActiveNetworkApprovalAttribution::Ambiguous => { - return NetworkDecision::deny(REASON_NOT_ALLOWED); - } + if request + .environment_id + .as_deref() + .is_some_and(|environment_id| call.environment_id != environment_id) + { + return NetworkDecision::deny(REASON_NOT_ALLOWED); } + Some(call) + } + None => None, + }; + let (owner_call, active_environment_id) = if let Some(call) = attributed_call { + let environment_id = request + .environment_id + .clone() + .unwrap_or_else(|| call.environment_id.clone()); + (Some(call), Some(environment_id)) + } else if let Some(environment_id) = request.environment_id.clone() { + let owner_call = match self.resolve_active_call_attribution().await { + ActiveNetworkApprovalAttribution::Single(call) => { + (call.environment_id == environment_id).then_some(call) + } + ActiveNetworkApprovalAttribution::None + | ActiveNetworkApprovalAttribution::Ambiguous => None, }; + (owner_call, Some(environment_id)) + } else { + match self.resolve_active_call_attribution().await { + ActiveNetworkApprovalAttribution::None => (None, None), + ActiveNetworkApprovalAttribution::Single(call) => { + let environment_id = call.environment_id.clone(); + (Some(call), Some(environment_id)) + } + ActiveNetworkApprovalAttribution::Ambiguous => { + return NetworkDecision::deny(REASON_NOT_ALLOWED); + } + } + }; let turn_context = Self::active_turn_context(session.as_ref()).await; let Some(environment_id) = active_environment_id.or_else(|| { turn_context @@ -789,19 +830,32 @@ pub(crate) async fn begin_network_approval( turn_id: &str, managed_network_active: bool, spec: Option, -) -> Option { +) -> Result, ToolError> { let NetworkApprovalSpec { network, mode, trigger, command, environment_id, - } = spec?; - if !managed_network_active || network.is_none() { - return None; + } = match spec { + Some(spec) => spec, + None => return Ok(None), + }; + let Some(network) = network else { + return Ok(None); + }; + if !managed_network_active { + return Ok(None); } let registration_id = Uuid::new_v4().to_string(); + let network = network + .scope_for_request(&environment_id, registration_id.clone()) + .map_err(|err| { + ToolError::Codex(codex_protocol::error::CodexErr::Io(io::Error::other( + format!("failed to create request-scoped network proxy: {err}"), + ))) + })?; let cancellation_token = CancellationToken::new(); session .services @@ -815,12 +869,21 @@ pub(crate) async fn begin_network_approval( cancellation_token.clone(), ) .await; + session + .services + .network_approval + .calls + .lock() + .await + .request_scopes + .insert(registration_id.clone(), network.clone()); - Some(ActiveNetworkApproval { + Ok(Some(ActiveNetworkApproval { registration_id: Some(registration_id), mode, cancellation_token, - }) + network, + })) } pub(crate) async fn finish_immediate_network_approval( diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index aaa0dfe9d1..2e5f39a23b 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -68,13 +68,17 @@ impl ToolOrchestrator { where T: ToolRuntime, { - let network_approval = begin_network_approval( + let network_approval = match begin_network_approval( &tool_ctx.session, &tool_ctx.turn.sub_id, managed_network_active, tool.network_approval_spec(req, tool_ctx), ) - .await; + .await + { + Ok(network_approval) => network_approval, + Err(err) => return (Err(err), None), + }; let attempt_tool_ctx = ToolCtx { session: tool_ctx.session.clone(), @@ -98,6 +102,9 @@ impl ToolOrchestrator { network_denial_cancellation_token: network_approval .as_ref() .map(ActiveNetworkApproval::cancellation_token), + network_proxy: network_approval + .as_ref() + .map(ActiveNetworkApproval::network), }; let run_result = tool .run(req, &attempt_with_network_approval, &attempt_tool_ctx) @@ -274,6 +281,7 @@ impl ToolOrchestrator { .permissions .windows_sandbox_private_desktop, network_denial_cancellation_token: None, + network_proxy: None, }; let initial_attempt_start = Instant::now(); @@ -456,6 +464,7 @@ impl ToolOrchestrator { .permissions .windows_sandbox_private_desktop, network_denial_cancellation_token: None, + network_proxy: None, }; // Second attempt. diff --git a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs index f1c6f43aa3..170068cde0 100644 --- a/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs +++ b/codex-rs/core/src/tools/runtimes/apply_patch_tests.rs @@ -232,6 +232,7 @@ async fn file_system_sandbox_context_uses_active_attempt() { windows_sandbox_level: WindowsSandboxLevel::RestrictedToken, windows_sandbox_private_desktop: true, network_denial_cancellation_token: None, + network_proxy: None, }; let sandbox = ApplyPatchRuntime::file_system_sandbox_context_for_attempt(&req, &attempt) @@ -300,6 +301,7 @@ async fn no_sandbox_attempt_has_no_file_system_context() { windows_sandbox_level: WindowsSandboxLevel::Disabled, windows_sandbox_private_desktop: false, network_denial_cancellation_token: None, + network_proxy: None, }; assert_eq!( diff --git a/codex-rs/core/src/tools/runtimes/mod_tests.rs b/codex-rs/core/src/tools/runtimes/mod_tests.rs index 9b485dbe58..7862a1e3c2 100644 --- a/codex-rs/core/src/tools/runtimes/mod_tests.rs +++ b/codex-rs/core/src/tools/runtimes/mod_tests.rs @@ -118,6 +118,7 @@ async fn explicit_escalation_prepares_exec_without_managed_network() -> anyhow:: windows_sandbox_level: WindowsSandboxLevel::Disabled, windows_sandbox_private_desktop: false, network_denial_cancellation_token: None, + network_proxy: None, }; let exec_request = attempt diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index c1d5191204..8d6fbe9e6d 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -317,10 +317,10 @@ impl<'a> ToolRuntime for UnifiedExecRunt req.sandbox_permissions, &file_system_sandbox_policy, ); - let managed_network = managed_network_for_sandbox_permissions( + let managed_network = attempt.network_proxy(managed_network_for_sandbox_permissions( req.network.as_ref(), launch_sandbox_permissions, - ); + )); let mut env = exec_env_for_sandbox_permissions(&req.env, launch_sandbox_permissions); if let Some(network) = managed_network { network diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index 77f68d96b5..5c0a1738dd 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -424,9 +424,17 @@ pub(crate) struct SandboxAttempt<'a> { pub windows_sandbox_level: codex_protocol::config_types::WindowsSandboxLevel, pub windows_sandbox_private_desktop: bool, pub network_denial_cancellation_token: Option, + pub(crate) network_proxy: Option<&'a NetworkProxy>, } impl<'a> SandboxAttempt<'a> { + pub(crate) fn network_proxy<'b>( + &'b self, + fallback: Option<&'b NetworkProxy>, + ) -> Option<&'b NetworkProxy> { + fallback.map(|fallback| self.network_proxy.unwrap_or(fallback)) + } + pub fn env_for( &self, command: SandboxCommand, @@ -434,6 +442,7 @@ impl<'a> SandboxAttempt<'a> { network: Option<&NetworkProxy>, environment_id: Option<&str>, ) -> Result { + let network = self.network_proxy(network); let request = self .manager .transform(SandboxTransformRequest { @@ -466,6 +475,7 @@ impl<'a> SandboxAttempt<'a> { network: Option<&NetworkProxy>, environment_id: Option<&str>, ) -> Result { + let network = self.network_proxy(network); let exec_server_permissions = effective_permission_profile( self.exec_server_permissions, command.additional_permissions.as_ref(), diff --git a/codex-rs/core/src/tools/sandboxing_tests.rs b/codex-rs/core/src/tools/sandboxing_tests.rs index c647e1c401..743706a273 100644 --- a/codex-rs/core/src/tools/sandboxing_tests.rs +++ b/codex-rs/core/src/tools/sandboxing_tests.rs @@ -226,6 +226,7 @@ fn exec_server_env_keeps_command_native_and_carries_sandbox_context() { windows_sandbox_level: codex_protocol::config_types::WindowsSandboxLevel::Disabled, windows_sandbox_private_desktop: false, network_denial_cancellation_token: None, + network_proxy: None, }; let command = SandboxCommand { program: "/bin/bash".into(), diff --git a/codex-rs/core/tests/suite/network_approval.rs b/codex-rs/core/tests/suite/network_approval.rs index f476fb7287..2c63cd071c 100644 --- a/codex-rs/core/tests/suite/network_approval.rs +++ b/codex-rs/core/tests/suite/network_approval.rs @@ -55,6 +55,129 @@ use tempfile::TempDir; const NETWORK_TEST_HOST: &str = "codex-network-test.invalid"; const NETWORK_TEST_TARGET: &str = "http://codex-network-test.invalid:80"; +const FIRST_CONCURRENT_NETWORK_TEST_HOST: &str = "first.codex-network-test.invalid"; +const SECOND_CONCURRENT_NETWORK_TEST_HOST: &str = "second.codex-network-test.invalid"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn guardian_receives_exact_triggers_for_concurrent_network_requests() -> Result<()> { + skip_if_wine_exec!(Ok(()), "uses the POSIX/Python network fixture"); + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + skip_if_windows!(Ok(())); + + let server = start_mock_server().await; + let test = managed_network_unified_exec_test(&server).await?; + let barrier_dir = TempDir::new_in(test.cwd.path())?; + let first_marker = barrier_dir.path().join("first"); + let second_marker = barrier_dir.path().join("second"); + let network_command = |marker: &PathBuf, peer_marker: &PathBuf, host: &str| { + format!( + "touch '{}' && while [ ! -e '{}' ]; do sleep 0.01; done && python3 -c \"import urllib.request; urllib.request.build_opener(urllib.request.ProxyHandler()).open('http://{host}', timeout=10).read()\"", + marker.display(), + peer_marker.display(), + ) + }; + let first_command = network_command( + &first_marker, + &second_marker, + FIRST_CONCURRENT_NETWORK_TEST_HOST, + ); + let second_command = network_command( + &second_marker, + &first_marker, + SECOND_CONCURRENT_NETWORK_TEST_HOST, + ); + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-network-concurrent"), + ev_function_call( + "exec-network-first", + "exec_command", + &serde_json::to_string(&network_exec_args( + LOCAL_ENVIRONMENT_ID, + &first_command, + ))?, + ), + ev_function_call( + "exec-network-second", + "exec_command", + &serde_json::to_string(&network_exec_args( + LOCAL_ENVIRONMENT_ID, + &second_command, + ))?, + ), + ev_completed("resp-network-concurrent"), + ]), + sse(vec![ + ev_response_created("resp-network-guardian-1"), + ev_assistant_message("msg-network-guardian-1", r#"{"outcome":"allow"}"#), + ev_completed("resp-network-guardian-1"), + ]), + sse(vec![ + ev_response_created("resp-network-guardian-2"), + ev_assistant_message("msg-network-guardian-2", r#"{"outcome":"allow"}"#), + ev_completed("resp-network-guardian-2"), + ]), + sse(vec![ + ev_response_created("resp-network-done"), + ev_assistant_message("msg-network-done", "done"), + ev_completed("resp-network-done"), + ]), + ], + ) + .await; + + submit_managed_network_turn( + &test, + "run both network requests", + vec![local(test.config.cwd.clone())], + ApprovalsReviewer::AutoReview, + ) + .await?; + wait_for_turn_complete(&test).await; + + let mut actual_triggers = responses + .requests() + .into_iter() + .filter(|request| { + request.body_json()["client_metadata"]["x-openai-subagent"].as_str() == Some("guardian") + }) + .map(|request| { + let user_texts = request.message_input_texts("user"); + let action: Value = serde_json::from_str( + user_texts + .iter() + .find(|text| text.contains("\"tool\": \"network_access\"")) + .context("expected network access JSON in Guardian request")? + .trim(), + )?; + let trigger = &action["trigger"]; + Ok(( + trigger["callId"] + .as_str() + .context("expected exact trigger call id")? + .to_string(), + trigger["command"] + .as_array() + .and_then(|command| command.last()) + .and_then(Value::as_str) + .context("expected exact trigger command")? + .to_string(), + )) + }) + .collect::>>()?; + actual_triggers.sort_unstable(); + let mut expected_triggers = vec![ + ("exec-network-first".to_string(), first_command), + ("exec-network-second".to_string(), second_command), + ]; + expected_triggers.sort_unstable(); + assert_eq!(actual_triggers, expected_triggers); + + Ok(()) +} #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn approved_network_host_for_one_environment_still_prompts_in_another() -> Result<()> { @@ -101,6 +224,7 @@ async fn approved_network_host_for_one_environment_still_prompts_in_another() -> &test, "fetch from the local environment", environments.clone(), + ApprovalsReviewer::User, ) .await?; let approval = expect_network_approval(&test, LOCAL_ENVIRONMENT_ID).await?; @@ -124,6 +248,7 @@ async fn approved_network_host_for_one_environment_still_prompts_in_another() -> &test, "fetch from the remote environment", environments.clone(), + ApprovalsReviewer::User, ) .await?; let approval = expect_network_approval(&test, REMOTE_ENVIRONMENT_ID).await?; @@ -227,9 +352,16 @@ async fn mount_exec_network_turn( } fn network_fetch_args(environment_id: &str) -> Value { + let command = format!( + "python3 -c \"import urllib.request; opener = urllib.request.build_opener(urllib.request.ProxyHandler()); print('OK:' + opener.open('http://{NETWORK_TEST_HOST}', timeout=2).read().decode(errors='replace'))\"" + ); + network_exec_args(environment_id, &command) +} + +fn network_exec_args(environment_id: &str, command: &str) -> Value { json!({ "shell": "/bin/sh", - "cmd": format!("python3 -c \"import urllib.request; opener = urllib.request.build_opener(urllib.request.ProxyHandler()); print('OK:' + opener.open('http://{NETWORK_TEST_HOST}', timeout=2).read().decode(errors='replace'))\""), + "cmd": command, "login": false, "yield_time_ms": 1_000, "environment_id": environment_id, @@ -240,6 +372,7 @@ async fn submit_managed_network_turn( test: &TestCodex, prompt: &str, environments: Vec, + approvals_reviewer: ApprovalsReviewer, ) -> Result<()> { let permission_profile = PermissionProfile::workspace_write_with( &[], @@ -264,7 +397,7 @@ async fn submit_managed_network_turn( thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { environments: Some(turn_environment_selections), approval_policy: Some(AskForApproval::OnFailure), - approvals_reviewer: Some(ApprovalsReviewer::User), + approvals_reviewer: Some(approvals_reviewer), sandbox_policy: Some(sandbox_policy), permission_profile, collaboration_mode: Some(codex_protocol::config_types::CollaborationMode { diff --git a/codex-rs/network-proxy/src/network_policy.rs b/codex-rs/network-proxy/src/network_policy.rs index e5792a07c0..3f64a56d59 100644 --- a/codex-rs/network-proxy/src/network_policy.rs +++ b/codex-rs/network-proxy/src/network_policy.rs @@ -84,6 +84,7 @@ pub struct NetworkPolicyRequest { pub method: Option, pub command: Option, pub exec_policy_hint: Option, + pub request_origin: Option, } pub struct NetworkPolicyRequestArgs { @@ -118,6 +119,7 @@ impl NetworkPolicyRequest { method, command, exec_policy_hint, + request_origin: None, } } } @@ -300,7 +302,10 @@ pub(crate) async fn evaluate_host_policy( HostBlockDecision::Allowed => (NetworkDecision::Allow, false), HostBlockDecision::Blocked(HostBlockReason::NotAllowed) => { if let Some(decider) = decider { - let decider_decision = map_decider_decision(decider.decide(request.clone()).await); + let mut request = request.clone(); + // Trust only the listener-scoped state, never request metadata from the client. + request.request_origin = state.request_origin(); + let decider_decision = map_decider_decision(decider.decide(request).await); let policy_override = matches!(decider_decision, NetworkDecision::Allow); (decider_decision, policy_override) } else { diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index 13998dd447..56edb621ee 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -229,6 +229,7 @@ impl NetworkProxyBuilder { reserved_listeners, policy_decider: self.policy_decider, environment_proxies: Arc::new(Mutex::new(HashMap::new())), + request_scope: None, }) } } @@ -334,6 +335,21 @@ struct EnvironmentProxy { socks_task: Option>>, } +struct RequestScopedProxy { + environment_id: String, + http_task: JoinHandle>, + socks_task: Option>>, +} + +impl Drop for RequestScopedProxy { + fn drop(&mut self) { + self.http_task.abort(); + if let Some(socks_task) = self.socks_task.as_ref() { + socks_task.abort(); + } + } +} + #[derive(Clone)] pub struct NetworkProxy { state: Arc, @@ -345,6 +361,7 @@ pub struct NetworkProxy { reserved_listeners: Option>, policy_decider: Option>, environment_proxies: Arc>>, + request_scope: Option>, } impl std::fmt::Debug for NetworkProxy { @@ -618,6 +635,94 @@ impl NetworkProxy { self.socks_addr } + /// Creates a proxy listener pair whose requests carry one execution-specific origin. + /// The listeners remain active until every clone of the returned proxy is dropped. + /// On Windows this currently returns the unscoped proxy because its sandbox firewall is shared. + pub fn scope_for_request(&self, environment_id: &str, request_origin: String) -> Result { + #[cfg(target_os = "windows")] + { + // Windows' offline firewall policy is shared by every sandboxed exec. Distinct + // endpoints therefore cannot be isolated safely between concurrent processes yet. + let _ = (environment_id, request_origin); + Ok(self.clone()) + } + + #[cfg(not(target_os = "windows"))] + { + anyhow::ensure!( + self.request_scope.is_none(), + "cannot create a request scope from an already scoped network proxy" + ); + + let runtime = tokio::runtime::Handle::try_current() + .context("failed to create request-scoped network proxy")?; + let listeners = reserve_loopback_ephemeral_listeners(self.socks_enabled) + .context("failed to reserve request-scoped network proxy")?; + let http_addr = listeners + .http_addr() + .context("failed to read request-scoped HTTP proxy address")?; + let socks_addr = listeners + .socks_addr(self.socks_addr) + .context("failed to read request-scoped SOCKS proxy address")?; + let ReservedListenerSet { + http_listener, + socks_listener, + } = listeners; + + let state = Arc::new(self.state.with_request_origin(request_origin)); + let http_state = Arc::clone(&state); + let http_decider = self.policy_decider.clone(); + let http_environment_id = Some(environment_id.to_string()); + let http_task = runtime.spawn(async move { + http_proxy::run_http_proxy_with_std_listener( + http_state, + http_listener, + http_decider, + http_environment_id, + ) + .await + }); + + let socks_task = if self.socks_enabled { + let socks_state = Arc::clone(&state); + let socks_decider = self.policy_decider.clone(); + let socks_environment_id = Some(environment_id.to_string()); + let socks5_udp_enabled = self.socks5_udp_enabled; + socks_listener.map(|listener| { + runtime.spawn(async move { + socks5::run_socks5_with_std_listener( + socks_state, + listener, + socks_decider, + socks_environment_id, + socks5_udp_enabled, + ) + .await + }) + }) + } else { + None + }; + + Ok(Self { + state, + http_addr, + socks_addr, + socks_enabled: self.socks_enabled, + socks5_udp_enabled: self.socks5_udp_enabled, + runtime_settings: Arc::clone(&self.runtime_settings), + reserved_listeners: None, + policy_decider: self.policy_decider.clone(), + environment_proxies: Arc::new(Mutex::new(HashMap::new())), + request_scope: Some(Arc::new(RequestScopedProxy { + environment_id: environment_id.to_string(), + http_task, + socks_task, + })), + }) + } + } + pub async fn current_cfg(&self) -> Result { self.state.current_cfg().await } @@ -706,6 +811,18 @@ impl NetworkProxy { } fn environment_proxy_addrs(&self, environment_id: &str) -> Result { + if let Some(request_scope) = self.request_scope.as_ref() { + anyhow::ensure!( + request_scope.environment_id == environment_id, + "request-scoped network proxy belongs to environment `{}`, not `{environment_id}`", + request_scope.environment_id + ); + return Ok(EnvironmentProxyAddrs { + http_addr: self.http_addr, + socks_addr: self.socks_addr, + }); + } + let mut proxies = self .environment_proxies .lock() @@ -823,6 +940,10 @@ impl NetworkProxy { } pub async fn run(&self) -> Result { + anyhow::ensure!( + self.request_scope.is_none(), + "request-scoped network proxy is already running" + ); let current_cfg = self.state.current_cfg().await?; if !current_cfg.network.enabled { warn!("network.enabled is false; skipping proxy listeners"); @@ -998,11 +1119,19 @@ impl Drop for NetworkProxyHandle { mod tests { use super::*; use crate::config::NetworkProxySettings; + #[cfg(not(target_os = "windows"))] + use crate::network_policy::NetworkDecision; use crate::state::network_proxy_state_for_policy; use pretty_assertions::assert_eq; use std::net::IpAddr; use std::net::Ipv4Addr; use std::path::Path; + #[cfg(not(target_os = "windows"))] + use std::time::Duration; + #[cfg(not(target_os = "windows"))] + use tokio::io::AsyncWriteExt; + #[cfg(not(target_os = "windows"))] + use tokio::time::timeout; #[tokio::test] async fn managed_proxy_builder_uses_loopback_ports() { @@ -1097,6 +1226,54 @@ mod tests { Ok(()) } + #[cfg(not(target_os = "windows"))] + #[tokio::test] + async fn request_scoped_proxy_stamps_policy_and_blocked_requests() -> Result<()> { + let (policy_tx, mut policy_rx) = tokio::sync::mpsc::unbounded_channel(); + let (blocked_tx, mut blocked_rx) = tokio::sync::mpsc::unbounded_channel(); + let state = Arc::new(network_proxy_state_for_policy( + NetworkProxySettings::default(), + )); + let proxy = NetworkProxy::builder() + .state(state) + .policy_decider( + move |request: crate::network_policy::NetworkPolicyRequest| { + let policy_tx = policy_tx.clone(); + async move { + let _ = policy_tx.send(request.request_origin); + NetworkDecision::deny("not_allowed") + } + }, + ) + .blocked_request_observer(move |request: crate::runtime::BlockedRequest| { + let blocked_tx = blocked_tx.clone(); + async move { + let _ = blocked_tx.send(request.request_origin); + } + }) + .build() + .await?; + let scoped = proxy.scope_for_request("local", "exec-1".to_string())?; + + assert_ne!(scoped.http_addr(), proxy.http_addr()); + let mut stream = tokio::net::TcpStream::connect(scoped.http_addr()).await?; + stream + .write_all( + b"GET http://example.com/ HTTP/1.1\r\nHost: example.com\r\nConnection: close\r\n\r\n", + ) + .await?; + + assert_eq!( + timeout(Duration::from_secs(2), policy_rx.recv()).await?, + Some(Some("exec-1".to_string())) + ); + assert_eq!( + timeout(Duration::from_secs(2), blocked_rx.recv()).await?, + Some(Some("exec-1".to_string())) + ); + Ok(()) + } + #[tokio::test] async fn managed_proxy_builder_does_not_reserve_socks_listener_when_disabled() { let settings = NetworkProxySettings { diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index e721d0de70..d6f40eaf30 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -100,6 +100,8 @@ pub struct BlockedRequest { pub source: Option, #[serde(skip_serializing_if = "Option::is_none")] pub port: Option, + #[serde(skip)] + pub request_origin: Option, pub timestamp: i64, } @@ -138,6 +140,7 @@ impl BlockedRequest { decision, source, port, + request_origin: None, timestamp: unix_timestamp(), } } @@ -208,6 +211,7 @@ pub struct NetworkProxyState { reloader: Arc, blocked_request_observer: Arc>>>, audit_metadata: NetworkProxyAuditMetadata, + request_origin: Option>, } impl std::fmt::Debug for NetworkProxyState { @@ -225,6 +229,7 @@ impl Clone for NetworkProxyState { reloader: self.reloader.clone(), blocked_request_observer: self.blocked_request_observer.clone(), audit_metadata: self.audit_metadata.clone(), + request_origin: self.request_origin.clone(), } } } @@ -275,9 +280,22 @@ impl NetworkProxyState { reloader, blocked_request_observer: Arc::new(RwLock::new(blocked_request_observer)), audit_metadata, + request_origin: None, } } + #[cfg(not(target_os = "windows"))] + pub(crate) fn with_request_origin(&self, request_origin: String) -> Self { + Self { + request_origin: Some(request_origin.into()), + ..self.clone() + } + } + + pub(crate) fn request_origin(&self) -> Option { + self.request_origin.as_deref().map(str::to_string) + } + pub async fn set_blocked_request_observer( &self, blocked_request_observer: Option>, @@ -429,8 +447,9 @@ impl NetworkProxyState { } } - pub async fn record_blocked(&self, entry: BlockedRequest) -> Result<()> { + pub async fn record_blocked(&self, mut entry: BlockedRequest) -> Result<()> { self.reload_if_needed().await?; + entry.request_origin = self.request_origin(); let blocked_for_observer = entry.clone(); let blocked_request_observer = self.blocked_request_observer.read().await.clone(); let violation_line = blocked_request_violation_log_line(&entry); @@ -1204,6 +1223,7 @@ mod tests { decision: Some("ask".to_string()), source: Some("decider".to_string()), port: Some(80), + request_origin: None, timestamp: 1_735_689_600, };