From 22a3f6d5d89c026c6b4f606ae5604f9e00059b23 Mon Sep 17 00:00:00 2001 From: acrognale-oai Date: Wed, 23 Sep 2026 00:30:17 +0000 Subject: [PATCH] Enforce application network policy across app-server requests (#47407) ## What changed - Load `application.network` requirements at startup and explicit config/account reloads, applying local policy before authentication and cloud configuration bootstrap. - Block traffic when policy loading fails, cancel requests that a new policy no longer permits, and reject configuration loads based on superseded policy snapshots. Local requirements edits take effect on the next explicit reload or restart. - Bind authenticated clients to the current account and revoke them when account ownership changes. Carry policy enforcement through backend requests, realtime connections, workload identity exchanges, and code-mode gRPC streams while preserving gRPC trailers. - Treat denied credential refreshes as policy errors without retrying them or invalidating stored credentials. ## Testing Add regression coverage for startup restrictions, reload ordering, cancellation after restrictive or malformed requirements edits, account revocation, gRPC destination denial, workload identity cancellation, and credential preservation after denied OAuth refresh. GitOrigin-RevId: 0cb4c720a906397e9df06675f01ecf60d741101f --- codex-rs/Cargo.lock | 2 + codex-rs/app-server/README.md | 7 + .../app-server/src/application_network.rs | 153 +++++++++++++++ .../src/application_network_tests.rs | 137 +++++++++++++ codex-rs/app-server/src/config_manager.rs | 107 +++++++++- .../src/config_manager_provider_tests.rs | 2 + .../src/config_manager_service_tests.rs | 1 + codex-rs/app-server/src/in_process.rs | 39 ++-- .../app-server/src/in_process_bootstrap.rs | 88 +++++++++ codex-rs/app-server/src/lib.rs | 64 +++--- .../request_processors/account_processor.rs | 27 ++- .../account_processor/rate_limit_resets.rs | 6 +- .../request_processors/feedback_processor.rs | 2 +- codex-rs/app-server/src/turn_cost_worker.rs | 25 ++- .../src/turn_cost_worker_chatgpt.rs | 3 +- codex-rs/app-server/tests/suite/logging.rs | 2 + codex-rs/app-server/tests/suite/v2/account.rs | 17 +- .../tests/suite/v2/account_thread_usage.rs | 5 + .../tests/suite/v2/application_network.rs | 109 +++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + .../suite/v2/rate_limits_identity_tests.rs | 33 +++- .../tests/suite/v2/remote_control_auth.rs | 1 + .../tests/suite/v2/workspace_routing.rs | 8 +- codex-rs/backend-client/src/client.rs | 36 +++- .../src/client_request_tests.rs | 11 +- codex-rs/cloud-config/src/backend.rs | 4 + codex-rs/cloud-config/src/bundle_loader.rs | 23 +-- codex-rs/cloud-config/src/service.rs | 44 ++++- codex-rs/cloud-config/src/service_tests.rs | 18 +- codex-rs/code-mode-host/Cargo.toml | 1 + codex-rs/code-mode-host/tests/grpc.rs | 2 + .../tests/grpc/network_policy_tests.rs | 96 +++++++++ .../code-mode/src/grpc_session/transport.rs | 183 ++++++++---------- codex-rs/config/src/cloud_config_bundle.rs | 20 +- codex-rs/config/src/loader/application.rs | 102 ++++++++++ codex-rs/config/src/loader/macos.rs | 14 ++ .../config/src/loader/managed_requirements.rs | 7 +- codex-rs/config/src/loader/mod.rs | 4 + codex-rs/core/src/client.rs | 1 + codex-rs/core/src/config/mod.rs | 18 +- codex-rs/core/src/realtime_conversation.rs | 13 +- codex-rs/core/tests/common/test_codex.rs | 45 +++-- .../core/tests/suite/auth_recovery_policy.rs | 93 +++++++++ codex-rs/core/tests/suite/mod.rs | 1 + .../tests/suite/realtime_sideband_endpoint.rs | 49 ++++- codex-rs/http-client/Cargo.toml | 1 + codex-rs/http-client/src/client_builder.rs | 47 ++++- codex-rs/http-client/src/client_tls.rs | 37 ++++ codex-rs/http-client/src/lib.rs | 2 + codex-rs/http-client/src/network_policy.rs | 10 + codex-rs/http-client/src/outbound_proxy.rs | 20 +- codex-rs/http-client/src/response.rs | 57 ++++++ .../src/route_aware_client_pool.rs | 25 ++- codex-rs/login/src/auth/change_state_tests.rs | 59 +++++- codex-rs/login/src/auth/default_client.rs | 20 +- codex-rs/login/src/auth/manager.rs | 53 ++++- codex-rs/login/src/auth/workload_identity.rs | 9 +- codex-rs/login/src/device_code_auth.rs | 12 +- codex-rs/login/src/oauth/client.rs | 8 +- codex-rs/login/src/oauth/client_tests.rs | 56 ++++++ codex-rs/login/src/outbound_proxy.rs | 22 +++ codex-rs/login/src/server.rs | 4 +- .../login/tests/suite/login_server_e2e.rs | 16 +- codex-rs/thread-manager-sample/src/main.rs | 2 + codex-rs/tui/src/app/background_requests.rs | 22 ++- codex-rs/tui/src/app/session_picker.rs | 4 + codex-rs/tui/src/app/tests.rs | 4 + codex-rs/tui/src/lib.rs | 8 +- codex-rs/tui/src/status/tests.rs | 29 ++- codex-rs/workload-identity/src/exchange.rs | 101 +++++----- codex-rs/workload-identity/src/lib.rs | 5 +- .../src/workload_identity_tests.rs | 92 +++++++++ 72 files changed, 1984 insertions(+), 365 deletions(-) create mode 100644 codex-rs/app-server/src/application_network.rs create mode 100644 codex-rs/app-server/src/application_network_tests.rs create mode 100644 codex-rs/app-server/src/in_process_bootstrap.rs create mode 100644 codex-rs/app-server/tests/suite/v2/application_network.rs create mode 100644 codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs create mode 100644 codex-rs/config/src/loader/application.rs create mode 100644 codex-rs/core/tests/suite/auth_recovery_policy.rs create mode 100644 codex-rs/http-client/src/client_tls.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 06a13bee36..d11ecb2830 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2681,6 +2681,7 @@ dependencies = [ "codex-code-mode", "codex-code-mode-protocol", "codex-code-mode-runtime", + "codex-http-client", "codex-otel", "codex-otel-trace-websocket", "codex-protocol", @@ -3593,6 +3594,7 @@ dependencies = [ "codex-utils-rustls-provider", "futures", "http 1.4.0", + "http-body 1.0.1", "native-tls", "opentelemetry", "opentelemetry_sdk", diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 67fa723cc7..8a572be0ab 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -30,6 +30,13 @@ the choice unset. Ephemeral threads cannot save it. Use `thread/metadata/update` for later changes. This preference does not select `turn/start.cyberAccessProgram` or grant access to an access program. +# Application network policy + +App-server loads application network policy at startup and existing explicit +config/account reloads. Local requirements-file edits take effect on +the next explicit reload or restart. Installing a new policy cancels requests +that it no longer permits; a failed policy load blocks network traffic. + # User verification cancellation (experimental) Local UI clients can cancel a native user-verification RPC by sending diff --git a/codex-rs/app-server/src/application_network.rs b/codex-rs/app-server/src/application_network.rs new file mode 100644 index 0000000000..acd74d8eb6 --- /dev/null +++ b/codex-rs/app-server/src/application_network.rs @@ -0,0 +1,153 @@ +//! Installs application requirements at startup and explicit config/account reloads. +//! Local policy limits cloud bootstrap; config construction rejects superseded policy snapshots. + +use super::ApplicationPolicyLoad; +use super::ApplicationPolicySnapshot; +use super::ConfigManager; +use codex_config::ApplicationRequirementsToml; +use codex_config::CloudConfigBundleLoader; +use codex_config::NetworkDomainPermissionToml; +use codex_http_client::DestinationPolicy; +use codex_http_client::NetworkPolicyDenied; +use std::io; +use std::sync::Arc; + +fn destination_policy(application: Option<&ApplicationRequirementsToml>) -> DestinationPolicy { + match application.and_then(|application| application.network.as_ref()) { + Some(network) if network.enabled => DestinationPolicy::Restricted { + allowed_hosts: network + .domains + .iter() + .filter(|&(_host, permission)| *permission == NetworkDomainPermissionToml::Allow) + .map(|(host, _permission)| host.clone()) + .collect(), + }, + Some(_) | None => DestinationPolicy::Unrestricted, + } +} + +impl ConfigManager { + pub(super) async fn refresh_local_network_policy( + &self, + ) -> io::Result { + let _reload = self + .network_policy_reload + .acquire() + .await + .map_err(io::Error::other)?; + self.load_local_network_policy(self.network_policy.policy().revision()) + .await + } + + async fn load_local_network_policy( + &self, + application_revision: codex_http_client::NetworkPolicyRevision, + ) -> io::Result { + let local_revision = self.local_network_policy.policy().revision(); + let result = async { + let sources = codex_config::loader::load_local_application_requirements( + codex_exec_server::LOCAL_FS.as_ref(), + &self.loader_overrides, + ) + .await?; + let application = sources.compose(Default::default())?; + self.local_network_policy + .publish(local_revision, destination_policy(application.as_ref())); + Ok(sources) + } + .await; + if result.is_err() { + self.local_network_policy.unavailable(local_revision); + self.fail_application_policy(application_revision)?; + } + result + } + + pub(crate) async fn refresh_application_network_policy( + &self, + ) -> io::Result { + let _reload = self + .network_policy_reload + .acquire() + .await + .map_err(io::Error::other)?; + let revision = self.network_policy.policy().revision(); + let local_sources = self.load_local_network_policy(revision).await?; + let cloud_config = self.current_cloud_config_bundle(); + let result = async { + let cloud = cloud_config.get().await.map_err(io::Error::other)?; + let application = local_sources.compose( + cloud + .as_ref() + .map(|bundle| bundle.requirements_toml.clone()) + .unwrap_or_default(), + )?; + Ok::<_, io::Error>((destination_policy(application.as_ref()), cloud)) + } + .await; + match result { + Ok((policy, cloud)) => { + let mut current = self.network_policy_snapshot.write().map_err(|_| { + io::Error::other("application network policy snapshot lock poisoned") + })?; + if !self.network_policy.publish(revision, policy.clone()) { + *current = None; + return Err(io::Error::other(NetworkPolicyDenied::Revoked)); + } + let next = ApplicationPolicySnapshot { + revision, + policy, + cloud: cloud.clone(), + }; + // Concurrent config loads may finish together if their actual inputs are unchanged. + let snapshot = match current.as_ref() { + Some(snapshot) if snapshot.as_ref() == &next => Arc::clone(snapshot), + _ => { + let snapshot = Arc::new(next); + *current = Some(Arc::clone(&snapshot)); + snapshot + } + }; + Ok(ApplicationPolicyLoad { + cloud_config: CloudConfigBundleLoader::new(async move { Ok(cloud) }), + snapshot, + }) + } + Err(error) => { + self.fail_application_policy(revision)?; + Err(error) + } + } + } + + pub(super) fn check_application_policy_load( + &self, + load: &ApplicationPolicyLoad, + ) -> io::Result<()> { + let current = self + .network_policy_snapshot + .read() + .map_err(|_| io::Error::other("application network policy snapshot lock poisoned"))?; + if self.network_policy.policy().revision() != load.snapshot.revision + || !current + .as_ref() + .is_some_and(|snapshot| Arc::ptr_eq(snapshot, &load.snapshot)) + { + return Err(io::Error::other(NetworkPolicyDenied::Revoked)); + } + Ok(()) + } + + fn fail_application_policy( + &self, + revision: codex_http_client::NetworkPolicyRevision, + ) -> io::Result<()> { + let mut current = self + .network_policy_snapshot + .write() + .map_err(|_| io::Error::other("application network policy snapshot lock poisoned"))?; + self.network_policy.unavailable(revision); + *current = None; + Ok(()) + } +} diff --git a/codex-rs/app-server/src/application_network_tests.rs b/codex-rs/app-server/src/application_network_tests.rs new file mode 100644 index 0000000000..c6a513b45f --- /dev/null +++ b/codex-rs/app-server/src/application_network_tests.rs @@ -0,0 +1,137 @@ +//! Checks that overlapping policy reloads cannot publish out of order. + +use super::ConfigManager; +use anyhow::Result; +use codex_config::CloudConfigBundleLoader; +use codex_config::LoaderOverrides; +use codex_config::test_support::CloudConfigBundleFixture; +use codex_http_client::NetworkPolicyDenied; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; +use tempfile::tempdir; +use tokio::sync::Semaphore; +use tokio::time::timeout; + +#[tokio::test] +async fn overlapping_policy_reloads_publish_in_order() -> Result<()> { + let home = tempdir()?; + let entered = Arc::new(Semaphore::new(/*permits*/ 0)); + let release = Arc::new(Semaphore::new(/*permits*/ 0)); + let calls = Arc::new(AtomicUsize::new(/*v*/ 0)); + let loader = CloudConfigBundleLoader::from_getter({ + let (entered, release, calls) = (entered.clone(), release.clone(), calls.clone()); + move || { + let (entered, release, calls) = (entered.clone(), release.clone(), calls.clone()); + async move { + let call = calls.fetch_add(/*val*/ 1, Ordering::SeqCst); + let host = if call == 0 { + entered.add_permits(/*n*/ 1); + release + .acquire() + .await + .expect("release first reload") + .forget(); + "old.example" + } else if call == 3 { + "old.example" + } else { + "new.example" + }; + Ok(Some( + CloudConfigBundleFixture::enterprise_requirement(format!( + "[application.network.domains]\n'{host}' = 'allow'" + )) + .into_bundle(), + )) + } + } + }); + let manager = ConfigManager::new_for_tests( + home.path().to_path_buf(), + Vec::new(), + LoaderOverrides::without_managed_config_for_tests(), + loader, + ); + let first = tokio::spawn({ + let manager = manager.clone(); + async move { manager.refresh_application_network_policy().await } + }); + entered.acquire().await?.forget(); + let mut second = tokio::spawn({ + let manager = manager.clone(); + async move { manager.refresh_application_network_policy().await } + }); + assert!( + timeout(Duration::from_millis(/*millis*/ 100), &mut second) + .await + .is_err() + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + release.add_permits(/*n*/ 1); + let first_load = first.await??; + let second_load = timeout(Duration::from_secs(/*secs*/ 5), second).await???; + assert_eq!(calls.load(Ordering::SeqCst), 2); + manager.check_application_policy_load(&second_load)?; + let stale = manager + .check_application_policy_load(&first_load) + .unwrap_err(); + assert_eq!( + stale.get_ref().and_then(|error| error.downcast_ref()), + Some(&NetworkPolicyDenied::Revoked) + ); + + let policy = manager.network_policy.policy(); + assert_eq!( + policy.acquire(&"https://old.example".parse()?).map(|_| ()), + Err(NetworkPolicyDenied::Destination) + ); + assert_eq!( + policy.acquire(&"https://new.example".parse()?).map(|_| ()), + Ok(()) + ); + let unchanged = manager.refresh_application_network_policy().await?; + manager.check_application_policy_load(&second_load)?; + manager.check_application_policy_load(&unchanged)?; + let restored = manager.refresh_application_network_policy().await?; + manager.check_application_policy_load(&restored)?; + assert!(manager.check_application_policy_load(&first_load).is_err()); + assert!(manager.check_application_policy_load(&second_load).is_err()); + Ok(()) +} + +#[tokio::test] +async fn cloud_config_change_supersedes_load_even_when_network_rules_are_unchanged() -> Result<()> { + let home = tempdir()?; + let calls = AtomicUsize::new(/*v*/ 0); + let loader = CloudConfigBundleLoader::from_getter(move || { + let model = if calls.fetch_add(/*val*/ 1, Ordering::SeqCst) == 0 { + "old" + } else { + "new" + }; + async move { + Ok(Some( + CloudConfigBundleFixture::enterprise_config(format!("model = '{model}'")) + .into_bundle(), + )) + } + }); + let manager = ConfigManager::new_for_tests( + home.path().to_path_buf(), + Vec::new(), + LoaderOverrides::without_managed_config_for_tests(), + loader, + ); + let first = manager.refresh_application_network_policy().await?; + let second = manager.refresh_application_network_policy().await?; + manager.check_application_policy_load(&second)?; + let stale = manager.check_application_policy_load(&first).unwrap_err(); + assert_eq!( + stale.get_ref().and_then(|error| error.downcast_ref()), + Some(&NetworkPolicyDenied::Revoked) + ); + Ok(()) +} diff --git a/codex-rs/app-server/src/config_manager.rs b/codex-rs/app-server/src/config_manager.rs index cb2c6d3023..f952414820 100644 --- a/codex-rs/app-server/src/config_manager.rs +++ b/codex-rs/app-server/src/config_manager.rs @@ -30,6 +30,9 @@ use toml::Value as TomlValue; use tracing::instrument; use tracing::warn; +#[path = "application_network.rs"] +mod application_network; + #[derive(Debug, thiserror::Error)] #[error( "Your organization's required model provider settings changed. Restart Codex to apply them; this request was not sent" @@ -47,6 +50,23 @@ pub(crate) struct ConfigManager { cloud_config_bundle: Arc>, arg0_paths: Arg0DispatchPaths, thread_config_loader: Arc, + network_policy: codex_http_client::NetworkPolicyController, + local_network_policy: codex_http_client::NetworkPolicyController, + network_policy_reload: Arc, + network_policy_snapshot: Arc>>>, +} + +/// Configuration and policy must finish loading against the same account and cloud snapshot. +pub(crate) struct ApplicationPolicyLoad { + cloud_config: CloudConfigBundleLoader, + snapshot: Arc, +} + +#[derive(PartialEq, Eq)] +struct ApplicationPolicySnapshot { + revision: codex_http_client::NetworkPolicyRevision, + policy: codex_http_client::DestinationPolicy, + cloud: Option, } impl ConfigManager { @@ -59,6 +79,7 @@ impl ConfigManager { arg0_paths: Arg0DispatchPaths, thread_config_loader: Arc, ) -> Self { + let network_policy = codex_http_client::NetworkPolicyController::default(); Self { codex_home, cli_overrides: Arc::new(RwLock::new(cli_overrides)), @@ -68,6 +89,10 @@ impl ConfigManager { cloud_config_bundle: Arc::new(RwLock::new(cloud_config_bundle)), arg0_paths, thread_config_loader, + network_policy, + local_network_policy: Default::default(), + network_policy_reload: Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1)), + network_policy_snapshot: Arc::default(), } } @@ -109,6 +134,16 @@ impl ConfigManager { chatgpt_base_url: String, http_client_factory: codex_http_client::HttpClientFactory, ) { + let endpoint = codex_backend_client::Client::new( + chatgpt_base_url.clone(), + http_client_factory.clone(), + ) + .config_bundle_url(); + let http_client_factory = http_client_factory.with_network_policy( + self.local_network_policy + .policy() + .restrict_to_endpoints(endpoint.parse().into_iter().collect()), + ); let loader = cloud_config_bundle_loader( auth_manager, chatgpt_base_url, @@ -181,6 +216,8 @@ impl ConfigManager { .transpose()?, ) .await?; + config.application_network_policy = refreshed_config.application_network_policy; + config.application_auth_route_config = refreshed_config.application_auth_route_config; self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); Ok(config) @@ -206,6 +243,7 @@ impl ConfigManager { /*default_zsh_path*/ None, ) .await?; + self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); Ok(config) @@ -216,6 +254,7 @@ impl ConfigManager { &self, current: &Config, ) -> std::io::Result<()> { + let policy_load = self.refresh_application_network_policy().await?; // Existing threads retain their session route; only managed // requirements can invalidate it. let requirements = load_managed_requirements_state( @@ -224,10 +263,11 @@ impl ConfigManager { codex_config::ConfigLoadOptions { loader_overrides: self.loader_overrides.clone(), strict_config: self.strict_config, - cloud_config_bundle: self.current_cloud_config_bundle(), + cloud_config_bundle: policy_load.cloud_config.clone(), }, ) .await?; + self.check_application_policy_load(&policy_load)?; let selection_changed = requirements .model_provider .as_ref() @@ -267,9 +307,33 @@ impl ConfigManager { Ok(()) } + /// Installs local policy before cloud access; failed policy loads block traffic while + /// non-strict startup remains available for configuration repair. + pub(crate) async fn load_startup_config( + &self, + fallback_cwd: Option, + ) -> std::io::Result { + match self.load_latest_config(fallback_cwd).await { + Ok(config) => Ok(config), + Err(error) + if self.strict_config + || crate::is_unsupported_untrusted_approval_policy_error(&error) => + { + Err(error) + } + Err(error) => { + warn!(%error, "configuration is unavailable; using default settings"); + self.load_default_config().await + } + } + } + pub(crate) async fn load_default_config(&self) -> std::io::Result { let mut loader_overrides = self.loader_overrides.clone(); loader_overrides.ignore_user_config = true; + loader_overrides.ignore_project_config = true; + loader_overrides.ignore_managed_requirements |= + self.refresh_local_network_policy().await.is_err(); let mut config = ConfigBuilder::default() .codex_home(self.codex_home.clone()) .cli_overrides(self.current_cli_overrides()) @@ -278,6 +342,7 @@ impl ConfigManager { .cloud_config_bundle(CloudConfigBundleLoader::default()) .build() .await?; + self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); Ok(config) @@ -359,6 +424,7 @@ impl ConfigManager { mut typesafe_overrides: ConfigOverrides, fallback_cwd: Option, ) -> std::io::Result { + let policy_load = self.refresh_application_network_policy().await?; let mut request_overrides = request_overrides.unwrap_or_default(); if let Some(value) = request_overrides.remove("bypass_hook_trust") { typesafe_overrides.bypass_hook_trust = Some(value.as_bool().ok_or_else(|| { @@ -377,17 +443,20 @@ impl ConfigManager { .map(|(key, value)| (key, json_to_toml(value))), ) .collect::>(); - let mut config = codex_core::config::ConfigBuilder::default() + let result = codex_core::config::ConfigBuilder::default() .codex_home(self.codex_home.clone()) .cli_overrides(merged_cli_overrides) .loader_overrides(self.loader_overrides.clone()) .strict_config(self.strict_config) .harness_overrides(typesafe_overrides) .fallback_cwd(fallback_cwd) - .cloud_config_bundle(self.current_cloud_config_bundle()) + .cloud_config_bundle(policy_load.cloud_config.clone()) .thread_config_loader(Arc::clone(&self.thread_config_loader)) .build() - .await?; + .await; + let mut config = result?; + self.check_application_policy_load(&policy_load)?; + self.apply_network_policy(&mut config); self.apply_runtime_feature_enablement(&mut config); self.apply_arg0_paths(&mut config); Ok(config) @@ -404,7 +473,8 @@ impl ConfigManager { &self, cwd: Option, ) -> std::io::Result { - load_config_layers_state( + let policy_load = self.refresh_application_network_policy().await?; + let result = load_config_layers_state( LOCAL_FS.as_ref(), &self.codex_home, cwd, @@ -412,11 +482,14 @@ impl ConfigManager { codex_config::ConfigLoadOptions { loader_overrides: self.loader_overrides.clone(), strict_config: self.strict_config, - cloud_config_bundle: self.current_cloud_config_bundle(), + cloud_config_bundle: policy_load.cloud_config.clone(), }, self.thread_config_loader.as_ref(), ) - .await + .await; + let layers = result?; + self.check_application_policy_load(&policy_load)?; + Ok(layers) } fn apply_runtime_feature_enablement(&self, config: &mut Config) { @@ -430,6 +503,22 @@ impl ConfigManager { .unwrap_or_default() } + fn apply_network_policy(&self, config: &mut Config) { + config.application_network_policy = self.network_policy.policy(); + config.application_auth_route_config = Some( + codex_login::AuthRouteConfig::from_http_client_factory( + config + .http_client_factory() + .with_network_policy(self.network_policy.policy()), + ) + .with_local_bootstrap_factory( + config + .http_client_factory() + .with_network_policy(self.local_network_policy.policy()), + ), + ); + } + fn apply_arg0_paths(&self, config: &mut Config) { config.codex_self_exe = self.arg0_paths.codex_self_exe.clone(); config.codex_linux_sandbox_exe = self.arg0_paths.codex_linux_sandbox_exe.clone(); @@ -465,6 +554,10 @@ impl ConfigManager { } } +#[cfg(test)] +#[path = "application_network_tests.rs"] +mod application_network_tests; + #[cfg(test)] #[path = "config_manager_provider_tests.rs"] mod provider_tests; diff --git a/codex-rs/app-server/src/config_manager_provider_tests.rs b/codex-rs/app-server/src/config_manager_provider_tests.rs index 253b9018d5..20fbfb4989 100644 --- a/codex-rs/app-server/src/config_manager_provider_tests.rs +++ b/codex-rs/app-server/src/config_manager_provider_tests.rs @@ -94,6 +94,8 @@ async fn provider_requirements_do_not_reload_thread_config() -> Result<()> { .load_retained_session_config(¤t.config_layer_stack, ¤t.cwd) .await?; assert_eq!(retained.model_provider, current.model_provider); + let factory = current.http_client_factory(); + assert_eq!(retained.http_client_factory(), factory); manager.cloud_config_bundle = Arc::new(RwLock::new( CloudConfigBundleFixture::loader_with_enterprise_requirement("model_provider = 'other'"), diff --git a/codex-rs/app-server/src/config_manager_service_tests.rs b/codex-rs/app-server/src/config_manager_service_tests.rs index 639839d706..c3947b0e76 100644 --- a/codex-rs/app-server/src/config_manager_service_tests.rs +++ b/codex-rs/app-server/src/config_manager_service_tests.rs @@ -737,6 +737,7 @@ async fn psp_feature_configures_first_party_routing() -> Result<()> { config.http_client_factory(), HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) .with_system_proxy_fallback() + .with_network_policy(config.application_network_policy.clone()) .with_chatgpt_cookies([HeaderValue::from_static("oai-chat-psp=true")]) ); assert_eq!( diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 20e44d3253..e5896139f6 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -89,7 +89,6 @@ use codex_core::config::Config; use codex_core::resolve_installation_id; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; -use codex_login::AuthManager; use codex_protocol::protocol::SessionSource; pub use codex_rollout::StateDbHandle; pub use codex_state::log_db::LogDbLayer; @@ -99,6 +98,9 @@ use tokio::time::timeout; use toml::Value as TomlValue; use tracing::warn; +#[path = "in_process_bootstrap.rs"] +mod bootstrap; + const IN_PROCESS_CONNECTION_ID: ConnectionId = ConnectionId(0); const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); // Covers both bounded runtime drains plus the analytics client's 25-second best-effort flush. @@ -158,7 +160,7 @@ pub struct InProcessStartArgs { pub config_warnings: Vec, /// Session source stamped into thread/session metadata. pub session_source: SessionSource, - /// Whether auth loading should honor the `CODEX_API_KEY` environment variable. + /// Whether serving auth should honor `CODEX_API_KEY`; workspace policy still uses stored auth. pub enable_codex_api_key_env: bool, /// Initialize params used for initial handshake. pub initialize: InitializeParams, @@ -380,7 +382,7 @@ pub async fn start(mut args: InProcessStartArgs) -> IoResult IoResult { - args.config.auth_config().validate()?; +async fn start_uninitialized(mut args: InProcessStartArgs) -> IoResult { + let config_manager = ConfigManager::new( + args.config.codex_home.to_path_buf(), + args.cli_overrides, + args.loader_overrides, + args.strict_config, + args.cloud_config_bundle, + args.arg0_paths.clone(), + args.thread_config_loader, + ); + let auth_manager = bootstrap::configure( + &config_manager, + &mut args.config, + args.enable_codex_api_key_env, + ) + .await?; let channel_capacity = args.channel_capacity.max(1); let installation_id = resolve_installation_id(&args.config.codex_home).await?; - let auth_manager = - AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env) - .await - .map_err(IoError::other)?; let (client_tx, mut client_rx) = mpsc::channel::(channel_capacity); let (event_tx, event_rx) = mpsc::channel::(channel_capacity); @@ -464,15 +476,6 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult(channel_capacity); let mut processor_handle = tokio::spawn(async move { let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { diff --git a/codex-rs/app-server/src/in_process_bootstrap.rs b/codex-rs/app-server/src/in_process_bootstrap.rs new file mode 100644 index 0000000000..99f52100f7 --- /dev/null +++ b/codex-rs/app-server/src/in_process_bootstrap.rs @@ -0,0 +1,88 @@ +//! Loads embedded application policy and builds serving authentication from effective requirements. +//! Workspace policy uses stored auth even when the caller uses an environment API key. +//! Caller credential storage remains intact; initialization receives the effective URL and residency. + +use crate::config_manager::ConfigManager; +use codex_core::config::Config; +use codex_login::AuthManager; +use std::io::Error as IoError; +use std::io::Result as IoResult; +use std::sync::Arc; + +pub(super) async fn configure( + config_manager: &ConfigManager, + config: &mut Arc, + enable_codex_api_key_env: bool, +) -> IoResult> { + let bootstrap_config = config_manager + .load_startup_config(Some(config.cwd.to_path_buf())) + .await?; + let caller_auth = config.auth_config(); + let bootstrap_auth = codex_login::AuthConfig { + chatgpt_base_url: Some(bootstrap_config.chatgpt_base_url.clone()), + forced_login_method: bootstrap_config.forced_login_method, + forced_chatgpt_workspace_id: bootstrap_config.forced_chatgpt_workspace_id.clone(), + managed_auth_policy: bootstrap_config + .config_layer_stack + .requirements() + .managed_auth_policy(), + auth_route_config: bootstrap_config.auth_route_config(), + ..caller_auth.clone() + }; + bootstrap_auth.validate()?; + let bootstrap_auth_manager = AuthManager::shared_from_auth_config( + bootstrap_auth, + /*enable_codex_api_key_env*/ false, + ) + .await + .map_err(IoError::other)?; + config_manager.replace_cloud_config_bundle_loader( + bootstrap_auth_manager, + bootstrap_config.chatgpt_base_url.clone(), + bootstrap_config.http_client_factory(), + ); + let policy_config = config_manager + .load_startup_config(Some(config.cwd.to_path_buf())) + .await?; + let auth_config = codex_login::AuthConfig { + chatgpt_base_url: Some(policy_config.chatgpt_base_url.clone()), + forced_login_method: policy_config.forced_login_method, + forced_chatgpt_workspace_id: policy_config.forced_chatgpt_workspace_id.clone(), + managed_auth_policy: policy_config + .config_layer_stack + .requirements() + .managed_auth_policy(), + auth_route_config: policy_config.auth_route_config(), + ..caller_auth + }; + auth_config.validate()?; + let policy_auth_manager = AuthManager::shared_from_auth_config( + auth_config.clone(), + /*enable_codex_api_key_env*/ false, + ) + .await + .map_err(IoError::other)?; + let auth_manager = if enable_codex_api_key_env { + AuthManager::shared_from_auth_config(auth_config, /*enable_codex_api_key_env*/ true) + .await + .map_err(IoError::other)? + } else { + policy_auth_manager.clone() + }; + let config = Arc::make_mut(config); + config.chatgpt_base_url = policy_config.chatgpt_base_url; + config.enforce_residency = policy_config.enforce_residency; + config.application_network_policy = policy_config.application_network_policy; + config.application_auth_route_config = policy_config.application_auth_route_config; + config.forced_login_method = policy_config.forced_login_method; + config.forced_chatgpt_workspace_id = policy_config.forced_chatgpt_workspace_id; + config_manager.replace_cloud_config_bundle_loader( + policy_auth_manager, + config.chatgpt_base_url.clone(), + config.http_client_factory(), + ); + config_manager + .sync_default_client_residency_requirement() + .await; + Ok(auth_manager) +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 08d3232b23..fc0fb1d9d0 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -533,30 +533,18 @@ pub async fn run_main_with_transport_options( arg0_paths.clone(), Arc::new(NoopThreadConfigLoader), ); - match config_manager - .load_latest_config(/*fallback_cwd*/ None) - .await - { - Ok(config) => { - let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false) - .await - .map_err(std::io::Error::other)?; - config_manager.replace_cloud_config_bundle_loader( - auth_manager, - config.chatgpt_base_url.clone(), - config.http_client_factory(), - ); - } - Err(err) if is_unsupported_untrusted_approval_policy_error(&err) => { - return Err(err); - } - Err(err) => { - warn!(error = %err, "Failed to preload config for cloud config bundle"); - // If this fails, we cannot install cloud/thread config loaders, so non-strict - // startup continues without managed cloud config. - } - }; + let bootstrap_config = config_manager + .load_startup_config(/*fallback_cwd*/ None) + .await?; + let bootstrap_auth = + AuthManager::shared_from_config(&bootstrap_config, /*enable_codex_api_key_env*/ false) + .await + .map_err(std::io::Error::other)?; + config_manager.replace_cloud_config_bundle_loader( + bootstrap_auth, + bootstrap_config.chatgpt_base_url.clone(), + bootstrap_config.http_client_factory(), + ); let mut config_warnings = Vec::new(); let mut plugin_startup_config = PluginStartupConfig::Current; let config = match config_manager @@ -584,6 +572,19 @@ pub async fn run_main_with_transport_options( } }; config.auth_config().validate()?; + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false) + .await + .map_err(std::io::Error::other)?; + config_manager.replace_cloud_config_bundle_loader( + auth_manager.clone(), + config.chatgpt_base_url.clone(), + config.http_client_factory(), + ); + config_manager + .sync_default_client_residency_requirement() + .await; + #[cfg(target_os = "macos")] let local_runtime_paths = local_runtime_paths.with_allowed_symlinked_codex_home( codex_config::allowed_symlinked_codex_home(&config.config_layer_stack, &config.codex_home), @@ -606,13 +607,19 @@ pub async fn run_main_with_transport_options( )) } }; + // Executor credentials belong to the selected environment, not the current account. + // Each request and connection still acquires a revocable application-policy permit. + let environment_http_client_factory = config + .http_client_factory() + .with_network_policy(config.application_network_policy.clone()); let environment_manager = if ignore_user_config { - EnvironmentManager::from_env(Some(local_runtime_paths), config.http_client_factory()).await + EnvironmentManager::from_env(Some(local_runtime_paths), environment_http_client_factory) + .await } else { EnvironmentManager::from_codex_home( codex_home.clone(), Some(local_runtime_paths), - config.http_client_factory(), + environment_http_client_factory, ) .await } @@ -803,11 +810,6 @@ pub async fn run_main_with_transport_options( } drop(unix_socket_startup_lock); - let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false) - .await - .map_err(std::io::Error::other)?; - let remote_control_enabled = remote_control_policy == RemoteControlPolicy::Allowed && remote_control_explicitly_requested && state_db.is_some(); diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index f30658f3d5..95ed6d571d 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -976,11 +976,12 @@ impl AccountRequestProcessor { } } + self.config_manager.clear_cloud_config_bundle_loader(); + if config.model_provider.is_amazon_bedrock() { clear_user_model_provider_if_bedrock(&self.config_manager, &config).await?; } - self.config_manager.clear_cloud_config_bundle_loader(); *self.workspace_routing.lock().await = None; Self::maybe_refresh_plugin_caches_for_current_config( @@ -1115,7 +1116,9 @@ impl AccountRequestProcessor { &self, params: GetAccountRateLimitsParams, ) -> Result { - let Some(auth) = self.auth_manager.auth().await else { + let Some((auth, http_client_factory)) = + self.auth_manager.auth_with_http_client_factory().await + else { return Err(invalid_request( "codex account authentication required to read rate limits", )); @@ -1130,7 +1133,7 @@ impl AccountRequestProcessor { let client = BackendClient::from_auth( self.config.chatgpt_base_url.clone(), &auth, - self.config.http_client_factory(), + http_client_factory, ); let usage_request = async { @@ -1230,7 +1233,9 @@ impl AccountRequestProcessor { }) .transpose()?; - let Some(auth) = self.auth_manager.auth().await else { + let Some((auth, http_client_factory)) = + self.auth_manager.auth_with_http_client_factory().await + else { return Err(invalid_request( "codex account authentication required to read token usage", )); @@ -1245,7 +1250,7 @@ impl AccountRequestProcessor { let client = BackendClient::from_auth( self.config.chatgpt_base_url.clone(), &auth, - self.config.http_client_factory(), + http_client_factory, ); if let Some(thread_id) = thread_id { let thread_id = thread_id.to_string(); @@ -1315,7 +1320,9 @@ impl AccountRequestProcessor { async fn get_workspace_messages_response( &self, ) -> Result { - let Some(auth) = self.auth_manager.auth().await else { + let Some((auth, http_client_factory)) = + self.auth_manager.auth_with_http_client_factory().await + else { return Err(invalid_request( "codex account authentication required to read workspace messages", )); @@ -1330,7 +1337,7 @@ impl AccountRequestProcessor { let client = BackendClient::from_auth( self.config.chatgpt_base_url.clone(), &auth, - self.config.http_client_factory(), + http_client_factory, ); let messages = tokio::time::timeout( ACCOUNT_WORKSPACE_MESSAGES_FETCH_TIMEOUT, @@ -1407,7 +1414,9 @@ impl AccountRequestProcessor { &self, params: SendAddCreditsNudgeEmailParams, ) -> Result { - let Some(auth) = self.auth_manager.auth().await else { + let Some((auth, http_client_factory)) = + self.auth_manager.auth_with_http_client_factory().await + else { return Err(invalid_request( "codex account authentication required to notify workspace owner", )); @@ -1422,7 +1431,7 @@ impl AccountRequestProcessor { let client = BackendClient::from_auth( self.config.chatgpt_base_url.clone(), &auth, - self.config.http_client_factory(), + http_client_factory, ); match client diff --git a/codex-rs/app-server/src/request_processors/account_processor/rate_limit_resets.rs b/codex-rs/app-server/src/request_processors/account_processor/rate_limit_resets.rs index 4c7930c10a..c923429f56 100644 --- a/codex-rs/app-server/src/request_processors/account_processor/rate_limit_resets.rs +++ b/codex-rs/app-server/src/request_processors/account_processor/rate_limit_resets.rs @@ -98,7 +98,9 @@ impl AccountRequestProcessor { } async fn rate_limit_reset_backend_client(&self) -> Result { - let Some(auth) = self.auth_manager.auth().await else { + let Some((auth, http_client_factory)) = + self.auth_manager.auth_with_http_client_factory().await + else { return Err(invalid_request( "codex account authentication required for rate limit reset credits", )); @@ -112,7 +114,7 @@ impl AccountRequestProcessor { Ok(BackendClient::from_auth( self.config.chatgpt_base_url.clone(), &auth, - self.config.http_client_factory(), + http_client_factory, )) } } diff --git a/codex-rs/app-server/src/request_processors/feedback_processor.rs b/codex-rs/app-server/src/request_processors/feedback_processor.rs index fa5fe4ec68..62ba96e9bc 100644 --- a/codex-rs/app-server/src/request_processors/feedback_processor.rs +++ b/codex-rs/app-server/src/request_processors/feedback_processor.rs @@ -93,6 +93,7 @@ impl FeedbackRequestProcessor { None => None, }; + let http_client_factory = self.config.http_client_factory(); let auth = self.auth_manager.auth_cached(); let turn_metadata = if let Some(conversation_id) = conversation_id && let Some(rollout_path) = self @@ -281,7 +282,6 @@ impl FeedbackRequestProcessor { } let session_source = self.thread_manager.session_source(); - let http_client_factory = self.config.http_client_factory(); let runtime_handle = tokio::runtime::Handle::current(); let upload_result = tokio::task::spawn_blocking(move || { diff --git a/codex-rs/app-server/src/turn_cost_worker.rs b/codex-rs/app-server/src/turn_cost_worker.rs index b11efb9eef..c36108b4d5 100644 --- a/codex-rs/app-server/src/turn_cost_worker.rs +++ b/codex-rs/app-server/src/turn_cost_worker.rs @@ -400,12 +400,18 @@ impl WorkerRuntime { ) -> Result>, RequestError> { match &self.backend { TurnCostBackend::OpenAi(auth_manager) => { - let Some(auth) = auth_manager.auth().await else { + let Some((auth, captured_factory)) = + auth_manager.auth_with_http_client_factory().await + else { return Ok(None); }; + let http_client_factory = self + .config + .http_client_factory() + .with_network_policy(captured_factory.network_policy().clone()); if auth.is_chatgpt_auth() { return self - .query_chatgpt_turn_costs(&auth, turn_ids) + .query_chatgpt_turn_costs(&auth, http_client_factory, turn_ids) .await .map(Some); } @@ -420,7 +426,7 @@ impl WorkerRuntime { let client = BackendClient::from_auth( self.config.chatgpt_base_url.clone(), &auth, - self.config.http_client_factory(), + http_client_factory, ); client .query_api_key_turn_costs(turn_ids, &provider.headers) @@ -428,6 +434,12 @@ impl WorkerRuntime { .map(Some) } TurnCostBackend::ModelProvider(model_provider) => { + let http_client_factory = self.config.http_client_factory().with_network_policy( + self.config + .application_network_policy + .clone() + .for_current_account(), + ); if model_provider.info().requires_openai_auth { let Some(auth) = model_provider.auth().await else { return Ok(None); @@ -445,11 +457,8 @@ impl WorkerRuntime { .await .map_err(|error| RequestError::Other(error.into()))?; let endpoint = provider.url_for_path("analytics/codex/turn-costs"); - let client = BackendClient::new( - provider.base_url.clone(), - self.config.http_client_factory(), - ) - .with_auth_provider(auth); + let client = BackendClient::new(provider.base_url.clone(), http_client_factory) + .with_auth_provider(auth); client .query_api_key_turn_costs_at(&endpoint, turn_ids, &provider.headers) .await diff --git a/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs b/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs index 4c300e0bcc..a9bd6834b2 100644 --- a/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs +++ b/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs @@ -14,6 +14,7 @@ impl WorkerRuntime { pub(super) async fn query_chatgpt_turn_costs( &self, auth: &CodexAuth, + http_client_factory: codex_http_client::HttpClientFactory, turn_ids: &[String], ) -> Result, RequestError> { let mut threads: BTreeMap> = BTreeMap::new(); @@ -32,7 +33,7 @@ impl WorkerRuntime { let client = BackendClient::from_auth( self.config.chatgpt_base_url.clone(), auth, - self.config.http_client_factory(), + http_client_factory, ); let costs = client.query_chatgpt_turn_costs(&threads).await?; let mut settled_costs = Vec::new(); diff --git a/codex-rs/app-server/tests/suite/logging.rs b/codex-rs/app-server/tests/suite/logging.rs index 9c62bb5b88..67da665568 100644 --- a/codex-rs/app-server/tests/suite/logging.rs +++ b/codex-rs/app-server/tests/suite/logging.rs @@ -59,11 +59,13 @@ async fn credentials_stay_out_of_persisted_and_feedback_logs() -> Result<()> { let initial_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("initial@example.com") + .chatgpt_user_id("logging-user") .chatgpt_account_id(account_id), )?; let refreshed_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("refreshed@example.com") + .chatgpt_user_id("logging-user") .chatgpt_account_id(account_id), )?; let server = MockServer::start().await; diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 873e2662cd..7063ec6ac5 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -79,7 +79,6 @@ const WORKSPACE_ID_SECOND_ALLOWED: &str = "123e4567-e89b-42d3-a456-426614174001" const WORKSPACE_ID_DISALLOWED: &str = "123e4567-e89b-42d3-a456-426614174002"; const WORKSPACE_ID_EMBEDDED: &str = "123e4567-e89b-42d3-a456-426614174010"; const WORKSPACE_ID_INITIAL: &str = "123e4567-e89b-42d3-a456-426614174011"; -const WORKSPACE_ID_REFRESHED: &str = "123e4567-e89b-42d3-a456-426614174012"; const WORKSPACE_ID_DEVICE: &str = "123e4567-e89b-42d3-a456-426614174013"; const WORKSPACE_ID_STALE: &str = "123e4567-e89b-42d3-a456-426614174014"; @@ -659,19 +658,15 @@ async fn external_auth_refreshes_on_unauthorized(model_path: &str) -> Result<()> let initial_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("initial@example.com") - .chatgpt_user_id("user") + .chatgpt_user_id("refresh-user") .plan_type("pro") .chatgpt_account_id(WORKSPACE_ID_INITIAL), )?; - let refreshed_workspace = if model_path == "/v1" { - WORKSPACE_ID_REFRESHED - } else { - WORKSPACE_ID_INITIAL - }; + let refreshed_workspace = WORKSPACE_ID_INITIAL; let refreshed_access_token = encode_id_token( &ChatGptIdTokenClaims::new() .email("refreshed@example.com") - .chatgpt_user_id("user") + .chatgpt_user_id("refresh-user") .plan_type("pro") .chatgpt_account_id(refreshed_workspace), )?; @@ -747,12 +742,16 @@ async fn external_auth_refreshes_on_unauthorized(model_path: &str) -> Result<()> .await?; let _: codex_app_server_protocol::TurnStartResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(turn_req)).await??; - let _turn_completed = timeout( + let turn_completed = timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), ) .await??; + let completed: TurnCompletedNotification = + serde_json::from_value(turn_completed.params.expect("turn/completed params"))?; + assert_eq!(completed.turn.status, TurnStatus::Completed); + let requests = responses_mock.requests(); assert_eq!(requests.len(), if model_path == "/v1" { 2 } else { 1 }); assert_eq!( diff --git a/codex-rs/app-server/tests/suite/v2/account_thread_usage.rs b/codex-rs/app-server/tests/suite/v2/account_thread_usage.rs index 8fd846e4d4..17e507ccf7 100644 --- a/codex-rs/app-server/tests/suite/v2/account_thread_usage.rs +++ b/codex-rs/app-server/tests/suite/v2/account_thread_usage.rs @@ -135,6 +135,11 @@ async fn account_thread_usage_supports_externally_managed_authentication() -> Re .plan_type("business") .chatgpt_account_id("external-workspace"), )?; + Mock::given(method("GET")) + .and(path("/api/codex/config/bundle")) + .respond_with(ResponseTemplate::new(/*s*/ 200).set_body_json(json!({}))) + .mount(&server) + .await; let mut app_server = TestAppServer::builder() .with_codex_home(codex_home.path()) .with_env_overrides(&[("OPENAI_API_KEY", None)]) diff --git a/codex-rs/app-server/tests/suite/v2/application_network.rs b/codex-rs/app-server/tests/suite/v2/application_network.rs new file mode 100644 index 0000000000..71460ab568 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/application_network.rs @@ -0,0 +1,109 @@ +//! Application policy is enforced through public thread and turn RPCs. + +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use app_test_support::write_chatgpt_auth; +use app_test_support::write_models_cache; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::TurnCompletedNotification; +use codex_app_server_protocol::TurnStatus; +use codex_config::types::AuthCredentialsStoreMode; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; +use pretty_assertions::assert_eq; +use std::time::Duration; +use tempfile::TempDir; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test] +async fn standalone_startup_applies_local_policy_before_fetching_stored_account_requirements() +-> Result<()> { + let backend = MockServer::start().await; + let bundle_path = "/backend-api/wham/config/bundle"; + Mock::given(method("GET")) + .and(path(bundle_path)) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(0) + .mount(&backend) + .await; + let home = TempDir::new()?; + std::fs::write( + home.path().join("config.toml"), + format!( + "chatgpt_base_url = '{}/backend-api/'\ncli_auth_credentials_store = 'file'\n", + backend.uri() + ), + )?; + std::fs::write( + home.path().join("requirements.toml"), + "[application.network]", + )?; + write_chatgpt_auth( + home.path(), + ChatGptAuthFixture::new("stored-token") + .account_id("selected") + .plan_type("enterprise"), + AuthCredentialsStoreMode::File, + )?; + let _app = TestAppServer::builder() + .with_codex_home(home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build_initialized() + .await?; + assert!( + backend + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.url.path() != bundle_path) + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn explicit_reloads_apply_local_edits_and_cancel_active_responses() -> Result<()> { + for requirements in ["[application.network]", "[malformed"] { + let (_release, gate) = tokio::sync::oneshot::channel(); + let (destination, _) = start_streaming_sse_server(vec![vec![StreamingSseChunk { + gate: Some(gate), + body: String::new(), + }]]) + .await; + let home = TempDir::new()?; + MockResponsesConfig::new(destination.uri()).write(home.path())?; + write_models_cache(home.path()).await?; + let mut app = TestAppServer::builder() + .with_codex_home(home.path()) + .build_initialized() + .await?; + let thread = app.start_thread(ThreadStartParams::default()).await?.thread; + let params = serde_json::json!({ + "threadId": thread.id, "input": [{"type": "text", "text": "Hello"}] + }); + app.send_request("turn/start", Some(params)).await?; + timeout( + Duration::from_secs(/*secs*/ 30), + destination.wait_for_request_count(/*count*/ 1), + ) + .await?; + std::fs::write(home.path().join("requirements.toml"), requirements)?; + app.send_request("config/read", Some(serde_json::json!({}))) + .await?; + let completed: TurnCompletedNotification = timeout( + Duration::from_secs(/*secs*/ 5), + app.read_notification("turn/completed"), + ) + .await??; + assert_eq!(completed.turn.status, TurnStatus::Failed); + destination.shutdown().await; + } + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 09856f5306..51f3e06b9a 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -7,6 +7,7 @@ mod analytics; mod app_installed; mod app_list; mod app_read; +mod application_network; mod attestation; mod auto_env; mod bedrock_setup; diff --git a/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs b/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs index b8356c5641..635ebfad03 100644 --- a/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs +++ b/codex-rs/app-server/tests/suite/v2/rate_limits_identity_tests.rs @@ -12,6 +12,7 @@ use axum::http::HeaderMap; use axum::routing::get; use codex_app_server_protocol::GetAccountRateLimitsResponse; use codex_app_server_protocol::LoginAccountResponse; +use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; use pretty_assertions::assert_eq; use serde_json::json; @@ -34,10 +35,15 @@ async fn identity_is_rechecked_after_backend_response( user: &str, verified: bool, ) -> Result<()> { + let old_token = encode_id_token( + &ChatGptIdTokenClaims::new() + .chatgpt_account_id("workspace-a") + .chatgpt_user_id("user-a"), + )?; let home = TempDir::new()?; write_chatgpt_auth( home.path(), - ChatGptAuthFixture::new("old-token") + ChatGptAuthFixture::new(&old_token) .account_id("workspace-a") .chatgpt_user_id("user-a") .plan_type("team"), @@ -52,13 +58,15 @@ async fn identity_is_rechecked_after_backend_response( home.path().join("config.toml"), format!("chatgpt_base_url = \"http://{}\"\n", listener.local_addr()?), )?; + let expected_authorization = format!("Bearer {old_token}"); let router = Router::new().route( "/api/codex/usage", get(move |headers: HeaderMap| { + let expected_authorization = expected_authorization.clone(); let entered = Arc::clone(&request_entered); let release = Arc::clone(&response_release); async move { - assert_eq!(headers["authorization"], "Bearer old-token"); + assert_eq!(headers["authorization"], expected_authorization); assert_eq!(headers["chatgpt-account-id"], "workspace-a"); entered.notify_one(); release.notified().await; @@ -77,6 +85,14 @@ async fn identity_is_rechecked_after_backend_response( .with_env_overrides(&[("OPENAI_API_KEY", None)]) .build_initialized_with_timeout(READ_TIMEOUT) .await?; + let initial_login = app + .send_chatgpt_auth_tokens_login_request( + old_token, + "workspace-a".into(), + Some("team".into()), + ) + .await?; + let _: LoginAccountResponse = app.read_response(initial_login).await?; let request = app .send_request( "account/rateLimits/read", @@ -96,6 +112,19 @@ async fn identity_is_rechecked_after_backend_response( let response: LoginAccountResponse = timeout(READ_TIMEOUT, app.read_response(login)).await??; assert_eq!(response, LoginAccountResponse::ChatgptAuthTokens {}); release.notify_one(); + if !verified { + let error = timeout( + READ_TIMEOUT, + app.read_stream_until_error_message(RequestId::Integer(request)), + ) + .await??; + server.abort(); + assert_eq!( + error.error.message, + "failed to fetch codex rate limits: application network permission was revoked" + ); + return Ok(()); + } let response: GetAccountRateLimitsResponse = timeout(READ_TIMEOUT, app.read_response(request)).await??; server.abort(); diff --git a/codex-rs/app-server/tests/suite/v2/remote_control_auth.rs b/codex-rs/app-server/tests/suite/v2/remote_control_auth.rs index f557b81211..3ca4250ff8 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_control_auth.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_control_auth.rs @@ -312,6 +312,7 @@ async fn client_revoke_retry_stays_with_original_owner() -> Result<()> { app.read_stream_until_error_message(RequestId::Integer(old)), ) .await??; + login(&mut app, refreshed_user, "account-a", "refreshed").await?; let id = app .send_remote_control_clients_revoke_request(RemoteControlClientsRevokeParams { environment_id: "environment-id".to_string(), diff --git a/codex-rs/app-server/tests/suite/v2/workspace_routing.rs b/codex-rs/app-server/tests/suite/v2/workspace_routing.rs index c87f31c711..67f2041f49 100644 --- a/codex-rs/app-server/tests/suite/v2/workspace_routing.rs +++ b/codex-rs/app-server/tests/suite/v2/workspace_routing.rs @@ -38,13 +38,14 @@ async fn config(home: &TempDir, backend: &MockServer) -> Result<()> { backend.uri() ), )?; + // Routing uses a plain HTTP fixture; restricted policies require HTTPS/WSS. Mock::given(method("GET")).and(path("/backend-api/wham/config/bundle")) .respond_with(|request: &wiremock::Request| { let account_id = request.headers.get("chatgpt-account-id").expect("workspace header") .to_str().expect("workspace header text"); ResponseTemplate::new(200).set_body_json(json!({"requirements_toml": {"enterprise_managed": [{ "id": "workspace-policy", "name": "Workspace policy", - "contents": format!("[application.network.domains]\n'{account_id}.example' = 'allow'"), + "contents": format!("[application.network]\nenabled = false\n[application.network.domains]\n'{account_id}.example' = 'allow'"), }]}})) }).mount(backend).await; Ok(()) @@ -116,7 +117,8 @@ async fn saved_workspace_is_discovered_once_and_not_the_default_account( .set_delay(Duration::from_secs(bundle_delay_secs)), ) .with_priority(1) - .expect(2) + // Startup and account hydration each retry this delayed response via proxy fallback. + .expect(4) .mount(&backend) .await; } @@ -277,7 +279,7 @@ async fn login_and_workspace_switch_notify_after_routing_is_ready_then_logout_cl assert_eq!( requirements["requirements"]["application"], json!({"network": { - "enabled": true, "domains": {format!("{account}.example"): "allow"} + "enabled": false, "domains": {format!("{account}.example"): "allow"} }}) ); } diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 3c9116d338..48da6b0b66 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -53,6 +53,7 @@ pub use thread_usage::ThreadUsageBreakdownGroup; #[derive(Debug)] pub enum RequestError { + Policy(codex_http_client::NetworkPolicyDenied), UnexpectedStatus { method: String, url: String, @@ -67,7 +68,7 @@ impl RequestError { pub fn status(&self) -> Option { match self { Self::UnexpectedStatus { status, .. } => Some(*status), - Self::Other(_) => None, + Self::Policy(_) | Self::Other(_) => None, } } @@ -79,6 +80,7 @@ impl RequestError { impl fmt::Display for RequestError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { + Self::Policy(denied) => denied.fmt(f), Self::UnexpectedStatus { method, url, @@ -97,6 +99,7 @@ impl fmt::Display for RequestError { impl std::error::Error for RequestError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { + Self::Policy(denied) => Some(denied), Self::UnexpectedStatus { .. } => None, Self::Other(err) => Some(err.as_ref()), } @@ -109,6 +112,20 @@ impl From for RequestError { } } +impl From for RequestError { + fn from(error: codex_http_client::HttpError) -> Self { + match error { + codex_http_client::HttpError::Policy(denied) => Self::Policy(denied), + error @ (codex_http_client::HttpError::Request(_) + | codex_http_client::HttpError::Route(_) + | codex_http_client::HttpError::Build(_) + | codex_http_client::HttpError::UnsupportedRedirectScheme(_) + | codex_http_client::HttpError::TooManyRedirects + | codex_http_client::HttpError::Timeout) => Self::Other(error.into()), + } + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] pub enum AddCreditsNudgeCreditType { @@ -302,7 +319,7 @@ impl Client { method: &str, url: &str, ) -> std::result::Result<(String, String), RequestError> { - let res = req.send().await.map_err(anyhow::Error::from)?; + let res = req.send().await?; let status = res.status(); let content_type = res .headers() @@ -310,7 +327,7 @@ impl Client { .and_then(|v| v.to_str().ok()) .unwrap_or("") .to_string(); - let body = res.text().await.map_err(anyhow::Error::from)?; + let body = res.text().await?; if !status.is_success() { return Err(RequestError::UnexpectedStatus { method: method.to_string(), @@ -511,15 +528,20 @@ impl Client { pub async fn get_config_bundle( &self, ) -> std::result::Result { - let url = match self.path_style { - PathStyle::CodexApi => format!("{}/api/codex/config/bundle", self.base_url), - PathStyle::ChatGptApi => format!("{}/wham/config/bundle", self.base_url), - }; + let url = self.config_bundle_url(); let (body, ct) = self.exec_bootstrap_get(&url).await?; self.decode_json::(&url, &ct, &body) .map_err(RequestError::from) } + /// Returns the normalized discovery endpoint used by this client. + pub fn config_bundle_url(&self) -> String { + match self.path_style { + PathStyle::CodexApi => format!("{}/api/codex/config/bundle", self.base_url), + PathStyle::ChatGptApi => format!("{}/wham/config/bundle", self.base_url), + } + } + /// Fetch authenticated Codex user settings from the active backend route. /// /// Uses `GET /api/codex/settings/user` for Codex API hosts and diff --git a/codex-rs/backend-client/src/client_request_tests.rs b/codex-rs/backend-client/src/client_request_tests.rs index 4dfe774fec..0d699587e7 100644 --- a/codex-rs/backend-client/src/client_request_tests.rs +++ b/codex-rs/backend-client/src/client_request_tests.rs @@ -19,17 +19,20 @@ use wiremock::matchers::path; use super::*; -#[test] -fn client_preserves_supplied_http_client_factory_policy() { +#[tokio::test] +async fn client_preserves_supplied_factory_and_typed_policy_denial() { + let controller = codex_http_client::NetworkPolicyController::default(); let client = Client::new( "https://example.test", - HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy) + .with_network_policy(controller.policy()), ); - assert_eq!( client.http.outbound_proxy_policy(), OutboundProxyPolicy::RespectSystemProxy ); + let denied = client.get_config_bundle().await; + assert!(matches!(denied, Err(RequestError::Policy(_)))); } #[test] diff --git a/codex-rs/cloud-config/src/backend.rs b/codex-rs/cloud-config/src/backend.rs index b8b456a8ac..fc56579940 100644 --- a/codex-rs/cloud-config/src/backend.rs +++ b/codex-rs/cloud-config/src/backend.rs @@ -25,6 +25,7 @@ impl RetryableFailureKind { #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) enum BundleRequestError { + Policy(codex_http_client::NetworkPolicyDenied), Retryable(RetryableFailureKind), Unauthorized { status_code: Option, @@ -72,6 +73,9 @@ impl BundleClient for BackendBundleClient { tracing::warn!(error = %err, "Failed to fetch cloud config bundle"); }) .map_err(|err| { + if let codex_backend_client::RequestError::Policy(denied) = &err { + return BundleRequestError::Policy(*denied); + } let status_code = err.status().map(|status| status.as_u16()); if err.is_unauthorized() { BundleRequestError::Unauthorized { diff --git a/codex-rs/cloud-config/src/bundle_loader.rs b/codex-rs/cloud-config/src/bundle_loader.rs index d38355189f..eb99944010 100644 --- a/codex-rs/cloud-config/src/bundle_loader.rs +++ b/codex-rs/cloud-config/src/bundle_loader.rs @@ -8,26 +8,9 @@ use codex_login::AuthConfig; use codex_login::AuthManager; use std::path::PathBuf; use std::sync::Arc; -use std::sync::Mutex; -use std::sync::OnceLock; use tokio::task::AbortHandle; use tokio::task::JoinHandle; -fn refresher_task_slot() -> &'static Mutex> { - static REFRESHER_TASK: OnceLock>> = OnceLock::new(); - REFRESHER_TASK.get_or_init(|| Mutex::new(None)) -} - -pub(crate) fn replace_refresh_task(slot: &Mutex>, next: AbortHandle) { - let mut guard = slot.lock().unwrap_or_else(|err| { - tracing::warn!("cloud config bundle refresher task slot was poisoned"); - err.into_inner() - }); - if let Some(previous) = guard.replace(next) { - previous.abort(); - } -} - struct CloudConfigBundleLoaderLifetime { service: Arc>, refresh_task: JoinHandle<()>, @@ -54,8 +37,7 @@ pub fn cloud_config_bundle_loader( codex_home, CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let (loader, refresh_task) = cloud_config_bundle_loader_for_service(service); - replace_refresh_task(refresher_task_slot(), refresh_task); + let (loader, _) = cloud_config_bundle_loader_for_service(service); loader } @@ -90,8 +72,7 @@ pub async fn cloud_config_bundle_loader_for_storage( ) -> std::io::Result { let service = cloud_config_bundle_service_for_storage(auth_config, enable_codex_api_key_env).await?; - let (loader, refresh_task) = cloud_config_bundle_loader_for_service(service); - replace_refresh_task(refresher_task_slot(), refresh_task); + let (loader, _) = cloud_config_bundle_loader_for_service(service); Ok(loader) } diff --git a/codex-rs/cloud-config/src/service.rs b/codex-rs/cloud-config/src/service.rs index 6df148010a..e27db63b0c 100644 --- a/codex-rs/cloud-config/src/service.rs +++ b/codex-rs/cloud-config/src/service.rs @@ -254,6 +254,13 @@ where .validate_and_cache_remote_bundle(&auth, trigger, attempt, bundle) .await; } + Err(BundleRequestError::Policy(denied)) => { + return Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::RequestFailed, + /*status_code*/ None, + denied.to_string(), + )); + } Err(BundleRequestError::Retryable(status)) => { last_status_code = status.status_code(); if self @@ -416,6 +423,13 @@ where *auth = refreshed_auth; return Ok(UnauthorizedRecoveryAction::RetrySameAttempt); } + Err(RefreshTokenError::Policy(error)) => { + return Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::Auth, + status_code, + error.to_string(), + )); + } Err(RefreshTokenError::Permanent(failed)) => { tracing::warn!( error = %failed, @@ -491,6 +505,12 @@ where "Timed out refreshing cloud config bundle cache from remote; keeping existing cache" ); emit_load_metric("refresh", "error", /*bundle*/ None); + self.publish_refresh_result(Err(CloudConfigBundleLoadError::new( + CloudConfigBundleLoadErrorCode::Timeout, + /*status_code*/ None, + "timed out refreshing cloud config bundle", + ))) + .await; } } } @@ -510,9 +530,7 @@ where { Ok(bundle) => { emit_load_metric("refresh", "success", bundle.as_ref()); - if let Some(latest_bundle) = self.latest_bundle.get() { - *latest_bundle.lock().await = Ok(bundle); - } + self.publish_refresh_result(Ok(bundle)).await; } Err(err) => { tracing::error!( @@ -521,16 +539,24 @@ where "Failed to refresh cloud config bundle cache from remote" ); emit_load_metric("refresh", "error", /*bundle*/ None); - if let Some(latest_bundle) = self.latest_bundle.get() { - let mut latest_bundle = latest_bundle.lock().await; - if latest_bundle.is_err() { - *latest_bundle = Err(err); - } - } + self.publish_refresh_result(Err(err)).await; } } true } + + async fn publish_refresh_result( + &self, + result: Result, CloudConfigBundleLoadError>, + ) { + let Some(latest) = self.latest_bundle.get() else { + return; + }; + let mut latest = latest.lock().await; + if result.is_ok() || latest.is_err() { + *latest = result; + } + } } #[cfg(test)] diff --git a/codex-rs/cloud-config/src/service_tests.rs b/codex-rs/cloud-config/src/service_tests.rs index 08d2d19314..55b68c49ed 100644 --- a/codex-rs/cloud-config/src/service_tests.rs +++ b/codex-rs/cloud-config/src/service_tests.rs @@ -1403,7 +1403,7 @@ async fn production_loader_refreshes_later_configs_and_preserves_failed_refreshe } #[tokio::test(start_paused = true)] -async fn refresh_stops_on_replacement_or_after_the_last_loader_clone() { +async fn each_loader_owns_refresh_until_its_last_clone_is_dropped() { let codex_home = tempdir().expect("tempdir"); let fetcher = Arc::new(StaticBundleClient::new(test_bundle())); let service = CloudConfigBundleService::new( @@ -1412,10 +1412,8 @@ async fn refresh_stops_on_replacement_or_after_the_last_loader_clone() { codex_home.path().to_path_buf(), CLOUD_CONFIG_BUNDLE_TIMEOUT, ); - let (loader, abort_handle) = + let (loader, original_task) = crate::bundle_loader::cloud_config_bundle_loader_for_service(service); - let task_slot = std::sync::Mutex::new(None); - crate::bundle_loader::replace_refresh_task(&task_slot, abort_handle); let cloned_loader = loader.clone(); assert_eq!(loader.get().await, Ok(Some(test_bundle()))); tokio::task::yield_now().await; @@ -1446,25 +1444,27 @@ async fn refresh_stops_on_replacement_or_after_the_last_loader_clone() { ); let (replacement_loader, replacement_handle) = crate::bundle_loader::cloud_config_bundle_loader_for_service(replacement_service); - let replacement_task = replacement_handle.clone(); - crate::bundle_loader::replace_refresh_task(&task_slot, replacement_handle); + let replacement_task = replacement_handle; assert_eq!(replacement_loader.get().await, Ok(Some(test_bundle()))); tokio::task::yield_now().await; tokio::time::advance(CLOUD_CONFIG_BUNDLE_CACHE_REFRESH_INTERVAL + Duration::from_millis(1)) .await; let refresh_deadline = std::time::Instant::now() + Duration::from_secs(5); - while replacement_fetcher.request_count.load(Ordering::SeqCst) < 2 { + while replacement_fetcher.request_count.load(Ordering::SeqCst) < 2 + || fetcher.request_count.load(Ordering::SeqCst) < 3 + { assert!( std::time::Instant::now() < refresh_deadline, - "the replacement refresher should stay active" + "both loader owners should keep refreshing" ); tokio::task::yield_now().await; } - assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 2); + assert_eq!(fetcher.request_count.load(Ordering::SeqCst), 3); drop(cloned_loader); tokio::task::yield_now().await; + assert!(original_task.is_finished()); assert!(!replacement_task.is_finished()); drop(replacement_loader); diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 194d05dff6..e494e92d41 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -38,6 +38,7 @@ uuid = { workspace = true, features = ["v4"] } [dev-dependencies] codex-code-mode = { workspace = true } +codex-http-client = { workspace = true } codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } diff --git a/codex-rs/code-mode-host/tests/grpc.rs b/codex-rs/code-mode-host/tests/grpc.rs index 0e204ff989..2ffe4a8524 100644 --- a/codex-rs/code-mode-host/tests/grpc.rs +++ b/codex-rs/code-mode-host/tests/grpc.rs @@ -46,6 +46,8 @@ use tonic::transport::Server; mod host; #[path = "support/large_tool_delegate.rs"] mod large_tool_delegate; +#[path = "grpc/network_policy_tests.rs"] +mod network_policy_tests; #[path = "support/recording_delegate.rs"] mod recording_delegate; diff --git a/codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs b/codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs new file mode 100644 index 0000000000..5200c358f8 --- /dev/null +++ b/codex-rs/code-mode-host/tests/grpc/network_policy_tests.rs @@ -0,0 +1,96 @@ +//! Destination admission and account revocation cover the gRPC request lifetime. + +use codex_http_client::DestinationPolicy; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyController; +use codex_http_client::OutboundProxyPolicy; +use pretty_assertions::assert_eq; + +use super::*; + +#[tokio::test] +async fn restricted_destination_is_rejected_before_connecting_to_the_code_mode_host() -> Result<()> +{ + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let controller = NetworkPolicyController::default(); + assert!(controller.publish( + controller.policy().revision(), + DestinationPolicy::Restricted { + allowed_hosts: ["allowed.example".into()].into() + } + )); + let provider = Arc::new(GrpcCodeModeSessionProvider::with_http_client_factory( + format!("https://{}", listener.local_addr()?), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(controller.policy()), + )); + let Err(error) = provider.create_session().await else { + anyhow::bail!("restricted endpoint created a session") + }; + assert!( + error.contains("destination denied by application network policy"), + "{error}" + ); + assert!( + timeout(Duration::from_millis(/*millis*/ 50), listener.accept()) + .await + .is_err() + ); + + Ok(()) +} + +#[tokio::test] +async fn account_change_revokes_active_rpc_and_renews_the_cached_client() -> Result<()> { + let host = HostHarness::start("grpc://127.0.0.1:0").await?; + let controller = NetworkPolicyController::default(); + assert!(controller.publish( + controller.policy().revision(), + DestinationPolicy::Unrestricted + )); + let provider = GrpcCodeModeSessionProvider::with_http_client_factory( + host.endpoint.clone(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(controller.policy()), + ); + let session = provider + .create_session() + .await + .map_err(anyhow::Error::msg)?; + let mut pending = request("await new Promise(() => {});"); + pending.yield_time_ms = Some(60_000); + let started = session + .execute(pending, Arc::new(NoopCodeModeSessionDelegate)) + .await + .map_err(anyhow::Error::msg)?; + + controller.policy().invalidate(); + assert!( + timeout(TEST_TIMEOUT, started.initial_response()) + .await + .context("account revocation did not interrupt the active RPC")? + .is_err() + ); + assert!(controller.publish( + controller.policy().revision(), + DestinationPolicy::Unrestricted + )); + + let replacement = provider + .create_session() + .await + .map_err(anyhow::Error::msg)?; + let response = execute( + &replacement, + request(r#"text("replacement")"#), + Arc::new(NoopCodeModeSessionDelegate), + ) + .await?; + assert_eq!( + response, + text_response("1", "replacement", response.code_mode_host_duration()) + ); + session.shutdown().await.map_err(anyhow::Error::msg)?; + replacement.shutdown().await.map_err(anyhow::Error::msg)?; + Ok(()) +} diff --git a/codex-rs/code-mode/src/grpc_session/transport.rs b/codex-rs/code-mode/src/grpc_session/transport.rs index bd9d4b835e..9b32d005b0 100644 --- a/codex-rs/code-mode/src/grpc_session/transport.rs +++ b/codex-rs/code-mode/src/grpc_session/transport.rs @@ -19,119 +19,106 @@ use super::GrpcClient; pub(super) type GrpcTransport = BoxCloneSyncService, Response, io::Error>; -pub(super) struct SharedTransport { - endpoint: TransportEndpoint, - client: tokio::sync::OnceCell, -} - -enum TransportEndpoint { +pub(super) enum SharedTransport { Url { endpoint: String, http_client_factory: HttpClientFactory, + client: tokio::sync::Mutex>, }, - Connected(Channel), + Connected(GrpcClient), } impl SharedTransport { pub(super) fn new(endpoint: String, http_client_factory: HttpClientFactory) -> Self { - Self { - endpoint: TransportEndpoint::Url { - endpoint, - http_client_factory, - }, - client: tokio::sync::OnceCell::new(), + Self::Url { + endpoint, + http_client_factory, + client: tokio::sync::Mutex::new(None), } } pub(super) fn with_channel(channel: Channel) -> Self { - Self { - endpoint: TransportEndpoint::Connected(channel), - client: tokio::sync::OnceCell::new(), - } + let transport = channel.map_err(io::Error::other); + Self::Connected( + CodeModeHostClient::new(BoxCloneSyncService::new(transport)) + .max_decoding_message_size(MAX_FRAME_BYTES) + .max_encoding_message_size(MAX_FRAME_BYTES), + ) } pub(super) async fn client(&self) -> Result { - self.client - .get_or_try_init(|| async { - let client = match &self.endpoint { - TransportEndpoint::Url { endpoint, .. } if endpoint.starts_with("unix:") => { - let channel = Endpoint::from_shared(endpoint.clone()) - .map_err(|error| { - format!("invalid gRPC code-mode Unix socket endpoint: {error}") - })? - .connect_lazy(); - let transport = channel.map_err(io::Error::other); - CodeModeHostClient::new(BoxCloneSyncService::new(transport)) - } - TransportEndpoint::Url { - endpoint, - http_client_factory, - } => { - let target = reqwest::Url::parse(endpoint) - .map_err(|error| format!("invalid gRPC code-mode host URL: {error}"))?; - if !matches!(target.scheme(), "http" | "https") { - return Err("gRPC code-mode host URL must use http or https".to_string()); - } - if !target.username().is_empty() || target.password().is_some() { - return Err( - "gRPC code-mode host URL must not include credentials".to_string(), - ); - } - if target.path() != "/" - || target.query().is_some() - || target.fragment().is_some() - { - return Err("gRPC code-mode host URL must not include a path, query, or fragment".to_string()); - } - let origin: Uri = endpoint - .parse() - .map_err(|error| format!("invalid gRPC code-mode host origin: {error}"))?; - let endpoint = endpoint.clone(); - let http_client_factory = http_client_factory.clone(); - let client = tokio::task::spawn_blocking(move || { - http_client_factory - .build_reqwest_client( - reqwest::Client::builder() - .http2_prior_knowledge() - .redirect(reqwest::redirect::Policy::none()), - &endpoint, - ClientRouteClass::Other, - ) - .map_err(|error| { - format!( - "failed to configure gRPC code-mode host transport: {error}" - ) - }) - }) + let (endpoint, http_client_factory, client) = match self { + Self::Url { + endpoint, + http_client_factory, + client, + } => (endpoint, http_client_factory, client), + Self::Connected(client) => return Ok(client.clone()), + }; + let bound_factory = http_client_factory.clone().with_network_policy( + http_client_factory + .network_policy() + .clone() + .for_current_account(), + ); + let mut cached = client.lock().await; + if let Some((factory, client)) = cached.as_ref() + && factory == &bound_factory + { + return Ok(client.clone()); + } + let client = if endpoint.starts_with("unix:") { + let channel = Endpoint::from_shared(endpoint.clone()) + .map_err(|error| format!("invalid gRPC code-mode Unix socket endpoint: {error}"))? + .connect_lazy(); + let transport = channel.map_err(io::Error::other); + CodeModeHostClient::new(BoxCloneSyncService::new(transport)) + } else { + let target = reqwest::Url::parse(endpoint) + .map_err(|error| format!("invalid gRPC code-mode host URL: {error}"))?; + if !matches!(target.scheme(), "http" | "https") { + return Err("gRPC code-mode host URL must use http or https".to_string()); + } + if !target.username().is_empty() || target.password().is_some() { + return Err("gRPC code-mode host URL must not include credentials".to_string()); + } + if target.path() != "/" || target.query().is_some() || target.fragment().is_some() { + return Err( + "gRPC code-mode host URL must not include a path, query, or fragment" + .to_string(), + ); + } + let origin: Uri = endpoint + .parse() + .map_err(|error| format!("invalid gRPC code-mode host origin: {error}"))?; + let client = codex_http_client::RouteAwareClientPool::with_builder( + bound_factory.clone(), + ClientRouteClass::Other, + codex_http_client::HttpClientBuilder::new() + .without_request_logging() + .http2_prior_knowledge() + .without_redirects(), + ); + let transport = service_fn(move |request: Request| { + let client = client.clone(); + async move { + let request = + request.map(|body| reqwest::Body::wrap_stream(body.into_data_stream())); + let request = reqwest::Request::try_from(request).map_err(io::Error::other)?; + let response = client + .execute(request) .await - .map_err(|error| { - format!("gRPC code-mode host transport task failed: {error}") - })??; - let transport = service_fn(move |request: Request| { - let client = client.clone(); - async move { - let request = request.map(|body| { - reqwest::Body::wrap_stream(body.into_data_stream()) - }); - let request = - reqwest::Request::try_from(request).map_err(io::Error::other)?; - let response: Response = - client.execute(request).await.map_err(io::Error::other)?.into(); - Ok::<_, io::Error>(response.map(Body::new)) - } - }); - CodeModeHostClient::with_origin(BoxCloneSyncService::new(transport), origin) - } - TransportEndpoint::Connected(channel) => { - let transport = channel.clone().map_err(io::Error::other); - CodeModeHostClient::new(BoxCloneSyncService::new(transport)) - } - }; - Ok(client - .max_decoding_message_size(MAX_FRAME_BYTES) - .max_encoding_message_size(MAX_FRAME_BYTES)) - }) - .await - .cloned() + .map_err(io::Error::other)? + .into_http_response(); + Ok::<_, io::Error>(response.map(Body::new)) + } + }); + CodeModeHostClient::with_origin(BoxCloneSyncService::new(transport), origin) + }; + let client = client + .max_decoding_message_size(MAX_FRAME_BYTES) + .max_encoding_message_size(MAX_FRAME_BYTES); + *cached = Some((bound_factory, client.clone())); + Ok(client) } } diff --git a/codex-rs/config/src/cloud_config_bundle.rs b/codex-rs/config/src/cloud_config_bundle.rs index 847b67725e..e8e7950a1c 100644 --- a/codex-rs/config/src/cloud_config_bundle.rs +++ b/codex-rs/config/src/cloud_config_bundle.rs @@ -55,10 +55,11 @@ pub struct CloudRequirementsTomlBundle { } impl CloudRequirementsTomlBundle { - pub(crate) fn into_layers(self, base_dir: &AbsolutePathBuf) -> Vec { - let Self { enterprise_managed } = self; - let mut layers = enterprise_managed + pub(crate) fn into_layers(self) -> Vec { + // Bundle fragments arrive highest-priority first; requirements merge in reverse order. + self.enterprise_managed .into_iter() + .rev() .map(|fragment| { RequirementsLayerEntry::from_toml( RequirementSource::EnterpriseManaged { @@ -67,13 +68,8 @@ impl CloudRequirementsTomlBundle { }, fragment.contents, ) - .with_base_dir(base_dir.clone()) }) - .collect::>(); - // Bundle fragments arrive highest-priority first, while requirements - // layers are merged lowest-priority to highest-priority. - layers.reverse(); - layers + .collect() } } @@ -132,7 +128,11 @@ impl CloudConfigBundleLayers { cloud_config_layers_from_fragments(config_enterprise_managed, base_dir)? }; - let enterprise_managed_requirements = requirements_toml.into_layers(base_dir); + let enterprise_managed_requirements = requirements_toml + .into_layers() + .into_iter() + .map(|layer| layer.with_base_dir(base_dir.clone())) + .collect::>(); Ok(Self { enterprise_managed_config, diff --git a/codex-rs/config/src/loader/application.rs b/codex-rs/config/src/loader/application.rs new file mode 100644 index 0000000000..ab0786c35f --- /dev/null +++ b/codex-rs/config/src/loader/application.rs @@ -0,0 +1,102 @@ +//! Reads application requirements independently of user, project, and cloud configuration. +//! Captured local layers retain their positions around cloud layers during composition. + +use crate::ApplicationRequirementsToml; +use crate::CloudRequirementsTomlBundle; +use crate::LoaderOverrides; +use crate::RequirementsLayerEntry; +use crate::compose_requirements; +use codex_file_system::ExecutorFileSystem; +use std::io; +use toml::Value; + +/// Local application sources captured without waiting for a cloud request. +#[derive(Clone, Debug, Default)] +pub struct LocalApplicationRequirements { + system: Option, + mdm: Option, + ignore_managed: bool, +} + +impl LocalApplicationRequirements { + /// Composes system, cloud, and MDM application policy in the normal requirements order. + pub fn compose( + &self, + cloud: CloudRequirementsTomlBundle, + ) -> io::Result> { + if self.ignore_managed { + return Ok(None); + } + let mut layers = self.system.clone().into_iter().collect::>(); + for layer in cloud.into_layers() { + layers.extend(application_layer(layer)?); + } + layers.extend(self.mdm.clone()); + Ok(compose_requirements(layers)? + .and_then(|requirements| requirements.into_toml().application)) + } +} + +/// Reads only the local sources that can set application requirements. +/// macOS policy reads require a successful preference refresh before using cached values. +pub async fn load_local_application_requirements( + fs: &dyn ExecutorFileSystem, + overrides: &LoaderOverrides, +) -> io::Result { + if overrides.ignore_managed_requirements { + return Ok(LocalApplicationRequirements { + ignore_managed: true, + ..Default::default() + }); + } + let system_path = super::system_requirements_toml_file_with_overrides(overrides)?; + let system = super::load_requirements_toml(fs, &system_path) + .await? + .map(application_layer) + .transpose()? + .flatten(); + #[cfg(target_os = "macos")] + let mdm = { + if overrides.macos_managed_config_requirements_base64.is_none() { + tokio::task::spawn_blocking(super::macos::synchronize_managed_preferences) + .await + .map_err(io::Error::other)??; + } + super::macos::load_managed_admin_requirements_layer( + overrides + .macos_managed_config_requirements_base64 + .as_deref(), + ) + .await? + .map(application_layer) + .transpose()? + .flatten() + }; + #[cfg(not(target_os = "macos"))] + let mdm = None; + let local = LocalApplicationRequirements { + system, + mdm, + ignore_managed: false, + }; + // Validate local requirements before cloud loading. + local.compose(CloudRequirementsTomlBundle::default())?; + Ok(local) +} + +fn application_layer(layer: RequirementsLayerEntry) -> io::Result> { + let (source, mut value, _base_dir) = layer.into_raw_parts()?; + let Some(application) = value + .as_table_mut() + .and_then(|table| table.remove("application")) + else { + return Ok(None); + }; + Ok(Some(RequirementsLayerEntry::from_toml_value( + source, + Value::Table(toml::map::Map::from_iter([( + "application".into(), + application, + )])), + ))) +} diff --git a/codex-rs/config/src/loader/macos.rs b/codex-rs/config/src/loader/macos.rs index 1dd279951c..592bb2570b 100644 --- a/codex-rs/config/src/loader/macos.rs +++ b/codex-rs/config/src/loader/macos.rs @@ -133,6 +133,20 @@ fn load_managed_admin_requirements() -> io::Result> { .transpose() } +pub(super) fn synchronize_managed_preferences() -> io::Result<()> { + #[link(name = "CoreFoundation", kind = "framework")] + unsafe extern "C" { + fn CFPreferencesAppSynchronize(application_id: CFStringRef) -> u8; + } + let application_id = CFString::new(MANAGED_PREFERENCES_APPLICATION_ID); + if unsafe { CFPreferencesAppSynchronize(application_id.as_concrete_TypeRef()) } == 0 { + return Err(io::Error::other( + "Failed to synchronize managed preferences", + )); + } + Ok(()) +} + fn load_managed_preference(key_name: &str) -> io::Result> { let key = CFString::new(key_name); let application = CFString::new(MANAGED_PREFERENCES_APPLICATION_ID); diff --git a/codex-rs/config/src/loader/managed_requirements.rs b/codex-rs/config/src/loader/managed_requirements.rs index c3efe1abee..f2b55fc315 100644 --- a/codex-rs/config/src/loader/managed_requirements.rs +++ b/codex-rs/config/src/loader/managed_requirements.rs @@ -40,7 +40,12 @@ pub async fn load_managed_requirements_state( let mut bundle_requirements = Vec::new(); if let Some(bundle) = cloud_config_bundle.get().await.map_err(io::Error::other)? { let base_dir = AbsolutePathBuf::from_absolute_path(codex_home)?; - bundle_requirements = bundle.requirements_toml.into_layers(&base_dir); + bundle_requirements = bundle + .requirements_toml + .into_layers() + .into_iter() + .map(|layer| layer.with_base_dir(base_dir.clone())) + .collect(); } let (requirements, _, _) = load_requirements_from_sources( fs, diff --git a/codex-rs/config/src/loader/mod.rs b/codex-rs/config/src/loader/mod.rs index 8dbaad22f2..615d394a89 100644 --- a/codex-rs/config/src/loader/mod.rs +++ b/codex-rs/config/src/loader/mod.rs @@ -1,3 +1,4 @@ +mod application; mod layer_io; mod local; #[cfg(target_os = "macos")] @@ -13,6 +14,9 @@ mod projectless_directory_tests; #[cfg(windows)] mod windows; +pub use application::LocalApplicationRequirements; +pub use application::load_local_application_requirements; + use self::layer_io::LoadedConfigLayers; use crate::CONFIG_TOML_FILE; use crate::CloudConfigBundleLayers; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index bc5f59e479..68900517ab 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2612,6 +2612,7 @@ async fn handle_unauthorized( ); Ok(UnauthorizedRecoveryExecution { mode, phase }) } + Err(RefreshTokenError::Policy(error)) => Err(CodexErr::Fatal(error.to_string())), Err(RefreshTokenError::Permanent(failed)) => { session_telemetry.record_auth_recovery( mode, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 566e9c4912..9078708388 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -606,6 +606,10 @@ pub enum ThreadStoreConfig { /// Application configuration loaded from disk and merged with overrides. #[derive(Debug, Clone, PartialEq)] pub struct Config { + /// App-server-owned destination policy; other runtimes remain unmanaged. + pub application_network_policy: codex_http_client::NetworkPolicy, + /// Auth bootstrap routing installed by the app-server configuration owner. + pub application_auth_route_config: Option, /// Provenance for how this [`Config`] was derived (merged layers + enforced /// requirements). pub config_layer_stack: ConfigLayerStack, @@ -1655,7 +1659,14 @@ impl Config { /// Returns auth routing resolved from the effective feature configuration. pub fn auth_route_config(&self) -> AuthRouteConfig { - AuthRouteConfig::from_http_client_factory(self.http_client_factory()) + self.application_auth_route_config + .clone() + .unwrap_or_else(|| { + AuthRouteConfig::from_http_client_factory( + self.http_client_factory() + .with_network_policy(self.application_network_policy.clone()), + ) + }) } /// Creates the HTTP client factory resolved from the effective feature configuration. @@ -1665,7 +1676,8 @@ impl Config { } else { OutboundProxyPolicy::ReqwestDefault }; - let mut factory = HttpClientFactory::new(outbound_proxy_policy); + let mut factory = HttpClientFactory::new(outbound_proxy_policy) + .with_network_policy(self.application_network_policy.clone()); if !self.respect_system_proxy && self.features.enabled(Feature::SystemProxyFallback) { factory = factory.with_system_proxy_fallback(); } @@ -4329,6 +4341,8 @@ impl Config { sqlite: codex_state::SqliteConfig::from_sqlite_home(sqlite_home), log_dir, config_layer_stack, + application_network_policy: Default::default(), + application_auth_route_config: None, history, ephemeral: ephemeral.unwrap_or_default(), extra_config: None, diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index eebea1fc75..46d2d797a9 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -1279,8 +1279,17 @@ async fn prepare_realtime_start( .model_client .auth_manager() .unwrap_or_else(|| Arc::clone(&sess.services.auth_manager)); - let auth = auth_manager.auth().await; let config = sess.get_config().await; + let (auth, http_client_factory) = match auth_manager.auth_with_http_client_factory().await { + Some((auth, captured_factory)) if auth.is_chatgpt_auth() || auth.is_api_key_auth() => ( + Some(auth), + config + .http_client_factory() + .with_network_policy(captured_factory.network_policy().clone()), + ), + Some((auth, _)) => (Some(auth), config.http_client_factory()), + None => (None, config.http_client_factory()), + }; let transport = params .transport .clone() @@ -1387,7 +1396,7 @@ async fn prepare_realtime_start( } Ok(PreparedRealtimeConversationStart { api_provider, - http_client_factory: config.http_client_factory(), + http_client_factory, realtime_sideband_base_url, extra_headers: Some(extra_headers), client_managed_handoffs: params.client_managed_handoffs, diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 19153db088..1564fe4252 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -35,6 +35,7 @@ use codex_extension_api::UserInstructionsProvider; use codex_extension_api::empty_extension_registry; use codex_features::Feature; use codex_home::CodexHomeUserInstructionsProvider; +use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::built_in_model_providers; @@ -327,9 +328,26 @@ pub fn turn_permission_fields( (sandbox_policy, Some(permission_profile)) } +enum TestAuth { + Cached(CodexAuth), + Manager(Arc), +} + +impl TestAuth { + fn manager_for_home(&self, home: &Path) -> Arc { + match self { + Self::Cached(auth) => codex_core::test_support::auth_manager_from_auth_with_home( + auth.clone(), + home.to_path_buf(), + ), + Self::Manager(manager) => manager.clone(), + } + } +} + pub struct TestCodexBuilder { config_mutators: Vec>, - auth: CodexAuth, + auth: TestAuth, analytics_events_client: Option, pre_build_hooks: Vec>, workspace_setups: Vec>, @@ -363,7 +381,12 @@ impl TestCodexBuilder { } pub fn with_auth(mut self, auth: CodexAuth) -> Self { - self.auth = auth; + self.auth = TestAuth::Cached(auth); + self + } + + pub fn with_auth_manager(mut self, auth_manager: Arc) -> Self { + self.auth = TestAuth::Manager(auth_manager); self } @@ -715,7 +738,6 @@ impl TestCodexBuilder { mut test_env: TestEnv, environment_manager: Arc, ) -> anyhow::Result { - let auth = self.auth.clone(); let state_db = codex_core::init_state_db(&config).await; let thread_store = self .thread_store @@ -728,10 +750,7 @@ impl TestCodexBuilder { config.codex_home.clone(), )) }); - let auth_manager = codex_core::test_support::auth_manager_from_auth_with_home( - auth.clone(), - config.codex_home.to_path_buf(), - ); + let auth_manager = self.auth.manager_for_home(config.codex_home.as_path()); let models_manager = self .models_manager .clone() @@ -787,10 +806,7 @@ impl TestCodexBuilder { let new_conversation = match (resume_from, user_shell_override) { (Some(path), Some(user_shell_override)) => { - let auth_manager = codex_core::test_support::auth_manager_from_auth_with_home( - auth, - config.codex_home.to_path_buf(), - ); + let auth_manager = self.auth.manager_for_home(config.codex_home.as_path()); Box::pin( codex_core::test_support::resume_thread_from_rollout_with_user_shell_override( thread_manager.as_ref(), @@ -804,10 +820,7 @@ impl TestCodexBuilder { .await? } (Some(path), None) => { - let auth_manager = codex_core::test_support::auth_manager_from_auth_with_home( - auth, - config.codex_home.to_path_buf(), - ); + let auth_manager = self.auth.manager_for_home(config.codex_home.as_path()); Box::pin(thread_manager.resume_thread_from_rollout( config.clone(), path, @@ -1407,7 +1420,7 @@ pub fn test_codex() -> TestCodexBuilder { .disable(Feature::ShellSnapshot) .expect("test config should allow ShellSnapshot override"); })], - auth: CodexAuth::from_api_key("dummy"), + auth: TestAuth::Cached(CodexAuth::from_api_key("dummy")), analytics_events_client: None, pre_build_hooks: vec![], workspace_setups: vec![], diff --git a/codex-rs/core/tests/suite/auth_recovery_policy.rs b/codex-rs/core/tests/suite/auth_recovery_policy.rs new file mode 100644 index 0000000000..ca3b1b5f44 --- /dev/null +++ b/codex-rs/core/tests/suite/auth_recovery_policy.rs @@ -0,0 +1,93 @@ +//! A denied credential refresh ends 401 recovery without changing stored credentials. + +use std::sync::Arc; + +use anyhow::Result; +use codex_core::TurnInputRequest; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyController; +use codex_http_client::NetworkPolicyDenied; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; +use codex_login::AuthManager; +use codex_login::AuthRouteConfig; +use codex_protocol::protocol::EventMsg; +use codex_protocol::user_input::UserInput; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; +use wiremock::Mock; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn denied_oauth_refresh_terminates_401_recovery_and_preserves_auth() -> Result<()> { + let home = Arc::new(TempDir::new()?); + let auth_json = serde_json::to_vec(&json!({ + "auth_mode": "chatgpt", + "tokens": { + "id_token": "e30.e30.signature", + "access_token": "current-access-token", + "refresh_token": "current-refresh-token", + "account_id": "workspace-one" + }, + "last_refresh": chrono::Utc::now() + }))?; + std::fs::write(home.path().join("auth.json"), &auth_json)?; + let auth = AuthManager::shared( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + AuthRouteConfig::from_http_client_factory( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(NetworkPolicyController::default().policy()), + ), + ) + .await; + let original = auth.auth_cached().expect("stored ChatGPT auth"); + let server = start_mock_server().await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(header("authorization", "Bearer current-access-token")) + .and(header("chatgpt-account-id", "workspace-one")) + .respond_with(ResponseTemplate::new(401)) + .expect(/*r*/ 2) // Initial request and the existing disk-reload recovery step. + .mount(&server) + .await; + let mut builder = test_codex().with_home(home).with_auth_manager(auth.clone()); + let test = builder.build_with_auto_env(&server).await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "hello".into(), + text_elements: Vec::new(), + }])) + .await?; + let mut policy_error = None; + wait_for_event(&test.codex, |event| { + if let EventMsg::Error(error) = event { + policy_error.get_or_insert_with(|| error.message.clone()); + } + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!( + policy_error, + Some(format!("Fatal error: {}", NetworkPolicyDenied::Unavailable)) + ); + assert_eq!(auth.auth().await, Some(original)); + assert_eq!( + std::fs::read(test.home.path().join("auth.json"))?, + auth_json + ); + server.verify().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 5309fa63c9..8b4d4ff2fc 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -45,6 +45,7 @@ mod apply_patch_serialization; #[cfg(not(target_os = "windows"))] mod approvals; mod audio_truncation; +mod auth_recovery_policy; mod auto_review; mod catalog_permission_messages; mod cli_stream; diff --git a/codex-rs/core/tests/suite/realtime_sideband_endpoint.rs b/codex-rs/core/tests/suite/realtime_sideband_endpoint.rs index 39227ce21d..c4b4356476 100644 --- a/codex-rs/core/tests/suite/realtime_sideband_endpoint.rs +++ b/codex-rs/core/tests/suite/realtime_sideband_endpoint.rs @@ -28,10 +28,12 @@ use test_case::test_case; enum EndpointSelection { Configured, PerCall, + Revoked, } #[test_case(EndpointSelection::Configured; "configured endpoint")] #[test_case(EndpointSelection::PerCall; "per-call endpoint")] +#[test_case(EndpointSelection::Revoked; "policy revocation")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn existing_call_uses_selected_endpoint_and_runtime_auth( selection: EndpointSelection, @@ -46,6 +48,12 @@ async fn existing_call_uses_selected_endpoint_and_runtime_auth( let configured_endpoint = start_websocket_server(events.clone()).await; let call_endpoint = start_websocket_server(events).await; let configured_base_url = configured_endpoint.uri().to_string(); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish( + policy.revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); let mut headers = HeaderMap::new(); headers.insert( "x-runtime-auth", @@ -55,6 +63,9 @@ async fn existing_call_uses_selected_endpoint_and_runtime_auth( .with_auth(CodexAuth::Headers(AuthHeaders::new(headers))) .with_config(move |config| { config.experimental_realtime_ws_base_url = Some(configured_base_url); + if matches!(selection, EndpointSelection::Revoked) { + config.application_network_policy = policy; + } }); let test = builder.build_with_auto_env(&server).await?; @@ -79,7 +90,7 @@ async fn existing_call_uses_selected_endpoint_and_runtime_auth( transport: Some(ConversationStartTransport::ExistingCall { call_id: "rtc_existing".to_string(), sideband_base_url: match selection { - EndpointSelection::Configured => None, + EndpointSelection::Configured | EndpointSelection::Revoked => None, EndpointSelection::PerCall => Some(call_endpoint.uri().to_string()), }, }), @@ -96,7 +107,9 @@ async fn existing_call_uses_selected_endpoint_and_runtime_auth( .await?; let (selected, unused) = match selection { - EndpointSelection::Configured => (&configured_endpoint, &call_endpoint), + EndpointSelection::Configured | EndpointSelection::Revoked => { + (&configured_endpoint, &call_endpoint) + } EndpointSelection::PerCall => (&call_endpoint, &configured_endpoint), }; let handshake = selected.single_handshake(); @@ -114,6 +127,38 @@ async fn existing_call_uses_selected_endpoint_and_runtime_auth( ); assert!(unused.handshakes().is_empty()); + if matches!(selection, EndpointSelection::Revoked) { + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + let error = wait_for_event_match(&test.codex, |event| match event { + EventMsg::RealtimeConversationRealtime(event) => match &event.payload { + codex_protocol::protocol::RealtimeEvent::Error(message) => Some(message.clone()), + _ => None, + }, + _ => None, + }) + .await; + assert!(error.contains("revoked"), "{error}"); + // Recovery must not revive a sideband whose denial was reported as terminal. + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); + assert!( + !selected + .wait_for_handshakes( + /*expected*/ 2, + std::time::Duration::from_secs(/*secs*/ 1) + ) + .await + ); + assert_eq!(selected.handshakes().len(), 1); + } + test.codex.submit(Op::RealtimeConversationClose).await?; wait_for_event(&test.codex, |event| { matches!(event, EventMsg::RealtimeConversationClosed(_)) diff --git a/codex-rs/http-client/Cargo.toml b/codex-rs/http-client/Cargo.toml index a18d653d96..857798c47c 100644 --- a/codex-rs/http-client/Cargo.toml +++ b/codex-rs/http-client/Cargo.toml @@ -9,6 +9,7 @@ bytes = { workspace = true } codex-utils-rustls-provider = { workspace = true } futures = { workspace = true } http = { workspace = true } +http-body = "1.0.1" native-tls = "0.2" opentelemetry = { workspace = true } reqwest = { workspace = true, features = ["json", "rustls-tls-native-roots", "stream"] } diff --git a/codex-rs/http-client/src/client_builder.rs b/codex-rs/http-client/src/client_builder.rs index 6ffd914f99..a3ad63d007 100644 --- a/codex-rs/http-client/src/client_builder.rs +++ b/codex-rs/http-client/src/client_builder.rs @@ -18,6 +18,7 @@ use crate::BuildRouteAwareHttpClientError; use crate::ClientRouteClass; use crate::HttpClient; use crate::HttpClientFactory; +use crate::HttpClientTlsConfig; use crate::OutboundProxyRoute; use crate::chatgpt_cloudflare_cookies::ChatGptCookieStore; use crate::client::HttpClientBackend; @@ -33,6 +34,7 @@ use crate::with_chatgpt_cloudflare_cookie_store; /// bypass the factory and are restricted to documented exceptional or legacy compatibility paths. #[derive(Clone)] pub struct HttpClientBuilder { + http2_prior_knowledge: bool, pub(crate) default_headers: Option, follow_redirects: bool, pub(crate) redirect_observed: Option>, @@ -41,6 +43,7 @@ pub struct HttpClientBuilder { chatgpt_cookie_store: Option>, pub(crate) request_logging: RequestLogging, tls_backend: TlsBackend, + tls: HttpClientTlsConfig, } #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -82,6 +85,25 @@ impl HttpClientFactory { } impl HttpClientBuilder { + /// Builds a strict pooled client with explicit TLS settings and the factory's proxy policy. + /// Route and transport construction failures are returned when sending a request. + pub fn build_with_tls( + mut self, + http_client_factory: &HttpClientFactory, + route_class: ClientRouteClass, + tls: HttpClientTlsConfig, + ) -> HttpClient { + self.tls = tls; + crate::RouteAwareClientPool::with_builder(http_client_factory.clone(), route_class, self) + .with_tls_backend_fallback() + .into_client() + } + + /// Uses HTTP/2 for SDKs such as gRPC that require framed bidirectional bodies. + pub fn http2_prior_knowledge(mut self) -> Self { + self.http2_prior_knowledge = true; + self + } pub fn new() -> Self { Self::default() } @@ -178,12 +200,20 @@ impl HttpClientBuilder { route: &OutboundProxyRoute, ) -> Result { self.chatgpt_cookie_store = http_client_factory.chatgpt_cookie_store(); + let explicit_roots = self.tls.root_certificate.is_some(); let (builder, request_logging, default_headers) = self.into_reqwest_parts(); - let inner = http_client_factory.build_reqwest_client_for_resolved_route( + let builder = crate::outbound_proxy::configure_builder_for_resolved_route( builder, route_class, route, )?; + let inner = if explicit_roots { + builder + .build() + .map_err(BuildRouteAwareHttpClientError::ExplicitTls)? + } else { + build_reqwest_client_with_custom_ca(builder)? + }; Ok(TransportClient::new( inner, request_logging, @@ -326,10 +356,21 @@ impl HttpClientBuilder { fn base_reqwest_builder(self) -> reqwest::ClientBuilder { let mut builder = reqwest::Client::builder(); - if self.tls_backend == TlsBackend::Rustls { + if self.http2_prior_knowledge { + builder = builder.http2_prior_knowledge(); + } + if self.tls_backend == TlsBackend::Rustls || self.tls.client_identity.is_some() { ensure_rustls_crypto_provider(); builder = builder.use_rustls_tls(); } + if let Some(certificate) = self.tls.root_certificate { + builder = builder + .tls_built_in_root_certs(false) + .add_root_certificate(certificate); + } + if let Some(identity) = self.tls.client_identity { + builder = builder.identity(identity).https_only(true); + } if !self.follow_redirects { builder = builder.redirect(reqwest::redirect::Policy::none()); } else if let Some(redirect_observed) = self.redirect_observed { @@ -354,6 +395,7 @@ impl HttpClientBuilder { impl Default for HttpClientBuilder { fn default() -> Self { Self { + http2_prior_knowledge: false, default_headers: None, follow_redirects: true, redirect_observed: None, @@ -362,6 +404,7 @@ impl Default for HttpClientBuilder { chatgpt_cookie_store: None, request_logging: RequestLogging::Enabled, tls_backend: TlsBackend::TransportDefault, + tls: HttpClientTlsConfig::default(), } } } diff --git a/codex-rs/http-client/src/client_tls.rs b/codex-rs/http-client/src/client_tls.rs new file mode 100644 index 0000000000..0a6ca00654 --- /dev/null +++ b/codex-rs/http-client/src/client_tls.rs @@ -0,0 +1,37 @@ +//! Validated TLS material for strict, route-aware HTTP clients. +//! Explicit roots replace shared environment roots; PEM identities require HTTPS and rustls. + +use std::fmt; + +use crate::HttpError; + +/// TLS configuration consumed by [`crate::HttpClientBuilder::build_with_tls`]. +/// Keeping this separate from the builder prevents combining it with legacy fallback APIs. +#[derive(Clone, Default)] +pub struct HttpClientTlsConfig { + pub(crate) root_certificate: Option, + pub(crate) client_identity: Option, +} + +impl fmt::Debug for HttpClientTlsConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HttpClientTlsConfig") + .field("explicit_root", &self.root_certificate.is_some()) + .field("client_identity", &self.client_identity.is_some()) + .finish() + } +} + +impl HttpClientTlsConfig { + /// Trusts this PEM certificate instead of system or environment-provided roots. + pub fn with_root_certificate_pem(mut self, pem: &[u8]) -> Result { + self.root_certificate = Some(reqwest::Certificate::from_pem(pem)?); + Ok(self) + } + + /// Installs a PEM certificate chain and private key and requires HTTPS requests. + pub fn with_client_identity_pem(mut self, pem: &[u8]) -> Result { + self.client_identity = Some(reqwest::Identity::from_pem(pem)?); + Ok(self) + } +} diff --git a/codex-rs/http-client/src/lib.rs b/codex-rs/http-client/src/lib.rs index 9d72958620..0a2bf7e5b3 100644 --- a/codex-rs/http-client/src/lib.rs +++ b/codex-rs/http-client/src/lib.rs @@ -2,6 +2,7 @@ mod chatgpt_cloudflare_cookies; mod chatgpt_hosts; mod client; mod client_builder; +mod client_tls; mod custom_ca; mod error; mod network_policy; @@ -20,6 +21,7 @@ pub use crate::chatgpt_hosts::is_allowed_chatgpt_host; pub use crate::client::HttpClient; pub use crate::client::HttpError; pub use crate::client_builder::HttpClientBuilder; +pub use crate::client_tls::HttpClientTlsConfig; pub use crate::custom_ca::BuildCustomCaTransportError; /// Test-only subprocess hook for custom CA coverage. /// diff --git a/codex-rs/http-client/src/network_policy.rs b/codex-rs/http-client/src/network_policy.rs index fa92421d60..4c3ce3e5f0 100644 --- a/codex-rs/http-client/src/network_policy.rs +++ b/codex-rs/http-client/src/network_policy.rs @@ -197,6 +197,16 @@ impl NetworkPolicy { } } + /// Narrows a bootstrap client's existing local policy to exact endpoint URLs. + /// This cannot grant access to a destination denied by the underlying policy. + pub fn restrict_to_endpoints(mut self, endpoints: BTreeSet) -> Self { + self.endpoints = Some(Arc::new(match &self.endpoints { + Some(existing) => existing.intersection(&endpoints).cloned().collect(), + None => endpoints, + })); + self + } + /// Binds a content client to the current account so retained credentials cannot /// be used after another workspace's policy has been installed. pub fn for_current_account(mut self) -> Self { diff --git a/codex-rs/http-client/src/outbound_proxy.rs b/codex-rs/http-client/src/outbound_proxy.rs index 429711f0bb..7094e8e8f0 100644 --- a/codex-rs/http-client/src/outbound_proxy.rs +++ b/codex-rs/http-client/src/outbound_proxy.rs @@ -323,7 +323,7 @@ impl HttpClientFactory { } /// Builds a reqwest client for a concrete outbound route. - pub fn build_reqwest_client( + pub(crate) fn build_reqwest_client( &self, builder: reqwest::ClientBuilder, request_url: &str, @@ -336,16 +336,6 @@ impl HttpClientFactory { self.outbound_proxy_policy, ) } - - pub(crate) fn build_reqwest_client_for_resolved_route( - &self, - builder: reqwest::ClientBuilder, - route_class: ClientRouteClass, - route: &OutboundProxyRoute, - ) -> Result { - let builder = configure_builder_for_resolved_route(builder, route_class, route)?; - build_reqwest_client_with_custom_ca(builder).map_err(Into::into) - } } fn resolve_proxy_route( @@ -450,6 +440,9 @@ pub enum BuildRouteAwareHttpClientError { #[error(transparent)] CustomCa(#[from] BuildCustomCaTransportError), + #[error("Failed to build HTTP client with explicit TLS configuration: {0}")] + ExplicitTls(reqwest::Error), + #[error("Failed to configure outbound proxy selected for {route_class}")] InvalidProxyConfig { route_class: ClientRouteClass }, } @@ -458,7 +451,8 @@ impl From for io::Error { fn from(error: BuildRouteAwareHttpClientError) -> Self { match error { BuildRouteAwareHttpClientError::CustomCa(error) => error.into(), - BuildRouteAwareHttpClientError::InvalidProxyConfig { .. } => io::Error::other(error), + BuildRouteAwareHttpClientError::InvalidProxyConfig { .. } + | BuildRouteAwareHttpClientError::ExplicitTls(_) => io::Error::other(error), } } } @@ -503,7 +497,7 @@ fn configure_proxy_for_route( configure_builder_for_resolved_route(builder, route_class, &route) } -fn configure_builder_for_resolved_route( +pub(crate) fn configure_builder_for_resolved_route( builder: reqwest::ClientBuilder, route_class: ClientRouteClass, route: &OutboundProxyRoute, diff --git a/codex-rs/http-client/src/response.rs b/codex-rs/http-client/src/response.rs index ef45a3102b..ccb64a794b 100644 --- a/codex-rs/http-client/src/response.rs +++ b/codex-rs/http-client/src/response.rs @@ -1,10 +1,15 @@ //! HTTP responses retain application authorization while their bodies are consumed. use std::ops::Deref; +use std::pin::Pin; +use std::task::Context; +use std::task::Poll; use bytes::Bytes; +use futures::FutureExt; use futures::Stream; use futures::StreamExt; +use futures::future::BoxFuture; use futures::stream; use serde::de::DeserializeOwned; @@ -19,6 +24,18 @@ pub struct HttpResponse { } impl HttpResponse { + /// Preserves HTTP frames, including gRPC trailers, while enforcing revocation. + pub fn into_http_response( + self, + ) -> http::Response + Send> { + let response: http::Response = self.inner.into(); + let revoked = self.permit.clone(); + response.map(|body| PolicyBody { + body: Some(Box::pin(body)), + permit: self.permit, + revoked: async move { revoked.revoked().await }.boxed(), + }) + } pub(crate) fn new(inner: reqwest::Response, permit: NetworkPermit) -> Self { Self { inner, permit } } @@ -66,6 +83,46 @@ impl HttpResponse { } } +struct PolicyBody { + body: Option>>, + permit: NetworkPermit, + revoked: BoxFuture<'static, ()>, +} + +impl> http_body::Body for PolicyBody { + type Data = Bytes; + type Error = HttpError; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, HttpError>>> { + let body = self.get_mut(); + let Some(inner) = body.body.as_mut() else { + return Poll::Ready(None); + }; + if body.permit.check().is_err() || body.revoked.poll_unpin(cx).is_ready() { + body.body = None; + return Poll::Ready(Some(Err(crate::NetworkPolicyDenied::Revoked.into()))); + } + inner + .as_mut() + .poll_frame(cx) + .map(|frame| frame.map(|frame| frame.map_err(HttpError::from))) + } + + fn is_end_stream(&self) -> bool { + self.body + .as_ref() + .is_none_or(http_body::Body::is_end_stream) + } + fn size_hint(&self) -> http_body::SizeHint { + self.body + .as_ref() + .map_or_else(http_body::SizeHint::default, http_body::Body::size_hint) + } +} + impl Deref for HttpResponse { type Target = reqwest::Response; diff --git a/codex-rs/http-client/src/route_aware_client_pool.rs b/codex-rs/http-client/src/route_aware_client_pool.rs index 8a2fec5339..7e9c795b0e 100644 --- a/codex-rs/http-client/src/route_aware_client_pool.rs +++ b/codex-rs/http-client/src/route_aware_client_pool.rs @@ -32,6 +32,7 @@ const MAX_CACHED_ROUTES: usize = 16; #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] enum CustomCaFallback { + LegacyDirect, #[default] Disabled, LegacyTransportDefault, @@ -128,7 +129,8 @@ impl RouteAwareRequestError { match self { Self::Route(RouteAwareClientPoolError::Build( - BuildRouteAwareHttpClientError::CustomCa(_), + BuildRouteAwareHttpClientError::CustomCa(_) + | BuildRouteAwareHttpClientError::ExplicitTls(_), )) => Some(RouteFailureClass::TlsError), Self::Route(RouteAwareClientPoolError::Build( BuildRouteAwareHttpClientError::InvalidProxyConfig { .. }, @@ -215,7 +217,7 @@ impl RouteAwareClientPool { } /// Exposes the ordinary request builder while sending through this policy-aware pool. - pub(crate) fn into_client(self) -> HttpClient { + pub fn into_client(self) -> HttpClient { HttpClient { backend: HttpClientBackend::Routed(Arc::new(self)), } @@ -332,6 +334,12 @@ impl RouteAwareClientPool { self } + /// Preserves the legacy sandbox client's direct routing and custom-CA fallback. + pub fn with_legacy_direct_proxy_and_custom_ca_fallback(mut self) -> Self { + self.custom_ca_fallback = CustomCaFallback::LegacyDirect; + self + } + /// Creates a pool that retains the Cloudflare cookies required by ChatGPT endpoints. pub fn with_chatgpt_cloudflare_cookies( http_client_factory: HttpClientFactory, @@ -447,9 +455,13 @@ impl RouteAwareClientPool { F: FnOnce(String) -> Fut, Fut: Future>, { - let route = resolve_route(request_url.to_string()) - .await - .map_err(RouteAwareClientPoolError::Resolve)?; + let route = if self.custom_ca_fallback == CustomCaFallback::LegacyDirect { + OutboundProxyRoute::Direct + } else { + resolve_route(request_url.to_string()) + .await + .map_err(RouteAwareClientPoolError::Resolve)? + }; if let Some(rustls_clients) = self.rustls_clients.as_ref() && let Ok(url) = reqwest::Url::parse(request_url) && rustls_clients.requires_rustls(&url, &route) @@ -475,6 +487,9 @@ impl RouteAwareClientPool { self.http_client_factory.outbound_proxy_policy(), self.custom_ca_fallback, ) { + (_, CustomCaFallback::LegacyDirect) => { + client_builder.build_with_custom_ca_fallback(ProxyRouting::Direct) + } (OutboundProxyPolicy::ReqwestDefault, CustomCaFallback::LegacyTransportDefault) => { client_builder.build_with_custom_ca_fallback(ProxyRouting::TransportDefault) } diff --git a/codex-rs/login/src/auth/change_state_tests.rs b/codex-rs/login/src/auth/change_state_tests.rs index 4022fbbfe8..25abff139a 100644 --- a/codex-rs/login/src/auth/change_state_tests.rs +++ b/codex-rs/login/src/auth/change_state_tests.rs @@ -1,5 +1,7 @@ use super::*; use base64::Engine; +use codex_http_client::DestinationPolicy; +use codex_http_client::NetworkPolicyController; use pretty_assertions::assert_eq; use serde_json::json; @@ -19,7 +21,16 @@ fn chatgpt_auth(user: &str, workspace: &str, token: &str) -> CodexAuth { #[test] fn auth_owner_generation_distinguishes_refreshes_from_identity_changes() { - let manager = AuthManager::from_optional_auth_for_testing(/*auth*/ None); + let mut manager = AuthManager::from_optional_auth_for_testing(/*auth*/ None); + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + Arc::get_mut(&mut manager).unwrap().auth_route_config = + AuthRouteConfig::from_http_client_factory( + manager + .http_client_factory() + .with_network_policy(policy.clone()), + ); + let endpoint = "https://example.com/".parse().unwrap(); let mut changes = manager.auth_change_state_receiver(); let mut legacy_changes = manager.auth_change_receiver(); let initial = chatgpt_auth("user-a", "workspace-a", "token-1"); @@ -34,7 +45,14 @@ fn auth_owner_generation_distinguishes_refreshes_from_identity_changes() { (Some(CodexAuth::from_api_key("key-1")), 7, 6), (Some(CodexAuth::from_api_key("key-2")), 8, 7), ] { + let previous_owner = changes.borrow().owner_generation; + let factory = manager.http_client_factory(); manager.set_cached_auth(auth); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + assert_eq!( + factory.network_policy().acquire(&endpoint).is_ok(), + owner_generation == previous_owner, + ); assert_eq!( *changes.borrow_and_update(), AuthChangeState { @@ -74,3 +92,42 @@ async fn auth_owner_generation_preserves_logout_and_switches_when_coalesced() { }, ); } + +#[tokio::test] +async fn captured_credentials_are_revoked_if_account_changes_before_client_construction() { + let mut manager = AuthManager::from_optional_auth_for_testing(/*auth*/ None); + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + Arc::get_mut(&mut manager).unwrap().auth_route_config = + AuthRouteConfig::from_http_client_factory( + manager + .http_client_factory() + .with_network_policy(policy.clone()), + ); + let initial = chatgpt_auth("user-a", "workspace-a", "token-a"); + manager.set_cached_auth(Some(initial.clone())); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let (auth, factory) = manager.auth_with_http_client_factory().await.unwrap(); + + manager.set_cached_auth(Some(chatgpt_auth("user-b", "workspace-b", "token-b"))); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let client = crate::default_client::create_client_for_route_async( + factory, + "https://example.com/".to_string(), + codex_http_client::ClientRouteClass::Api, + crate::default_client::ClientRedirectPolicy::Default, + ) + .await + .unwrap(); + assert_eq!(auth.get_account_id(), initial.get_account_id()); + assert!(matches!( + client + .get("https://example.com/") + .bearer_auth(auth.get_token().unwrap()) + .send() + .await, + Err(codex_http_client::RouteAwareRequestError::Policy( + codex_http_client::NetworkPolicyDenied::Revoked + )) + )); +} diff --git a/codex-rs/login/src/auth/default_client.rs b/codex-rs/login/src/auth/default_client.rs index 9a5ace9502..461f658d18 100644 --- a/codex-rs/login/src/auth/default_client.rs +++ b/codex-rs/login/src/auth/default_client.rs @@ -250,6 +250,22 @@ pub fn create_client_for_route( ClientRedirectPolicy::Default => default_http_client_builder(), ClientRedirectPolicy::Reject => default_http_client_builder().without_redirects(), }; + if http_client_factory.network_policy().is_managed() { + let mut pool = codex_http_client::RouteAwareClientPool::with_builder( + http_client_factory.clone(), + route_class, + builder, + ); + if is_sandboxed() { + pool = pool.with_legacy_direct_proxy_and_custom_ca_fallback(); + } else if matches!( + http_client_factory.outbound_proxy_policy(), + OutboundProxyPolicy::ReqwestDefault + ) { + pool = pool.with_legacy_custom_ca_fallback(); + } + return Ok(pool.into_client()); + } if matches!( http_client_factory.outbound_proxy_policy(), OutboundProxyPolicy::ReqwestDefault @@ -355,7 +371,7 @@ pub(crate) fn create_raw_auth_client( auth_route_config: &AuthRouteConfig, ) -> Result { auth_route_config - .http_client_factory() + .authentication_factory(endpoint) .build_client_without_request_logging(endpoint, ClientRouteClass::Auth) } @@ -365,7 +381,7 @@ pub(crate) fn create_default_auth_client( auth_route_config: &AuthRouteConfig, ) -> Result { create_client_for_route( - auth_route_config.http_client_factory(), + &auth_route_config.authentication_factory(endpoint), endpoint, ClientRouteClass::Auth, ClientRedirectPolicy::Default, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index e67b152973..14cfee9221 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -222,6 +222,9 @@ pub enum RefreshTokenError { Permanent(#[from] RefreshTokenFailedError), #[error(transparent)] Transient(#[from] std::io::Error), + /// Denied by application policy; neither retry nor cache as a credential failure. + #[error(transparent)] + Policy(#[from] codex_http_client::NetworkPolicyDenied), } /// Error returned when constructing an [`AuthManager`] from resolved configuration. @@ -295,7 +298,16 @@ impl RefreshTokenError { pub fn failed_reason(&self) -> Option { match self { Self::Permanent(error) => Some(error.reason), - Self::Transient(_) => None, + Self::Transient(_) | Self::Policy(_) => None, + } + } +} + +impl From for RefreshTokenError { + fn from(error: codex_http_client::HttpError) -> Self { + match error { + codex_http_client::HttpError::Policy(denied) => Self::Policy(denied), + error => Self::Transient(std::io::Error::other(error)), } } } @@ -305,6 +317,9 @@ impl From for std::io::Error { match err { RefreshTokenError::Permanent(failed) => std::io::Error::other(failed), RefreshTokenError::Transient(inner) => inner, + RefreshTokenError::Policy(error) => { + std::io::Error::new(std::io::ErrorKind::PermissionDenied, error) + } } } } @@ -1633,6 +1648,9 @@ async fn request_chatgpt_token_refresh( { Ok(response) => Ok(response), Err(OAuthError::Rejected(rejection)) => { + if let Some(codex_http_client::HttpError::Policy(denied)) = rejection.body_read_error { + return Err(denied.into()); + } let status = rejection.status; let detail = &rejection.detail; tracing::error!(%status, ?detail, "Failed to refresh token"); @@ -1651,7 +1669,8 @@ async fn request_chatgpt_token_refresh( ))) } } - Err(error @ (OAuthError::Transport(_) | OAuthError::InvalidResponse)) => { + Err(OAuthError::Transport(error)) => Err(error.into()), + Err(error @ OAuthError::InvalidResponse) => { Err(RefreshTokenError::Transient(std::io::Error::other(error))) } } @@ -2114,6 +2133,21 @@ fn default_agent_identity_authapi_base_url() -> Option { } impl AuthManager { + /// Returns the application policy associated with this account owner. + pub fn application_network_policy(&self) -> codex_http_client::NetworkPolicy { + self.auth_route_config + .http_client_factory() + .network_policy() + .clone() + } + + /// Creates content clients bound to the account currently owned by this manager. + pub fn http_client_factory(&self) -> HttpClientFactory { + self.auth_route_config + .http_client_factory() + .clone() + .with_network_policy(self.application_network_policy().for_current_account()) + } /// Create a new manager loading the initial auth using the provided /// preferred auth method. Errors loading auth are swallowed; `auth()` will /// simply return `None` in that case so callers can treat it as an @@ -2372,6 +2406,14 @@ impl AuthManager { self.auth_cached() } + /// Refreshes auth, then captures credentials and their account-bound factory together. + /// The auth read lock prevents an identity change between the two snapshots. + pub async fn auth_with_http_client_factory(&self) -> Option<(CodexAuth, HttpClientFactory)> { + self.auth().await; + let cached = self.inner.read().ok()?; + Some((cached.auth.clone()?, self.http_client_factory())) + } + pub async fn agent_identity_auth( &self, policy: AgentIdentityAuthPolicy, @@ -2552,6 +2594,7 @@ impl AuthManager { cached_auth } RefreshTokenError::Transient(_) => None, + RefreshTokenError::Policy(_) => cached_auth, } } }; @@ -2591,6 +2634,12 @@ impl AuthManager { !Self::auths_equal_for_refresh(previous, new_auth.as_ref()); let owner_changed = auth_changed_for_refresh && !same_owner(previous, new_auth.as_ref()); + if owner_changed { + self.auth_route_config + .http_client_factory() + .network_policy() + .invalidate(); + } if auth_changed_for_refresh { guard.permanent_refresh_failure = None; } diff --git a/codex-rs/login/src/auth/workload_identity.rs b/codex-rs/login/src/auth/workload_identity.rs index e71a57b836..06cf649027 100644 --- a/codex-rs/login/src/auth/workload_identity.rs +++ b/codex-rs/login/src/auth/workload_identity.rs @@ -161,7 +161,8 @@ fn resolve_config( assertion_file, environment: auth_environment, federation_rule_id, - http_client_factory: auth_route_config.http_client_factory().clone(), + http_client_factory: auth_route_config + .authentication_factory(auth_environment.token_url()?.as_str()), token_url: auth_environment.token_url()?, workload_identity_context, })) @@ -466,6 +467,12 @@ impl ExternalAuth for WorkloadIdentityExternalAuth { } fn classify_error(&self, error: std::io::Error) -> RefreshTokenError { + if let Some(WorkloadIdentityError::Policy(denied)) = error + .get_ref() + .and_then(|source| source.downcast_ref::()) + { + return RefreshTokenError::Policy(*denied); + } if error .get_ref() .and_then(|source| source.downcast_ref::()) diff --git a/codex-rs/login/src/device_code_auth.rs b/codex-rs/login/src/device_code_auth.rs index 263abd5aaf..f9f8f1993b 100644 --- a/codex-rs/login/src/device_code_auth.rs +++ b/codex-rs/login/src/device_code_auth.rs @@ -164,10 +164,11 @@ fn print_device_code_prompt(verification_url: &str, code: &str) { pub async fn request_device_code(opts: &ServerOptions) -> std::io::Result { let base_url = opts.issuer.trim_end_matches('/'); - // The route selected for the issuer is reused for all device-auth endpoint paths; the endpoint - // paths are not resolved separately. - let client = create_raw_auth_client(base_url, &opts.auth_route_config)?; let api_base_url = format!("{base_url}/api/accounts"); + let client = create_raw_auth_client( + &format!("{api_base_url}/deviceauth/usercode"), + &opts.auth_route_config, + )?; let uc = request_user_code(&client, &api_base_url, &opts.client_id).await?; Ok(DeviceCode { @@ -183,8 +184,11 @@ pub async fn complete_device_code_login( device_code: DeviceCode, ) -> std::io::Result<()> { let base_url = opts.issuer.trim_end_matches('/'); - let client = create_raw_auth_client(base_url, &opts.auth_route_config)?; let api_base_url = format!("{base_url}/api/accounts"); + let client = create_raw_auth_client( + &format!("{api_base_url}/deviceauth/token"), + &opts.auth_route_config, + )?; let code_resp = poll_for_token( &client, diff --git a/codex-rs/login/src/oauth/client.rs b/codex-rs/login/src/oauth/client.rs index 49ec82f117..f9a1d453d5 100644 --- a/codex-rs/login/src/oauth/client.rs +++ b/codex-rs/login/src/oauth/client.rs @@ -119,10 +119,10 @@ impl<'a> OAuthClient<'a> { .await, ))); } - response - .json() - .await - .map_err(|_| OAuthError::InvalidResponse) + response.json().await.map_err(|error| match error { + error @ codex_http_client::HttpError::Policy(_) => OAuthError::Transport(error), + _ => OAuthError::InvalidResponse, + }) } } diff --git a/codex-rs/login/src/oauth/client_tests.rs b/codex-rs/login/src/oauth/client_tests.rs index d81e2d2698..17e93ed5ce 100644 --- a/codex-rs/login/src/oauth/client_tests.rs +++ b/codex-rs/login/src/oauth/client_tests.rs @@ -245,3 +245,59 @@ fn decode_form(body: &[u8]) -> Value { ) .unwrap() } + +#[tokio::test] +async fn success_body_revocation_preserves_the_policy_error() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let controller = codex_http_client::NetworkPolicyController::default(); + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); + let http = codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ) + .with_network_policy(controller.policy()) + .build_client(&endpoint, codex_http_client::ClientRouteClass::Auth) + .unwrap(); + let server = async { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + while !request.ends_with(b"refresh_token=refresh") { + let mut buffer = [0; 4096]; + let read = stream.read(&mut buffer).await.unwrap(); + assert!(read > 0); + request.extend_from_slice(&buffer[..read]); + } + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 100\r\n\r\n{\"access_token\":") + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(/*millis*/ 50)).await; + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + }; + let form = oauth(&http, &endpoint, ErrorBodyLimit::Unlimited); + let (result, ()) = tokio::time::timeout(Duration::from_secs(/*secs*/ 5), async { + tokio::join!( + form.refresh::(RefreshTokenGrant { + refresh_token: "refresh", + resource: None + }), + server + ) + }) + .await + .unwrap(); + assert!(matches!( + result, + Err(OAuthError::Transport(codex_http_client::HttpError::Policy( + codex_http_client::NetworkPolicyDenied::Revoked + ))) + )); +} diff --git a/codex-rs/login/src/outbound_proxy.rs b/codex-rs/login/src/outbound_proxy.rs index 84dfccf989..430b99d453 100644 --- a/codex-rs/login/src/outbound_proxy.rs +++ b/codex-rs/login/src/outbound_proxy.rs @@ -7,6 +7,7 @@ use codex_http_client::HttpClientFactory; #[derive(Debug, Clone, PartialEq, Eq)] pub struct AuthRouteConfig { http_client_factory: HttpClientFactory, + local_bootstrap_factory: Option, } impl AuthRouteConfig { @@ -14,9 +15,30 @@ impl AuthRouteConfig { pub fn from_http_client_factory(http_client_factory: HttpClientFactory) -> Self { Self { http_client_factory, + local_bootstrap_factory: None, } } + /// Installs the configuration owner's local-only policy for auth discovery. + /// Only auth-owned endpoint constructors can access this factory. + pub fn with_local_bootstrap_factory(mut self, factory: HttpClientFactory) -> Self { + self.local_bootstrap_factory = Some(factory); + self + } + + pub(crate) fn authentication_factory(&self, endpoint: &str) -> HttpClientFactory { + let Some(factory) = &self.local_bootstrap_factory else { + return self.http_client_factory.clone(); + }; + let endpoints = endpoint.parse().into_iter().collect(); + factory.clone().with_network_policy( + factory + .network_policy() + .clone() + .restrict_to_endpoints(endpoints), + ) + } + /// Returns the HTTP client factory represented by this routing configuration. pub fn http_client_factory(&self) -> &HttpClientFactory { &self.http_client_factory diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 084a8ba574..35999a2b73 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -715,7 +715,7 @@ pub(crate) async fn exchange_code_for_tokens( auth_route_config: &AuthRouteConfig, ) -> io::Result<(ExchangedTokens, HttpClient)> { let token_endpoint = format!("{}/oauth/token", issuer.trim_end_matches('/')); - let factory = auth_route_config.http_client_factory(); + let factory = auth_route_config.authentication_factory(&token_endpoint); let allows_fallback = factory.allows_system_proxy_fallback(); let redirect_observed = Arc::new(AtomicBool::new(false)); let mut builder = HttpClientBuilder::new().without_request_logging(); @@ -733,7 +733,7 @@ pub(crate) async fn exchange_code_for_tokens( "starting oauth token exchange" ); let (mut client, mut result) = send_code_exchange_request( - factory, + &factory, builder.clone(), &token_endpoint, client_id, diff --git a/codex-rs/login/tests/suite/login_server_e2e.rs b/codex-rs/login/tests/suite/login_server_e2e.rs index 8e771735f5..1764794f74 100644 --- a/codex-rs/login/tests/suite/login_server_e2e.rs +++ b/codex-rs/login/tests/suite/login_server_e2e.rs @@ -123,11 +123,25 @@ async fn end_to_end_login_flow_persists_auth_json() -> Result<()> { // Run server in background let server_home = codex_home.clone(); + let content = codex_http_client::NetworkPolicyController::default(); + let local = codex_http_client::NetworkPolicyController::default(); + let local_policy = local.policy(); + local.publish( + local_policy.revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); + let factory = codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ); + let routes = codex_login::AuthRouteConfig::from_http_client_factory( + factory.clone().with_network_policy(content.policy()), + ) + .with_local_bootstrap_factory(factory.with_network_policy(local_policy)); let opts = ServerOptions { codex_home: server_home, cli_auth_credentials_store_mode: AuthCredentialsStoreMode::File, - auth_route_config: codex_login::test_support::transport_default_auth_route_config(), + auth_route_config: routes, client_id: codex_login::CLIENT_ID.to_string(), issuer, port: 0, diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 44f1197fa6..308ce8b2e5 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -187,6 +187,8 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R .clone(); let mut config = Config { + application_network_policy: Default::default(), + application_auth_route_config: None, config_layer_stack: ConfigLayerStack::default(), startup_warnings: Vec::new(), bypass_hook_trust: false, diff --git a/codex-rs/tui/src/app/background_requests.rs b/codex-rs/tui/src/app/background_requests.rs index 4f2c2fdadb..bfa33203da 100644 --- a/codex-rs/tui/src/app/background_requests.rs +++ b/codex-rs/tui/src/app/background_requests.rs @@ -1349,16 +1349,32 @@ mod tests { let server = wiremock::MockServer::start().await; let thread_id = ThreadId::new(); app.config.chatgpt_base_url = server.uri(); + app.cli_kv_overrides = vec![( + "chatgpt_base_url".to_string(), + toml::Value::String(server.uri()), + )]; app.config.cli_auth_credentials_store_mode = AuthCredentialsStoreMode::File; + wiremock::Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/api/codex/config/bundle")) + .respond_with(wiremock::ResponseTemplate::new(/*s*/ 200).set_body_string("{}")) + .mount(&server) + .await; write_chatgpt_auth( app.config.codex_home.as_path(), ChatGptAuthFixture::new("chatgpt-token").account_id("account-123"), AuthCredentialsStoreMode::File, ) .expect("write ChatGPT authentication"); - let app_server = crate::start_embedded_app_server_for_picker(&app.config) - .await - .expect("start authenticated embedded app server"); + let app_server = crate::start_app_server_for_picker( + &app.config, + &crate::AppServerTarget::Embedded, + app.cli_kv_overrides.clone(), + app.loader_overrides.clone(), + /*state_db*/ None, + app.environment_manager.clone(), + ) + .await + .expect("start authenticated embedded app server"); write_chatgpt_auth( app.config.codex_home.as_path(), ChatGptAuthFixture::new("different-token").account_id("different-account"), diff --git a/codex-rs/tui/src/app/session_picker.rs b/codex-rs/tui/src/app/session_picker.rs index 018a3f5ad1..de09c62277 100644 --- a/codex-rs/tui/src/app/session_picker.rs +++ b/codex-rs/tui/src/app/session_picker.rs @@ -13,12 +13,16 @@ impl App { // picker on a fresh event-loop iteration; its auth manager needs more stack headroom. let picker_config = self.config.clone(); let picker_target = self.app_server_target.clone(); + let picker_cli_kv_overrides = self.cli_kv_overrides.clone(); + let picker_loader_overrides = self.loader_overrides.clone(); let picker_state_db = self.state_db.clone(); let picker_environment_manager = Arc::clone(&self.environment_manager); let picker_app_server = tokio::spawn(async move { crate::start_app_server_for_picker( &picker_config, &picker_target, + picker_cli_kv_overrides, + picker_loader_overrides, picker_state_db, picker_environment_manager, ) diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 5be7b53631..4e17f2df3a 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -7914,6 +7914,8 @@ async fn in_app_resume_uses_configured_or_explicit_cwd() -> Result<()> { crate::start_app_server_for_picker( &config, &crate::AppServerTarget::Embedded, + Vec::new(), + LoaderOverrides::without_managed_config_for_tests(), /*state_db*/ None, Arc::clone(&environment_manager), ) @@ -8059,6 +8061,8 @@ async fn remembered_current_cwd_stays_at_launch_across_in_app_resumes() -> Resul let mut app_server = Box::pin(crate::start_app_server_for_picker( &config, &crate::AppServerTarget::Embedded, + Vec::new(), + LoaderOverrides::without_managed_config_for_tests(), state_db.clone(), Arc::new(EnvironmentManager::default_for_tests()), )) diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 0aed8985ae..e37d17ca94 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -571,6 +571,8 @@ async fn start_app_server( pub(crate) async fn start_app_server_for_picker( config: &Config, target: &AppServerTarget, + cli_kv_overrides: Vec<(String, toml::Value)>, + loader_overrides: LoaderOverrides, state_db: Option, environment_manager: Arc, ) -> color_eyre::Result { @@ -580,8 +582,8 @@ pub(crate) async fn start_app_server_for_picker( &mut target, Arg0DispatchPaths::default(), config.clone(), - Vec::new(), - LoaderOverrides::default(), + cli_kv_overrides, + loader_overrides, /*strict_config*/ false, CloudConfigBundleLoader::default(), codex_feedback::CodexFeedback::new(), @@ -2691,6 +2693,8 @@ requires_openai_auth = {requires_openai_auth} let mut app_server = start_app_server_for_picker( &final_config, &AppServerTarget::Embedded, + Vec::new(), + LoaderOverrides::without_managed_config_for_tests(), state_db, Arc::new(EnvironmentManager::default_for_tests()), ) diff --git a/codex-rs/tui/src/status/tests.rs b/codex-rs/tui/src/status/tests.rs index d72fbdc815..bb65053072 100644 --- a/codex-rs/tui/src/status/tests.rs +++ b/codex-rs/tui/src/status/tests.rs @@ -369,6 +369,22 @@ async fn status_snapshot_includes_reasoning_details() { #[tokio::test] async fn status_snapshot_shows_chatgpt_plan_without_email() { let temp_home = TempDir::new().expect("temp home"); + let profile_path = temp_home.path().join("work.config.toml"); + let loader_overrides = LoaderOverrides { + user_config_path: Some(profile_path.abs()), + user_config_profile: Some("work".parse().expect("profile")), + ..LoaderOverrides::without_managed_config_for_tests() + }; + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::path("/api/codex/config/bundle")) + .respond_with(wiremock::ResponseTemplate::new(/*s*/ 200).set_body_string("{}")) + .mount(&server) + .await; + std::fs::write( + &profile_path, + format!("chatgpt_base_url = '{}'", server.uri()), + ) + .expect("configure local ChatGPT fixture"); let mut config = test_config(&temp_home).await; config.model = Some("gpt-5.1-codex-max".to_string()); config.model_provider_id = "openai".to_string(); @@ -384,9 +400,16 @@ async fn status_snapshot_shows_chatgpt_plan_without_email() { write_models_cache(temp_home.path()) .await .expect("write models cache"); - let mut app_server = crate::start_embedded_app_server_for_picker(&config) - .await - .expect("start embedded app server"); + let mut app_server = crate::start_app_server_for_picker( + &config, + &crate::AppServerTarget::Embedded, + Vec::new(), + loader_overrides, + /*state_db*/ None, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ) + .await + .expect("start embedded app server"); let bootstrap = app_server .bootstrap(&config) .await diff --git a/codex-rs/workload-identity/src/exchange.rs b/codex-rs/workload-identity/src/exchange.rs index 018a1e12c5..d91078650c 100644 --- a/codex-rs/workload-identity/src/exchange.rs +++ b/codex-rs/workload-identity/src/exchange.rs @@ -26,6 +26,7 @@ const TRANSIENT_FAILURE_RETRY_DELAY: Duration = Duration::from_secs(30); /// Exchanges assertions and retains only the current short-lived access token in memory. pub struct WorkloadIdentityExchange { + network_policy: codex_http_client::NetworkPolicy, client: HttpClient, completed_attempts: AtomicU64, config: WorkloadIdentityConfig, @@ -61,6 +62,7 @@ impl WorkloadIdentityExchange { .map_err(|_| WorkloadIdentityError::HttpClientConfiguration) }?; Ok(Self { + network_policy: http_client_factory.network_policy().clone(), client, completed_attempts: AtomicU64::new(0), config, @@ -170,54 +172,59 @@ impl WorkloadIdentityExchange { } async fn exchange_uncached(&self) -> Result { - let assertion = read_assertion(&self.config.assertion_file).await?; - let body = { - let mut serializer = url::form_urlencoded::Serializer::new(String::new()); - serializer - .append_pair("grant_type", JWT_BEARER_GRANT_TYPE) - .append_pair("assertion", &assertion) - .append_pair("federation_rule_id", &self.config.federation_rule_id); - if let Some(context) = &self.config.workload_identity_context { - serializer.append_pair("workload_identity_context", context); - } - serializer.finish() - }; - let response = self - .client - .post(self.token_url.as_str()) - .header("content-type", "application/x-www-form-urlencoded") - .body(body) - .timeout(REQUEST_TIMEOUT) - .send() - .await - .map_err(|_| WorkloadIdentityError::ExchangeUnavailable)?; - if !response.status().is_success() { - return Err(WorkloadIdentityError::ExchangeRejected( - response.status().as_u16(), - )); - } - if response - .content_length() - .is_some_and(|length| length > MAX_RESPONSE_BYTES as u64) - { - return Err(WorkloadIdentityError::InvalidExchangeResponse); - } + let permit = self.network_policy.acquire(&self.token_url)?; + permit + .run(async { + let assertion = read_assertion(&self.config.assertion_file).await?; + let body = { + let mut serializer = url::form_urlencoded::Serializer::new(String::new()); + serializer + .append_pair("grant_type", JWT_BEARER_GRANT_TYPE) + .append_pair("assertion", &assertion) + .append_pair("federation_rule_id", &self.config.federation_rule_id); + if let Some(context) = &self.config.workload_identity_context { + serializer.append_pair("workload_identity_context", context); + } + serializer.finish() + }; + let response = self + .client + .post(self.token_url.as_str()) + .header("content-type", "application/x-www-form-urlencoded") + .body(body) + .timeout(REQUEST_TIMEOUT) + .send() + .await + .map_err(|_| WorkloadIdentityError::ExchangeUnavailable)?; + if !response.status().is_success() { + return Err(WorkloadIdentityError::ExchangeRejected( + response.status().as_u16(), + )); + } + if response + .content_length() + .is_some_and(|length| length > MAX_RESPONSE_BYTES as u64) + { + return Err(WorkloadIdentityError::InvalidExchangeResponse); + } - let mut response = response; - let mut bytes = Vec::new(); - while let Some(chunk) = response - .chunk() - .await - .map_err(|_| WorkloadIdentityError::ExchangeUnavailable)? - { - if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(bytes.len()) { - return Err(WorkloadIdentityError::InvalidExchangeResponse); - } - bytes.extend_from_slice(&chunk); - } - serde_json::from_slice::(&bytes) - .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)? - .into_token() + let mut response = response; + let mut bytes = Vec::new(); + while let Some(chunk) = response + .chunk() + .await + .map_err(|_| WorkloadIdentityError::ExchangeUnavailable)? + { + if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(bytes.len()) { + return Err(WorkloadIdentityError::InvalidExchangeResponse); + } + bytes.extend_from_slice(&chunk); + } + serde_json::from_slice::(&bytes) + .map_err(|_| WorkloadIdentityError::InvalidExchangeResponse)? + .into_token() + }) + .await? } } diff --git a/codex-rs/workload-identity/src/lib.rs b/codex-rs/workload-identity/src/lib.rs index 252d8e8d61..e693772b19 100644 --- a/codex-rs/workload-identity/src/lib.rs +++ b/codex-rs/workload-identity/src/lib.rs @@ -39,6 +39,8 @@ impl WorkloadIdentityConfig { #[derive(Clone, Debug, Error)] pub enum WorkloadIdentityError { + #[error(transparent)] + Policy(#[from] codex_http_client::NetworkPolicyDenied), #[error("the workload identity federation rule ID must not be empty")] InvalidFederationRuleId, #[error("the workload identity assertion file path must be absolute")] @@ -77,7 +79,8 @@ impl WorkloadIdentityError { | std::io::ErrorKind::WouldBlock ), Self::ExchangeUnavailable | Self::ExchangeRejected(408 | 429 | 500..=599) => true, - Self::InvalidFederationRuleId + Self::Policy(_) + | Self::InvalidFederationRuleId | Self::AssertionFileMustBeAbsolute | Self::InvalidAssertion | Self::AssertionTooLarge diff --git a/codex-rs/workload-identity/src/workload_identity_tests.rs b/codex-rs/workload-identity/src/workload_identity_tests.rs index 8ebd5ce481..59ca8f75a3 100644 --- a/codex-rs/workload-identity/src/workload_identity_tests.rs +++ b/codex-rs/workload-identity/src/workload_identity_tests.rs @@ -368,3 +368,95 @@ async fn exchange_rejects_oversized_assertions_and_incomplete_responses() { Err(WorkloadIdentityError::InvalidExchangeResponse) )); } + +#[tokio::test] +async fn managed_exchange_denies_before_reading_assertion_or_sending() { + let server = MockServer::start().await; + let (_home, assertion_path) = assertion_file("assertion"); + std::fs::remove_file(&assertion_path).unwrap(); + let controller = codex_http_client::NetworkPolicyController::default(); + let exchange = WorkloadIdentityExchange::new( + WorkloadIdentityConfig::new( + "rule".into(), + assertion_path, + /*workload_identity_context*/ None, + ) + .unwrap(), + Url::parse(&format!("{}/oauth/token", server.uri())).unwrap(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(controller.policy()), + ) + .unwrap(); + assert!(matches!( + exchange.resolve().await, + Err(WorkloadIdentityError::Policy( + codex_http_client::NetworkPolicyDenied::Unavailable + )) + )); + assert!(server.received_requests().await.unwrap().is_empty()); + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + assert!(matches!( + exchange.resolve().await, + Err(WorkloadIdentityError::Policy( + codex_http_client::NetworkPolicyDenied::Destination + )) + )); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn managed_exchange_revokes_an_active_token_response() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + success("token", /*expires_in*/ 600).set_delay(Duration::from_secs(/*secs*/ 30)), + ) + .expect(/*r*/ 1) + .mount(&server) + .await; + let (_home, assertion_path) = assertion_file("assertion"); + let controller = codex_http_client::NetworkPolicyController::default(); + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); + let exchange = WorkloadIdentityExchange::new( + WorkloadIdentityConfig::new( + "rule".into(), + assertion_path, + /*workload_identity_context*/ None, + ) + .unwrap(), + Url::parse(&format!("{}/oauth/token", server.uri())).unwrap(), + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(controller.policy()), + ) + .unwrap(); + let revoke = async { + while server.received_requests().await.unwrap().is_empty() { + tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await; + } + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + }; + let (result, ()) = tokio::time::timeout(Duration::from_secs(/*secs*/ 5), async { + tokio::join!(exchange.resolve(), revoke) + }) + .await + .unwrap(); + assert!(matches!( + result, + Err(WorkloadIdentityError::Policy( + codex_http_client::NetworkPolicyDenied::Revoked + )) + )); +}