diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 219abe0199..2afcb80afe 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1720,11 +1720,6 @@ impl Session { } else { None }; - let blocked_request_observer = config - .permissions - .network - .as_ref() - .map(|_| build_blocked_request_observer(Arc::clone(&network_approval))); let network_policy_decider = network_policy_decider_session .as_ref() @@ -1747,7 +1742,7 @@ impl Session { config.permissions.permission_profile(), config.effective_local_windows_sandbox_type(), network_policy_decider.as_ref().map(Arc::clone), - blocked_request_observer.as_ref().map(Arc::clone), + Some(build_blocked_request_observer(Arc::clone(&network_approval))), managed_network_requirements_configured, network_proxy_audit_metadata.clone(), ) diff --git a/codex-rs/core/src/tools/network_approval.rs b/codex-rs/core/src/tools/network_approval.rs index 5e1dcc514b..4632966bf9 100644 --- a/codex-rs/core/src/tools/network_approval.rs +++ b/codex-rs/core/src/tools/network_approval.rs @@ -45,6 +45,8 @@ use std::io; use std::sync::Arc; use std::sync::Mutex as SyncMutex; use std::sync::OnceLock as SyncOnceLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use tokio::sync::Mutex; use tokio::sync::Notify; use tokio::sync::OnceCell; @@ -72,12 +74,20 @@ pub(crate) struct NetworkApprovalSpec { pub network_policy: Option, } +#[derive(Clone, Debug)] +struct NetworkApprovalProxy { + proxy: NetworkProxy, + // Only an executor-owned remote proxy can drain its decisions. A shared listener + // can also attribute untagged traffic to this execution, even after a config refresh. + shared_ingress_for_remote: Option>, +} + #[derive(Clone, Debug)] pub(crate) struct DeferredNetworkApproval { registration_id: String, cancellation_token: CancellationToken, finish_outcome: Arc>>, - _execution_proxy: Option, + _execution_proxy: Option, } impl DeferredNetworkApproval { @@ -93,6 +103,20 @@ impl DeferredNetworkApproval { self.cancellation_token.is_cancelled() } + pub(crate) async fn remote_decisions_drained(&self) { + if let Some(NetworkApprovalProxy { + proxy, + shared_ingress_for_remote: Some(shared_ingress), + }) = &self._execution_proxy + { + proxy.wait_for_remote_policy_deciders().await; + if !shared_ingress.load(Ordering::Acquire) { + return; + } + } + std::future::pending().await + } + async fn finish(&self, service: &NetworkApprovalService) -> Result<(), ToolError> { let outcome = self .finish_outcome @@ -109,7 +133,7 @@ impl DeferredNetworkApproval { pub(crate) struct ActiveNetworkApproval { registration_id: Option, cancellation_token: CancellationToken, - execution_proxy: NetworkProxy, + execution_proxy: NetworkApprovalProxy, } impl ActiveNetworkApproval { @@ -118,7 +142,7 @@ impl ActiveNetworkApproval { } pub(crate) fn execution_proxy(&self) -> &NetworkProxy { - &self.execution_proxy + &self.execution_proxy.proxy } pub(crate) fn into_deferred(self) -> Option { @@ -279,6 +303,7 @@ struct PendingNetworkApprovals { } pub(crate) struct NetworkApprovalService { + shared_ingress_registered: Arc, // Pending-request cleanup must publish a tool outcome synchronously before cancelling it. calls: SyncMutex, // Owner cleanup runs from Drop, so this lock cannot require an async task. @@ -461,6 +486,7 @@ impl Drop for PendingHostApprovalOwner<'_> { impl Default for NetworkApprovalService { fn default() -> Self { Self { + shared_ingress_registered: Arc::default(), calls: SyncMutex::new(NetworkApprovalCallState::default()), pending_host_approvals: SyncMutex::default(), session_policy_commit_lock: Mutex::new(()), @@ -1155,6 +1181,9 @@ impl NetworkApprovalService { pub(crate) fn build_blocked_request_observer( network_approval: Arc, ) -> Arc { + network_approval + .shared_ingress_registered + .store(true, Ordering::Release); Arc::new(move |blocked: BlockedRequest| { let network_approval = Arc::clone(&network_approval); async move { @@ -1212,12 +1241,16 @@ pub(crate) async fn begin_network_approval( ) }); let registration_id = Uuid::new_v4().to_string(); - let Some(execution_proxy) = execution_network_proxy( + let environment = environments + .turn_environments() + .find(|environment| environment.selection.environment_id == environment_id); + let executor_owns_proxy = network_policy.is_some() + && turn.config.permissions.network.is_none() + && environment.is_some_and(|environment| environment.environment().is_remote()); + let Some(proxy) = execution_network_proxy( session, turn, - environments - .turn_environments() - .find(|environment| environment.selection.environment_id == environment_id), + environment, network, &environment_id, &permission_profile, @@ -1249,7 +1282,11 @@ pub(crate) async fn begin_network_approval( Ok(Some(ActiveNetworkApproval { registration_id: Some(registration_id), cancellation_token, - execution_proxy, + execution_proxy: NetworkApprovalProxy { + proxy, + shared_ingress_for_remote: executor_owns_proxy + .then(|| Arc::clone(&session.services.network_approval.shared_ingress_registered)), + }, })) } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 10ea9c7bb8..caa45622a2 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -353,30 +353,28 @@ async fn network_denial_message_for_session( } } -async fn wait_for_late_network_denial(network_cancelled: Option) -> bool { - let Some(network_cancelled) = network_cancelled else { - return false; - }; - if network_cancelled.is_cancelled() { - return true; - } - - tokio::select! { - _ = network_cancelled.cancelled() => true, - _ = tokio::time::sleep(LATE_NETWORK_DENIAL_GRACE_PERIOD) => false, - } +async fn wait_for_late_network_denial(network_cancelled: &CancellationToken) -> bool { + tokio::time::timeout( + LATE_NETWORK_DENIAL_GRACE_PERIOD, + network_cancelled.cancelled(), + ) + .await + .is_ok() } async fn finish_deferred_network_approval_after_process_exit_for_session( session: Option<&Arc>, deferred: Option, ) -> Result<(), String> { - wait_for_late_network_denial( - deferred - .as_ref() - .map(DeferredNetworkApproval::cancellation_token), - ) - .await; + if let Some(approval) = &deferred { + let network_cancelled = approval.cancellation_token(); + // Finishing below unregisters the call before we report the outcome. + tokio::select! { + biased; + _ = approval.remote_decisions_drained() => {} + _ = wait_for_late_network_denial(&network_cancelled) => {} + } + } finish_deferred_network_approval_for_session(session, deferred).await } @@ -436,7 +434,8 @@ fn terminate_process_on_network_denial( let denied = tokio::select! { _ = network_cancelled.cancelled() => true, _ = process_exited.cancelled() => { - wait_for_late_network_denial(Some(network_cancelled.clone())).await + // A yielded process stays registered until its next poll. + wait_for_late_network_denial(&network_cancelled).await } }; if !denied { 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 f17801b2fa..602894ff49 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -442,16 +442,121 @@ async fn network_denial_fallback_message_names_sandbox_network_proxy() { ); } -#[tokio::test] -async fn late_network_denial_grace_observes_cancellation_after_exit() { - let cancellation = CancellationToken::new(); - let cancellation_for_task = cancellation.clone(); - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(10)).await; - cancellation_for_task.cancel(); - }); +#[tokio::test(start_paused = true)] +async fn late_network_denial_grace_observes_cancellation_after_exit() -> anyhow::Result<()> { + use crate::environment_selection::EnvironmentAttachment; + use crate::environment_selection::TurnEnvironmentSnapshot; + use crate::environment_selection::TurnEnvironmentState; + use crate::guardian::GuardianNetworkAccessTrigger; + use crate::tools::network_approval::NetworkApprovalSpec; + use crate::tools::network_approval::begin_network_approval; + use crate::tools::network_approval::build_blocked_request_observer; + use codex_network_proxy::BlockedRequest; + use codex_network_proxy::BlockedRequestArgs; + use codex_network_proxy::EnvironmentNetworkPolicy; + use codex_tools::ToolName; - assert!(wait_for_late_network_denial(Some(cancellation)).await); + let (session, mut turn) = crate::session::tests::make_session_and_context().await; + let session = Arc::new(session); + Arc::make_mut(&mut turn.config).permissions.network = None; + let mut remote = turn + .initial_environments + .primary() + .expect("primary environment") + .clone(); + let remote_id = codex_exec_server::REMOTE_ENVIRONMENT_ID; + let (attachment, _guard, _) = EnvironmentAttachment::new( + remote_id, + Arc::new(codex_exec_server::Environment::create_for_tests(Some( + "ws://127.0.0.1:1".into(), + ))?), + ); + remote.attachment = attachment; + remote.selection.environment_id = remote_id.to_string(); + remote.config_mut().permission_profile = + codex_protocol::models::PermissionProfileSnapshot::legacy( + codex_protocol::models::PermissionProfile::workspace_write(), + ); + remote.config_mut().network_policy = Some(EnvironmentNetworkPolicy::from_config( + &codex_network_proxy::NetworkProxyConfig { + enabled: true, + ..Default::default() + }, + /*managed_allowed_domains_only*/ false, + )); + let environments = TurnEnvironmentSnapshot { + environments: vec![TurnEnvironmentState::Ready(remote)], + }; + let remote = environments.primary().expect("remote environment"); + + for register_shared_listener in [false, true] { + let approval = begin_network_approval( + &session, + &turn, + &environments, + /*managed_network_active*/ true, + Some(NetworkApprovalSpec { + network: None, + environment_id: remote_id.to_string(), + tool_name: ToolName::plain("exec_command"), + command: "pwd".into(), + trigger: GuardianNetworkAccessTrigger { + call_id: "exec".into(), + tool_name: "exec_command".into(), + command: vec!["pwd".into()], + cwd: remote.cwd().clone(), + sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, + additional_permissions: None, + justification: None, + tty: None, + }, + permission_profile: remote.permission_profile().clone(), + network_policy: remote.config().network_policy.clone(), + }), + ) + .await + .expect("build network proxy") + .expect("network approval"); + let callback = approval + .execution_proxy() + .remote_policy_decider(/*allow_local_binding*/ true) + .expect("remote policy callback"); + let finish = finish_deferred_network_approval_after_process_exit_for_session( + Some(&session), + approval.into_deferred(), + ); + tokio::pin!(finish); + assert!(futures::poll!(&mut finish).is_pending()); + + let observer = register_shared_listener.then(|| { + build_blocked_request_observer(Arc::clone(&session.services.network_approval)) + }); + drop(callback); + if let Some(observer) = observer { + // A shared listener started during the command and can report an untagged denial. + tokio::time::advance(Duration::from_millis(10)).await; + assert!(futures::poll!(&mut finish).is_pending()); + let blocked = BlockedRequest::new(BlockedRequestArgs { + host: "blocked.example".into(), + protocol: "http".into(), + port: Some(80), + reason: "not_allowed".into(), + decision: Some("deny".into()), + source: Some("decider".into()), + client: None, + method: None, + mode: None, + }); + let expected = crate::network_policy_decision::denied_network_policy_message(&blocked) + .expect("blocked request must describe the denied host"); + observer.on_blocked_request(blocked).await; + assert_eq!(finish.await, Err(expected)); + } else { + assert_eq!(futures::poll!(&mut finish), std::task::Poll::Ready(Ok(()))); + } + } + + Ok(()) } #[tokio::test] diff --git a/codex-rs/network-proxy/src/proxy/execution_scope.rs b/codex-rs/network-proxy/src/proxy/execution_scope.rs index e01e115e78..1a61d78f53 100644 --- a/codex-rs/network-proxy/src/proxy/execution_scope.rs +++ b/codex-rs/network-proxy/src/proxy/execution_scope.rs @@ -17,6 +17,17 @@ impl Drop for ExecutionScope { } impl NetworkProxy { + /// Waits until every remote decision callback and in-flight request for this + /// execution has been dropped. Call after the executor has started, when no + /// further callbacks will be installed for this scope. An unscoped proxy + /// cannot establish this boundary, so its wait remains pending. + pub async fn wait_for_remote_policy_deciders(&self) { + match &self.execution_scope { + Some(scope) => scope.lifetime_tx.closed().await, + None => std::future::pending().await, + } + } + /// Returns a proxy with execution-specific attribution and attachment policy. pub fn for_execution( &self,