diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d11ecb2830..448cba9d98 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1980,6 +1980,7 @@ version = "0.0.0" dependencies = [ "codex-app-server-protocol", "codex-git-utils", + "codex-http-client", "codex-login", "codex-model-provider", "codex-plugin", @@ -2375,12 +2376,23 @@ dependencies = [ "aws-config", "aws-credential-types", "aws-sigv4", + "aws-smithy-runtime-api", + "aws-smithy-types", "aws-types", "bytes", + "codex-http-client", + "codex-utils-rustls-provider", "http 1.4.0", + "http-body 1.0.1", + "http-body-util", "pretty_assertions", + "rcgen", + "rustls", + "tempfile", "thiserror 2.0.18", "tokio", + "tokio-rustls", + "url", ] [[package]] @@ -4082,8 +4094,10 @@ dependencies = [ name = "codex-otel" version = "0.0.0" dependencies = [ + "bytes", "chrono", "codex-api", + "codex-http-client", "codex-protocol", "codex-utils-absolute-path", "codex-utils-string", @@ -4093,6 +4107,7 @@ dependencies = [ "libc", "opentelemetry", "opentelemetry-appender-tracing", + "opentelemetry-http", "opentelemetry-otlp", "opentelemetry-semantic-conventions", "opentelemetry_sdk", @@ -5155,6 +5170,7 @@ dependencies = [ "anyhow", "base64 0.22.1", "chrono", + "codex-http-client", "codex-otel", "codex-protocol", "codex-utils-absolute-path", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 4512261e77..1a43dd6f77 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -327,6 +327,8 @@ async-stream = "0.3.6" aws-config = "1" aws-credential-types = "1" aws-sigv4 = "1" +aws-smithy-runtime-api = { version = "1", features = ["client", "http-1x"] } +aws-smithy-types = { version = "1", features = ["http-body-1-x"] } aws-types = "1" axum = { version = "0.8", default-features = false } base64 = "0.22.1" diff --git a/codex-rs/analytics/Cargo.toml b/codex-rs/analytics/Cargo.toml index 8546427469..26d84441d9 100644 --- a/codex-rs/analytics/Cargo.toml +++ b/codex-rs/analytics/Cargo.toml @@ -13,6 +13,7 @@ path = "src/lib.rs" workspace = true [dependencies] +codex-http-client = { workspace = true } codex-app-server-protocol = { workspace = true } codex-git-utils = { workspace = true } codex-login = { workspace = true } diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 373e4e2bea..a00c76da30 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -61,7 +61,6 @@ use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::item_event_to_server_notification; use codex_login::AuthManager; use codex_login::CodexAuth; -use codex_login::default_client::create_client; use codex_plugin::PluginId; use codex_plugin::PluginTelemetryMetadata; use codex_protocol::ThreadId; @@ -867,7 +866,8 @@ async fn send_track_events( return; } - let Some(auth) = auth_manager.auth().await else { + let Some((auth, http_client_factory)) = auth_manager.auth_with_http_client_factory().await + else { return; }; if auth.is_api_key_auth() { @@ -880,7 +880,7 @@ async fn send_track_events( } for events in track_event_request_batches(events) { - send_track_events_request(&auth, destination, events).await; + send_track_events_request(&auth, destination, events, &http_client_factory).await; } } @@ -911,6 +911,7 @@ async fn send_track_events_request( auth: &CodexAuth, destination: &AnalyticsEventsDestination, events: Vec, + http_client_factory: &codex_http_client::HttpClientFactory, ) { if events.is_empty() { return; @@ -928,7 +929,21 @@ async fn send_track_events_request( #[cfg(debug_assertions)] AnalyticsEventsDestination::CaptureFile { .. } => return, }; - let response = create_client() + let client = match codex_login::default_client::create_client_for_route_async( + http_client_factory.clone(), + url.clone(), + codex_http_client::ClientRouteClass::Api, + codex_login::default_client::ClientRedirectPolicy::Default, + ) + .await + { + Ok(client) => client, + Err(error) => { + tracing::warn!(%error, "failed to build events client"); + return; + } + }; + let response = client .post(url) .timeout(ANALYTICS_EVENTS_TIMEOUT) .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()) diff --git a/codex-rs/analytics/src/client_tests.rs b/codex-rs/analytics/src/client_tests.rs index b36829758e..9112070c9c 100644 --- a/codex-rs/analytics/src/client_tests.rs +++ b/codex-rs/analytics/src/client_tests.rs @@ -88,6 +88,8 @@ use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::TurnStatus as AppServerTurnStatus; use codex_app_server_protocol::TurnSteerParams; use codex_app_server_protocol::TurnSteerResponse; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; #[cfg(debug_assertions)] use codex_login::AuthManager; use codex_utils_absolute_path::test_support::PathBufExt; @@ -367,7 +369,8 @@ async fn capture_file_writes_exact_serialized_request() { let expected_event = serde_json::to_value(&event).expect("serialize expected event"); let auth = codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(); - send_track_events_request(&auth, &destination, vec![event]).await; + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + send_track_events_request(&auth, &destination, vec![event], &factory).await; let contents = fs::read_to_string(&capture_path).expect("read capture file"); let lines = contents.lines().collect::>(); @@ -394,7 +397,8 @@ async fn capture_file_writes_final_batches_as_separate_lines() { ]; for batch in track_event_request_batches(events) { - send_track_events_request(&auth, &destination, batch).await; + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + send_track_events_request(&auth, &destination, batch, &factory).await; } let contents = fs::read_to_string(&capture_path).expect("read capture file"); diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 8a572be0ab..ffb03c0ec5 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -201,6 +201,12 @@ Bedrock login return an error without changing configuration or saved credential exporter configuration before selecting another credential source. `aws.credential_export` and `aws.profile` cannot be configured together. +Application network restrictions apply to each AWS credential and region HTTP request and to +the Bedrock destination. Static access keys with an explicit region need no credential discovery. +AWS profile `credential_process` commands are run by the AWS SDK; their network traffic is outside +the application's HTTP policy. Configured credential exporters and AWS reauthentication commands +require unrestricted application policy; policy revocation cancels their active work. + ## Stored thread attachments - `thread/attachment/add` — add a durable resource reference to a stored thread without loading it. Repeated writes with the same attachment type and identity key return the existing attachment. diff --git a/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs b/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs index 9e7cb31856..a3fe8f9108 100644 --- a/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs +++ b/codex-rs/app-server/src/request_processors/account_processor/bedrock_setup.rs @@ -98,13 +98,17 @@ impl AccountRequestProcessor { if profile.is_empty() { return Err(invalid_request("AWS profile name must not be empty.")); } - codex_aws_auth::validate_aws_profile(profile, region) - .await - .map_err(|err| { - invalid_request(format!( - "failed to load credentials for AWS profile `{profile}`: {err}" - )) - })?; + codex_aws_auth::validate_aws_profile( + profile, + region, + self.config.http_client_factory(), + ) + .await + .map_err(|err| { + invalid_request(format!( + "failed to load credentials for AWS profile `{profile}`: {err}" + )) + })?; Some(profile.to_string()) } BedrockSetupParams::Environment { .. } => { diff --git a/codex-rs/app-server/tests/suite/v2/otel.rs b/codex-rs/app-server/tests/suite/v2/otel.rs index 197af70f4d..1f5836dde6 100644 --- a/codex-rs/app-server/tests/suite/v2/otel.rs +++ b/codex-rs/app-server/tests/suite/v2/otel.rs @@ -60,6 +60,26 @@ async fn account_switch_reloads_telemetry_collectors_and_preserves_trace_context .start_thread(ThreadStartParams::default()) .await?; + let contains = |requests: &[wiremock::Request], path: &str, value: &str| { + requests.iter().any(|request| { + request.url.path() == path && String::from_utf8_lossy(&request.body).contains(value) + }) + }; + timeout(TEST_TIMEOUT, async { + let mut poll = tokio::time::interval(Duration::from_millis(/*millis*/ 25)); + loop { + poll.tick().await; + let requests = collector.received_requests().await.unwrap_or_default(); + if contains(&requests, "/initial/logs", INITIAL_EMAIL) + && contains(&requests, "/initial/traces", PARENT_TRACE_ID) + { + break; + } + } + }) + .await + .context("initial account telemetry must be exported before its permission is revoked")?; + let next_endpoint = format!("{}/next", collector.uri()); write_otel_config(codex_home.path(), &next_endpoint)?; let access_token = encode_id_token( @@ -84,9 +104,21 @@ async fn account_switch_reloads_telemetry_collectors_and_preserves_trace_context app_server.wait_for_json_log_event("codex.app_server.otel_reloaded"), ) .await??; - app_server - .start_thread(ThreadStartParams::default()) - .await?; + timeout(TEST_TIMEOUT, async { + let mut poll = tokio::time::interval(Duration::from_secs(/*secs*/ 1)); + loop { + poll.tick().await; + app_server + .start_thread(ThreadStartParams::default()) + .await?; + let requests = collector.received_requests().await.unwrap_or_default(); + if contains(&requests, "/next/logs", NEXT_EMAIL) { + break Ok::<(), anyhow::Error>(()); + } + } + }) + .await + .context("new account telemetry must reach the replacement collector")??; let status = timeout(TEST_TIMEOUT, app_server.shutdown_gracefully()).await??; assert!(status.success(), "app-server did not shut down cleanly"); diff --git a/codex-rs/aws-auth/Cargo.toml b/codex-rs/aws-auth/Cargo.toml index 6bb5a69ae9..f9b5e7139a 100644 --- a/codex-rs/aws-auth/Cargo.toml +++ b/codex-rs/aws-auth/Cargo.toml @@ -13,14 +13,25 @@ path = "src/lib.rs" workspace = true [dependencies] +codex-http-client = { workspace = true } aws-config = { workspace = true, features = ["credentials-login"] } aws-credential-types = { workspace = true } aws-sigv4 = { workspace = true } +aws-smithy-runtime-api = { workspace = true } +aws-smithy-types = { workspace = true } aws-types = { workspace = true } bytes = { workspace = true } http = { workspace = true } +http-body = "1.0.1" +http-body-util = "0.1.3" +url = { workspace = true } thiserror = { workspace = true } [dev-dependencies] pretty_assertions = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } +rcgen = { workspace = true } +rustls = { workspace = true } +tempfile = { workspace = true } +tokio-rustls = { workspace = true } +codex-utils-rustls-provider = { workspace = true } +tokio = { workspace = true, features = ["io-util", "macros", "net", "rt-multi-thread", "time"] } diff --git a/codex-rs/aws-auth/src/config.rs b/codex-rs/aws-auth/src/config.rs index 3d62832e0e..7b2f0e5ef6 100644 --- a/codex-rs/aws-auth/src/config.rs +++ b/codex-rs/aws-auth/src/config.rs @@ -1,24 +1,29 @@ +//! Loads AWS SDK configuration with the application's HTTP client. + use aws_config::BehaviorVersion; use aws_config::SdkConfig; use aws_credential_types::provider::SharedCredentialsProvider; use aws_types::region::Region; +use codex_http_client::HttpClientFactory; use crate::AwsAuthConfig; use crate::AwsAuthError; -pub(crate) async fn load_sdk_config(config: &AwsAuthConfig) -> Result { +pub(crate) async fn load_sdk_config( + config: &AwsAuthConfig, + factory: &HttpClientFactory, +) -> Result { if config.service.trim().is_empty() { return Err(AwsAuthError::EmptyService); } - - let mut loader = aws_config::defaults(BehaviorVersion::latest()); + let mut loader = aws_config::defaults(BehaviorVersion::latest()) + .http_client(crate::transport::http_client(factory.clone())); if let Some(profile) = config.profile.as_ref() { loader = loader.profile_name(profile); } if let Some(region) = config.region.as_ref() { loader = loader.region(Region::new(region.clone())); } - Ok(loader.load().await) } diff --git a/codex-rs/aws-auth/src/discovery.rs b/codex-rs/aws-auth/src/discovery.rs index 10dbff017c..d22752d8fb 100644 --- a/codex-rs/aws-auth/src/discovery.rs +++ b/codex-rs/aws-auth/src/discovery.rs @@ -1,4 +1,3 @@ -use aws_config::default_provider::region::DefaultRegionChain; use aws_config::profile::ProfileFileCredentialsProvider; use aws_config::provider_config::ProviderConfig; use aws_credential_types::provider::ProvideCredentials; @@ -44,32 +43,27 @@ pub async fn discover_aws_profiles() -> Result, AwsAuthError> { } /// Resolves credentials only from the selected AWS profile. -pub async fn validate_aws_profile(profile: &str, region: &str) -> Result<(), AwsAuthError> { - profile_credentials_provider(profile, Some(region)) - .await +pub async fn validate_aws_profile( + profile: &str, + region: &str, + factory: codex_http_client::HttpClientFactory, +) -> Result<(), AwsAuthError> { + profile_credentials_provider(profile, region, factory) .provide_credentials() .await?; Ok(()) } -pub(crate) async fn profile_credentials_provider( +pub(crate) fn profile_credentials_provider( profile: &str, - region: Option<&str>, + region: &str, + factory: codex_http_client::HttpClientFactory, ) -> ProfileFileCredentialsProvider { - let region = match region { - Some(region) => Some(Region::new(region.to_string())), - None => { - DefaultRegionChain::builder() - .profile_name(profile) - .build() - .region() - .await - } - }; - let provider_config = ProviderConfig::without_region().with_region(region); - + let config = ProviderConfig::without_region() + .with_region(Some(Region::new(region.to_string()))) + .with_http_client(crate::transport::http_client(factory)); ProfileFileCredentialsProvider::builder() - .configure(&provider_config) + .configure(&config) .profile_name(profile) .build() } diff --git a/codex-rs/aws-auth/src/lib.rs b/codex-rs/aws-auth/src/lib.rs index f79ef91b60..8f5e2d3eab 100644 --- a/codex-rs/aws-auth/src/lib.rs +++ b/codex-rs/aws-auth/src/lib.rs @@ -1,6 +1,7 @@ mod config; mod discovery; mod signing; +mod transport; use std::sync::Arc; use std::time::SystemTime; @@ -49,6 +50,8 @@ impl std::fmt::Debug for AwsAccessKeys { /// /// Implementations should return current credentials for every call so request signing can /// observe credential refreshes. Errors must not contain credentials or command output. +/// These caller-supplied providers require unrestricted network policy because their I/O +/// runs outside the shared AWS HTTP transport. pub trait AwsCredentialsProvider: std::fmt::Debug + Send + Sync { fn credentials( &self, @@ -56,7 +59,7 @@ pub trait AwsCredentialsProvider: std::fmt::Debug + Send + Sync { } #[derive(Debug)] -struct AwsCredentialsProviderAdapter

(Arc

); +struct AwsCredentialsProviderAdapter

(Arc

, codex_http_client::NetworkPolicy); #[derive(Debug, Error)] #[error("{0}")] @@ -70,7 +73,12 @@ impl ProvideCredentials for AwsCredentialsProviderAda Self: 'a, { aws_credential_types::provider::future::ProvideCredentials::new(async move { - let access_keys = self.0.credentials().await.map_err(|error| { + let permit = self.1.acquire_for_unsupported_sdk().map_err( + aws_credential_types::provider::error::CredentialsError::invalid_configuration, + )?; + let access_keys = permit.run(self.0.credentials()).await + .map_err(aws_credential_types::provider::error::CredentialsError::invalid_configuration)? + .map_err(|error| { let error = ProvidedCredentialsError(error); match error.0.kind() { std::io::ErrorKind::InvalidData @@ -113,6 +121,8 @@ pub struct AwsSignedRequest { /// Errors returned by credential loading or SigV4 signing. #[derive(Debug, Error)] pub enum AwsAuthError { + #[error(transparent)] + Policy(#[from] codex_http_client::NetworkPolicyDenied), #[error("AWS service name must not be empty")] EmptyService, #[error("AWS profile must be configured")] @@ -125,6 +135,8 @@ pub enum AwsAuthError { ProfileLoad(#[from] aws_config::profile::ProfileFileLoadError), #[error("failed to load AWS credentials: {0}")] Credentials(#[from] aws_credential_types::provider::error::CredentialsError), + #[error("request URL is not a valid URL: {0}")] + InvalidUrl(#[source] url::ParseError), #[error("request URL is not a valid URI: {0}")] InvalidUri(#[source] http::uri::InvalidUri), #[error("failed to construct HTTP request for signing: {0}")] @@ -142,6 +154,7 @@ pub enum AwsAuthError { /// Loaded AWS auth context that can sign outbound HTTP requests. #[derive(Clone)] pub struct AwsAuthContext { + network_policy: codex_http_client::NetworkPolicy, credentials_provider: SharedCredentialsProvider, region: String, service: String, @@ -157,14 +170,15 @@ impl std::fmt::Debug for AwsAuthContext { } impl AwsAuthContext { - pub async fn load(config: AwsAuthConfig) -> Result { - let sdk_config = config::load_sdk_config(&config).await?; - let credentials_provider = config::credentials_provider(&sdk_config)?; - let region = config::resolved_region(&sdk_config)?; - + pub async fn load( + config: AwsAuthConfig, + factory: codex_http_client::HttpClientFactory, + ) -> Result { + let sdk_config = config::load_sdk_config(&config, &factory).await?; Ok(Self { - credentials_provider, - region, + network_policy: factory.network_policy().clone(), + credentials_provider: config::credentials_provider(&sdk_config)?, + region: config::resolved_region(&sdk_config)?, service: config.service.trim().to_string(), }) } @@ -172,39 +186,52 @@ impl AwsAuthContext { pub async fn load_with_access_keys( config: AwsAuthConfig, access_keys: AwsAccessKeys, + factory: codex_http_client::HttpClientFactory, ) -> Result { - let mut context = Self::load(config).await?; - context.credentials_provider = - SharedCredentialsProvider::new(aws_credential_types::Credentials::new( - access_keys.access_key_id, - access_keys.secret_access_key, - access_keys.session_token, - /*expires_after*/ None, - "codex-managed-bedrock-access-keys", - )); - Ok(context) + if config.service.trim().is_empty() { + return Err(AwsAuthError::EmptyService); + } + let region = match config.region.as_ref() { + Some(region) => region.clone(), + None => config::resolved_region(&config::load_sdk_config(&config, &factory).await?)?, + }; + Ok(Self { + network_policy: factory.network_policy().clone(), + credentials_provider: SharedCredentialsProvider::new( + aws_credential_types::Credentials::new( + access_keys.access_key_id, + access_keys.secret_access_key, + access_keys.session_token, + /*expires_after*/ None, + "codex-managed-bedrock-access-keys", + ), + ), + region, + service: config.service.trim().to_string(), + }) } pub async fn load_with_credentials_provider( config: AwsAuthConfig, provider: Arc, + factory: codex_http_client::HttpClientFactory, ) -> Result { - let mut context = Self::load(config).await?; - context.credentials_provider = - SharedCredentialsProvider::new(AwsCredentialsProviderAdapter(provider)); + let mut context = Self::load(config, factory).await?; + context.credentials_provider = SharedCredentialsProvider::new( + AwsCredentialsProviderAdapter(provider, context.network_policy.clone()), + ); Ok(context) } - pub async fn load_profile(config: AwsAuthConfig) -> Result { - let profile = config - .profile - .as_deref() - .ok_or(AwsAuthError::MissingProfile)?; - let credentials_provider = SharedCredentialsProvider::new( - discovery::profile_credentials_provider(profile, config.region.as_deref()).await, + pub async fn load_profile( + config: AwsAuthConfig, + factory: codex_http_client::HttpClientFactory, + ) -> Result { + let profile = config.profile.clone().ok_or(AwsAuthError::MissingProfile)?; + let mut context = Self::load(config, factory.clone()).await?; + context.credentials_provider = SharedCredentialsProvider::new( + discovery::profile_credentials_provider(&profile, &context.region, factory), ); - let mut context = Self::load(config).await?; - context.credentials_provider = credentials_provider; Ok(context) } @@ -225,7 +252,28 @@ impl AwsAuthContext { request: AwsRequestToSign, time: SystemTime, ) -> Result { - let credentials = self.credentials_provider.provide_credentials().await?; + let url = request + .url + .parse::() + .map_err(AwsAuthError::InvalidUrl)?; + let permit = self.network_policy.acquire(&url)?; + let credentials = permit + .run(self.credentials_provider.provide_credentials()) + .await? + .map_err(|error| { + let mut source: &(dyn std::error::Error + 'static) = &error; + loop { + if let Some(policy) = + source.downcast_ref::() + { + return AwsAuthError::Policy(*policy); + } + match source.source() { + Some(next) => source = next, + None => return AwsAuthError::Credentials(error), + } + } + })?; signing::sign_request(&credentials, &self.region, &self.service, request, time) } } @@ -249,11 +297,13 @@ impl AwsAuthError { aws_credential_types::provider::error::CredentialsError::ProviderTimedOut(_) | aws_credential_types::provider::error::CredentialsError::ProviderError(_) ), - AwsAuthError::EmptyService + AwsAuthError::Policy(_) + | AwsAuthError::EmptyService | AwsAuthError::MissingProfile | AwsAuthError::MissingCredentialsProvider | AwsAuthError::MissingRegion | AwsAuthError::ProfileLoad(_) + | AwsAuthError::InvalidUrl(_) | AwsAuthError::InvalidUri(_) | AwsAuthError::BuildHttpRequest(_) | AwsAuthError::InvalidHeaderValue(_) @@ -277,6 +327,7 @@ mod tests { fn test_context(session_token: Option<&str>) -> AwsAuthContext { AwsAuthContext { + network_policy: Default::default(), credentials_provider: SharedCredentialsProvider::new(Credentials::new( "AKIDEXAMPLE", "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", @@ -289,7 +340,7 @@ mod tests { } } - fn test_request() -> AwsRequestToSign { + pub(super) fn test_request() -> AwsRequestToSign { let mut headers = HeaderMap::new(); headers.insert( http::header::CONTENT_TYPE, @@ -346,15 +397,17 @@ mod tests { } } - let credentials = - AwsCredentialsProviderAdapter(Arc::new(TestCredentialsProvider(Ok(AwsAccessKeys { + let credentials = AwsCredentialsProviderAdapter( + Arc::new(TestCredentialsProvider(Ok(AwsAccessKeys { access_key_id: "access-key-id".to_string(), secret_access_key: "secret-access-key".to_string(), session_token: Some("session-token".to_string()), - })))) - .provide_credentials() - .await - .expect("exported credentials should be available"); + }))), + Default::default(), + ) + .provide_credentials() + .await + .expect("exported credentials should be available"); assert_eq!( credentials, @@ -375,10 +428,13 @@ mod tests { (std::io::ErrorKind::NotFound, false), (std::io::ErrorKind::PermissionDenied, false), ] { - let error = AwsCredentialsProviderAdapter(Arc::new(TestCredentialsProvider(Err(kind)))) - .provide_credentials() - .await - .expect_err("credential export failure should be propagated"); + let error = AwsCredentialsProviderAdapter( + Arc::new(TestCredentialsProvider(Err(kind))), + Default::default(), + ) + .provide_credentials() + .await + .expect_err("credential export failure should be propagated"); let error = AwsAuthError::Credentials(error); assert_eq!( ( @@ -442,24 +498,40 @@ mod tests { #[tokio::test] async fn load_rejects_invalid_configuration() { - let err = AwsAuthContext::load(AwsAuthConfig { + let config = AwsAuthConfig { profile: None, region: None, service: " ".to_string(), - }) + }; + let err = AwsAuthContext::load( + config, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) .await .expect_err("empty service should be rejected"); assert_eq!(err.to_string(), "AWS service name must not be empty"); - let err = AwsAuthContext::load_profile(AwsAuthConfig { + let config = AwsAuthConfig { profile: None, region: Some("us-east-1".to_string()), service: "bedrock".to_string(), - }) + }; + let err = AwsAuthContext::load_profile( + config, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) .await .expect_err("profile auth should require a configured profile"); assert_eq!(err.to_string(), "AWS profile must be configured"); } } + +#[cfg(test)] +#[path = "network_policy_tests.rs"] +mod network_policy_tests; diff --git a/codex-rs/aws-auth/src/network_policy_tests.rs b/codex-rs/aws-auth/src/network_policy_tests.rs new file mode 100644 index 0000000000..adf4978fd7 --- /dev/null +++ b/codex-rs/aws-auth/src/network_policy_tests.rs @@ -0,0 +1,469 @@ +//! Exercises policy on AWS SDK HTTP and the SDK's credential selection. + +use super::AwsAuthError::Policy; +use super::*; +use aws_config::provider_config::ProviderConfig; +use codex_http_client::DestinationPolicy; +use codex_http_client::NetworkPolicyDenied::Revoked; +use pretty_assertions::assert_eq; +use std::error::Error; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::time::timeout; + +#[tokio::test] +async fn real_imds_credentials_stop_after_policy_revocation() -> Result<(), Box> { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + let factory = codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ) + .with_network_policy(policy.clone().for_current_account()); + let http_client = crate::transport::http_client(factory); + let imds_client = aws_config::imds::Client::builder() + .configure(&ProviderConfig::default().with_http_client(http_client)) + .endpoint(format!("http://{}", listener.local_addr()?)) + .unwrap() + .build(); + let credentials_provider = aws_config::imds::credentials::ImdsCredentialsProvider::builder() + .profile("test-role") + .imds_client(imds_client.clone()) + .build(); + let context = AwsAuthContext { + network_policy: policy.clone().for_current_account(), + credentials_provider: SharedCredentialsProvider::new(credentials_provider), + region: "us-east-1".into(), + service: "bedrock".into(), + }; + let request = super::tests::test_request(); + let server = async { + let mut credentials_sent = false; + loop { + let (mut socket, _) = listener.accept().await?; + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(socket.read_u8().await?); + } + let body = if headers.starts_with(b"PUT /latest/api/token ") { + "test-token" + } else if credentials_sent { + return Ok::<_, std::io::Error>(socket); + } else { + credentials_sent = true; + r#"{"Code":"Success","AccessKeyId":"test-key","SecretAccessKey":"test-secret","Token":"test-session","Expiration":"2099-01-01T00:00:00Z"}"# + }; + let response = format!( + "HTTP/1.1 200 OK\r\nx-aws-ec2-metadata-token-ttl-seconds: 21600\r\nConnection: close\r\n\r\n{body}" + ); + socket.write_all(response.as_bytes()).await?; + } + }; + controller.publish(policy.revision(), DestinationPolicy::Unrestricted); + let signing = async { + context.sign(request.clone()).await.unwrap(); + context.sign(request.clone()).await + }; + tokio::pin!(signing); + let timeout_duration = std::time::Duration::from_secs(/*secs*/ 5); + let mut socket = tokio::select! { + result = &mut signing => panic!("signing finished before revocation: {result:?}"), + socket = timeout(timeout_duration, server) => socket??, + }; + let restricted = DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }; + controller.publish(policy.revision(), restricted); + let revoked = timeout(timeout_duration, signing).await?; + assert!(matches!(revoked, Err(Policy(Revoked)))); + assert_eq!(timeout(timeout_duration, socket.read(&mut [0])).await??, 0); + // The model remains allowed while metadata traffic is denied before connecting. + controller.publish( + policy.revision(), + DestinationPolicy::Restricted { + allowed_hosts: ["bedrock-runtime.us-east-1.amazonaws.com".to_string()].into(), + }, + ); + context.sign(request.clone()).await?; + let fresh = AwsAuthContext { + credentials_provider: SharedCredentialsProvider::new( + aws_config::imds::credentials::ImdsCredentialsProvider::builder() + .profile("test-role") + .imds_client(imds_client) + .build(), + ), + ..context.clone() + }; + let denied = fresh.sign(request.clone()).await; + assert!( + matches!( + &denied, + Err(Policy(codex_http_client::NetworkPolicyDenied::Destination)) + ), + "{denied:?}" + ); + assert!( + timeout( + std::time::Duration::from_millis(/*millis*/ 100), + listener.accept() + ) + .await + .is_err() + ); + Ok(()) +} + +#[tokio::test] +async fn allowlisted_static_keys_send_sigv4_and_denied_destination_sends_nothing() +-> Result<(), Box> { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let certificate = rcgen::generate_simple_self_signed(vec!["127.0.0.1".to_string()])?; + let trust = codex_http_client::HttpClientTlsConfig::default() + .with_root_certificate_pem(certificate.cert.pem().as_bytes())?; + let tls = tokio_rustls::TlsAcceptor::from(Arc::new( + rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert( + vec![certificate.cert.der().clone()], + certificate.signing_key.into(), + )?, + )); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let url = format!("https://{}/v1/responses", listener.local_addr()?); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish( + policy.revision(), + DestinationPolicy::Restricted { + allowed_hosts: ["127.0.0.1".to_string()].into(), + }, + ); + let factory = codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ) + .with_network_policy(policy.for_current_account()); + let context = AwsAuthContext::load_with_access_keys( + AwsAuthConfig { + profile: Some("unused-profile".into()), + region: Some("us-east-1".into()), + service: "bedrock".into(), + }, + AwsAccessKeys { + access_key_id: "static-key".into(), + secret_access_key: "static-secret".into(), + session_token: None, + }, + factory.clone(), + ) + .await?; + let request = AwsRequestToSign { + url: url.clone(), + ..super::tests::test_request() + }; + let signed = context.sign(request.clone()).await?; + let client = codex_http_client::HttpClientBuilder::new().build_with_tls( + &factory, + codex_http_client::ClientRouteClass::Api, + trust, + ); + let server = async { + let (socket, _) = listener.accept().await?; + let mut socket = tls.accept(socket).await?; + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(socket.read_u8().await?); + } + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}") + .await?; + Ok::<_, Box>(String::from_utf8(headers)?) + }; + let (response, headers) = tokio::join!( + client + .post(&signed.url) + .headers(signed.headers) + .body(request.body.clone()) + .send(), + server + ); + assert!(response?.status().is_success()); + assert!(headers?.contains("Credential=static-key/")); + controller.publish( + controller.policy().revision(), + DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + assert!(matches!( + context.sign(request).await, + Err(Policy(codex_http_client::NetworkPolicyDenied::Destination)) + )); + assert!( + timeout( + std::time::Duration::from_millis(/*millis*/ 100), + listener.accept() + ) + .await + .is_err() + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sdk_credentials_preserve_profile_processes_and_environment_precedence() +-> Result<(), Box> { + for case in [ + "restricted-static", + "restricted-process", + "unrestricted-process", + "unrestricted-offset-process", + "restricted-role-process", + "unrestricted-role-process", + "env-precedence", + "profile-precedence", + ] { + let home = tempfile::tempdir()?; + let marker = home.path().join("process-started"); + let expiration = if case == "unrestricted-offset-process" { + "2099-01-01T01:30:00+01:30" + } else { + "2099-01-01T00:00:00Z" + }; + let json = format!( + r#"{{"Version":1,"AccessKeyId":"process-key","SecretAccessKey":"process-secret","Expiration":"{expiration}"}}"# + ); + std::fs::write(home.path().join("credentials-output.json"), &json)?; + let command = if cfg!(windows) { + "echo started > process-started & type credentials-output.json".to_string() + } else { + format!( + "printf started > '{}' && printf '%s' '{json}'", + marker.to_string_lossy().replace('\'', "'\\''") + ) + }; + let contents = if matches!(case, "restricted-static" | "profile-precedence") { + "[managed]\naws_access_key_id = profile-key\naws_secret_access_key = profile-secret\n" + .to_string() + } else { + format!( + "[managed]\ncredential_process = {command}\naws_access_key_id = fallback-key\naws_secret_access_key = fallback-secret\n" + ) + }; + let credentials = home.path().join("credentials"); + let config = home.path().join("config"); + std::fs::write(&credentials, contents)?; + std::fs::write( + &config, + if case.contains("role-process") { + "[profile managed] # SDK accepts section comments\nrole_arn = arn:aws:iam::123456789012:role/test\nsource_profile = managed # self source\n" + } else { + "" + }, + )?; + let role_request = if case == "unrestricted-role-process" { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}", listener.local_addr()?); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await?; + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(socket.read_u8().await?); + } + let headers = String::from_utf8(headers).map_err(std::io::Error::other)?; + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::()) + }) + .transpose() + .map_err(std::io::Error::other)? + .unwrap_or_default(); + // Drain the POST body before closing so Windows does not reset the response. + socket.read_exact(&mut vec![0; content_length]).await?; + let body = "assumed-keyassumed-secretassumed-token2099-01-01T00:00:00Z"; + socket.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await?; + Ok::<_, std::io::Error>(headers) + }); + Some((endpoint, server)) + } else { + None + }; + let mut child = std::process::Command::new(std::env::current_exe()?); + child + .current_dir(home.path()) + .args([ + "--ignored", + "--exact", + "network_policy_tests::sdk_profile_probe", + ]) + .env("CODEX_AWS_PROFILE_PROBE_CASE", case) + .env("AWS_CONFIG_FILE", config) + .env("AWS_SHARED_CREDENTIALS_FILE", credentials) + .env("AWS_EC2_METADATA_DISABLED", "true") + .env_remove("AWS_ACCESS_KEY_ID") + .env_remove("AWS_SECRET_ACCESS_KEY") + .env_remove("SECRET_ACCESS_KEY") + .env_remove("AWS_SESSION_TOKEN"); + if matches!(case, "env-precedence" | "profile-precedence") { + child + .env("AWS_ACCESS_KEY_ID", "env-key") + .env("AWS_SECRET_ACCESS_KEY", "env-secret"); + } + if let Some((endpoint, _)) = &role_request { + child.env("AWS_ENDPOINT_URL_STS", endpoint); + } + let output = tokio::task::spawn_blocking(move || child.output()).await??; + assert!( + output.status.success(), + "{case}: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!(marker.exists(), case.contains("process"), "{case}"); + if let Some((_, server)) = role_request { + assert!(server.await??.contains("Credential=process-key/")); + } + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sdk_default_credentials_keep_endpoint_modes_from_environment_and_profile() +-> Result<(), Box> { + for case in ["endpoint-modes-env", "endpoint-modes-profile"] { + let home = tempfile::tempdir()?; + let credentials = home.path().join("credentials"); + let config = home.path().join("config"); + std::fs::write( + &credentials, + "[source]\naws_access_key_id = source-key\naws_secret_access_key = source-secret\n", + )?; + let mut profile = "[profile managed]\nrole_arn = arn:aws:iam::123456789012:role/test\nsource_profile = source\n".to_string(); + if case == "endpoint-modes-profile" { + profile.push_str("use_fips_endpoint = true\nuse_dualstack_endpoint = true\n"); + } + std::fs::write(&config, profile)?; + let proxy = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let proxy_url = format!("http://{}", proxy.local_addr()?); + let server = tokio::spawn(async move { + let (mut socket, _) = proxy.accept().await?; + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(socket.read_u8().await?); + } + socket + .write_all(b"HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\n\r\n") + .await?; + String::from_utf8(headers).map_err(std::io::Error::other) + }); + let mut child = std::process::Command::new(std::env::current_exe()?); + child + .args([ + "--ignored", + "--exact", + "network_policy_tests::sdk_profile_probe", + ]) + .env("CODEX_AWS_PROFILE_PROBE_CASE", case) + .env("AWS_CONFIG_FILE", config) + .env("AWS_SHARED_CREDENTIALS_FILE", credentials) + .env("AWS_EC2_METADATA_DISABLED", "true") + .env("AWS_MAX_ATTEMPTS", "1") + .env("HTTPS_PROXY", &proxy_url) + .env("https_proxy", proxy_url) + .env("NO_PROXY", "") + .env("no_proxy", ""); + for variable in [ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_ENDPOINT_URL", + "AWS_ENDPOINT_URL_STS", + "AWS_USE_FIPS_ENDPOINT", + "AWS_USE_DUALSTACK_ENDPOINT", + ] { + child.env_remove(variable); + } + if case == "endpoint-modes-env" { + child + .env("AWS_USE_FIPS_ENDPOINT", "true") + .env("AWS_USE_DUALSTACK_ENDPOINT", "true"); + } + let output = tokio::task::spawn_blocking(move || child.output()).await??; + assert!( + output.status.success(), + "{case}: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + let request = timeout(std::time::Duration::from_secs(/*secs*/ 5), server).await???; + assert!(request.starts_with("CONNECT sts-fips.us-east-1.api.aws:443 ")); + } + Ok(()) +} + +#[test] +#[ignore = "isolated AWS environment fixture"] +fn sdk_profile_probe() -> Result<(), Box> { + let case = std::env::var("CODEX_AWS_PROFILE_PROBE_CASE")?; + tokio::runtime::Runtime::new()?.block_on(async { + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + let destination = if case.starts_with("unrestricted") { + DestinationPolicy::Unrestricted + } else { + let mut allowed_hosts = ["bedrock-runtime.us-east-1.amazonaws.com".to_string()] + .into_iter() + .collect::>(); + if case.starts_with("endpoint-modes-") { + allowed_hosts.insert("sts-fips.us-east-1.api.aws".to_string()); + } + DestinationPolicy::Restricted { allowed_hosts } + }; + controller.publish(policy.revision(), destination); + let factory = codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ) + .with_network_policy(policy.for_current_account()); + let config = AwsAuthConfig { + profile: Some("managed".into()), + region: Some("us-east-1".into()), + service: "bedrock".into(), + }; + let context = if case == "env-precedence" || case.starts_with("endpoint-modes-") { + AwsAuthContext::load(config, factory).await? + } else { + AwsAuthContext::load_profile(config, factory).await? + }; + let signed = context.sign(super::tests::test_request()).await; + if case.starts_with("endpoint-modes-") { + assert!( + matches!(signed, Err(AwsAuthError::Credentials(_))), + "{signed:?}" + ); + } else if case == "restricted-role-process" { + assert!( + matches!( + signed, + Err(Policy(codex_http_client::NetworkPolicyDenied::Destination)) + ), + "{signed:?}" + ); + } else { + let key = match case.as_str() { + "restricted-static" | "profile-precedence" => "profile-key", + "env-precedence" => "env-key", + "unrestricted-role-process" => "assumed-key", + _ => "process-key", + }; + assert!( + signed?.headers[http::header::AUTHORIZATION] + .to_str()? + .contains(&format!("Credential={key}/")) + ); + } + Ok::<_, Box>(()) + }) +} diff --git a/codex-rs/aws-auth/src/transport.rs b/codex-rs/aws-auth/src/transport.rs new file mode 100644 index 0000000000..c32b42fc07 --- /dev/null +++ b/codex-rs/aws-auth/src/transport.rs @@ -0,0 +1,119 @@ +//! Routes AWS credential and region HTTP through the application-owned client. + +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::Mutex; +use std::task::Context; +use std::task::Poll; +use std::time::Duration; + +use aws_smithy_runtime_api::client::http::HttpConnector; +use aws_smithy_runtime_api::client::http::HttpConnectorFuture; +use aws_smithy_runtime_api::client::http::SharedHttpClient; +use aws_smithy_runtime_api::client::http::SharedHttpConnector; +use aws_smithy_runtime_api::client::http::http_client_fn; +use aws_smithy_runtime_api::client::orchestrator::HttpRequest; +use aws_smithy_runtime_api::client::orchestrator::HttpResponse; +use aws_smithy_runtime_api::client::result::ConnectorError; +use aws_smithy_types::body::SdkBody; +use bytes::Bytes; +use codex_http_client::ClientRouteClass; +use codex_http_client::HttpClient; +use codex_http_client::HttpClientBuilder; +use codex_http_client::HttpClientFactory; +use codex_http_client::HttpError; +use codex_http_client::RouteAwareClientPool; +use http_body_util::BodyExt; + +pub(crate) fn http_client(factory: HttpClientFactory) -> SharedHttpClient { + let connectors = Mutex::new(HashMap::new()); + http_client_fn(move |settings, _| { + let timeouts = (settings.connect_timeout(), settings.read_timeout()); + connectors + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .entry(timeouts) + .or_insert_with(|| { + // Smithy's default connector does not follow redirects. Keep that behavior + // so signed headers and metadata tokens cannot be sent to another origin. + let mut builder = HttpClientBuilder::new() + .without_redirects() + .without_request_logging(); + if let Some(timeout) = timeouts.0 { + builder = builder.connect_timeout(timeout); + } + SharedHttpConnector::new(CredentialConnector { + client: RouteAwareClientPool::with_builder( + factory.clone(), + ClientRouteClass::Auth, + builder, + ) + .into_client(), + read_timeout: timeouts.1, + }) + }) + .clone() + }) +} + +#[derive(Debug)] +struct CredentialConnector { + client: HttpClient, + read_timeout: Option, +} + +impl HttpConnector for CredentialConnector { + fn call(&self, request: HttpRequest) -> HttpConnectorFuture { + let client = self.client.clone(); + let read_timeout = self.read_timeout; + HttpConnectorFuture::new(async move { + let request = request + .try_into_http1x() + .map_err(|error| ConnectorError::user(Box::new(error)))?; + let (parts, body) = request.into_parts(); + let mut request = client + .request(parts.method, parts.uri.to_string()) + .headers(parts.headers) + .version(parts.version); + if let Some(timeout) = read_timeout { + request = request.timeout(timeout); + } + request = match body.bytes() { + Some(bytes) => request.body(Bytes::copy_from_slice(bytes)), + None => request.body_stream(body.into_data_stream()), + }; + let response = request.send().await.map_err(|error| match error { + HttpError::Policy(error) => ConnectorError::user(Box::new(error)), + error if error.is_timeout() => ConnectorError::timeout(Box::new(error)), + error if error.is_connect() => ConnectorError::io(Box::new(error)), + error => ConnectorError::user(Box::new(error)), + })?; + HttpResponse::try_from( + response + .into_http_response() + .map(|body| SdkBody::from_body_1_x(SyncBody(Mutex::new(Box::pin(body))))), + ) + .map_err(|error| ConnectorError::user(Box::new(error))) + }) + } +} + +// Smithy requires Sync bodies; the shared client's cancellation future is only Send. +// Polling stays synchronous, while the body retains policy checks for every frame. +struct SyncBody(Mutex>>); + +impl http_body::Body for SyncBody { + type Data = B::Data; + type Error = B::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + self.0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_mut() + .poll_frame(cx) + } +} diff --git a/codex-rs/cli/src/exec_server_auth.rs b/codex-rs/cli/src/exec_server_auth.rs index ca4a84f37d..96193ea9b5 100644 --- a/codex-rs/cli/src/exec_server_auth.rs +++ b/codex-rs/cli/src/exec_server_auth.rs @@ -19,6 +19,7 @@ use http::HeaderMap; /// Creates a SigV4 provider, preferring an explicit profile over the default credential chain. pub(super) async fn aws_sigv4_auth_provider( mut config: AwsAuthConfig, + http_client_factory: codex_http_client::HttpClientFactory, ) -> Result { config.profile = config .profile @@ -29,9 +30,9 @@ pub(super) async fn aws_sigv4_auth_provider( .map(|region| region.trim().to_string()) .filter(|region| !region.is_empty()); let context = if config.profile.is_some() { - AwsAuthContext::load_profile(config).await + AwsAuthContext::load_profile(config, http_client_factory).await } else { - AwsAuthContext::load(config).await + AwsAuthContext::load(config, http_client_factory).await }?; Ok(Arc::new(AwsSigV4AuthProvider { context })) } diff --git a/codex-rs/cli/src/exec_server_auth_tests.rs b/codex-rs/cli/src/exec_server_auth_tests.rs index dd964f8169..2271d6e09b 100644 --- a/codex-rs/cli/src/exec_server_auth_tests.rs +++ b/codex-rs/cli/src/exec_server_auth_tests.rs @@ -19,6 +19,9 @@ async fn test_provider() -> AwsSigV4AuthProvider { secret_access_key: "test-secret-key".to_string(), session_token: Some("test-session-token".to_string()), }, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), ) .await .expect("load fixture signing context"); @@ -84,11 +87,17 @@ async fn invalid_signing_request_is_a_permanent_auth_error() { #[tokio::test] async fn invalid_signing_configuration_is_rejected() { - let error = aws_sigv4_auth_provider(AwsAuthConfig { + let config = AwsAuthConfig { profile: None, region: Some("us-east-1".to_string()), service: " ".to_string(), - }) + }; + let error = aws_sigv4_auth_provider( + config, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) .await .err() .expect("empty service should fail configuration"); diff --git a/codex-rs/cli/src/exec_server_command.rs b/codex-rs/cli/src/exec_server_command.rs index 77e3016944..633ac84d68 100644 --- a/codex-rs/cli/src/exec_server_command.rs +++ b/codex-rs/cli/src/exec_server_command.rs @@ -168,11 +168,14 @@ impl ExecServerCommand { let direct_transport = self.remote_transport == ExecServerRemoteTransport::Direct; let (_otel, telemetry) = exec_server_telemetry::init(Some(&config)); let auth_provider = if self.aws_sigv4 { - exec_server_auth::aws_sigv4_auth_provider(codex_aws_auth::AwsAuthConfig { - profile: self.aws_profile, - region: self.aws_region, - service: self.aws_service, - }) + exec_server_auth::aws_sigv4_auth_provider( + codex_aws_auth::AwsAuthConfig { + profile: self.aws_profile, + region: self.aws_region, + service: self.aws_service, + }, + config.http_client_factory(), + ) .await? } else { load_exec_server_remote_auth_provider( diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index e494e92d41..d56661f7d0 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -17,6 +17,7 @@ path = "src/lib.rs" workspace = true [dependencies] +codex-http-client = { workspace = true } anyhow = { workspace = true } axum = { workspace = true, features = ["http1", "tokio", "ws"] } clap = { workspace = true, features = ["derive"] } @@ -38,7 +39,6 @@ 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/src/main.rs b/codex-rs/code-mode-host/src/main.rs index d7be1f60cb..fb205fc1bf 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -90,6 +90,9 @@ async fn main() -> anyhow::Result<()> { fn build_trace_provider(endpoint: &str) -> anyhow::Result { OtelProvider::try_new(&OtelSettings { + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), environment: "code-mode-host".to_string(), service_name: "codex-code-mode-host".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), diff --git a/codex-rs/core/src/otel_init.rs b/codex-rs/core/src/otel_init.rs index 9e99b54d19..bd9df090ff 100644 --- a/codex-rs/core/src/otel_init.rs +++ b/codex-rs/core/src/otel_init.rs @@ -81,6 +81,7 @@ pub fn build_provider( let runtime_metrics = config.features.enabled(Feature::RuntimeMetrics); OtelProvider::try_new(&OtelSettings { + http_client_factory: config.http_client_factory(), service_name: service_name.to_string(), service_version: service_version.to_string(), codex_home: config.codex_home.to_path_buf(), diff --git a/codex-rs/core/tests/suite/skills_extension.rs b/codex-rs/core/tests/suite/skills_extension.rs index d2cd7e22b4..f468184a40 100644 --- a/codex-rs/core/tests/suite/skills_extension.rs +++ b/codex-rs/core/tests/suite/skills_extension.rs @@ -2020,6 +2020,7 @@ async fn assert_catalog_model_switch(max_context_tokens: Option) -> Resul let codex_home = Arc::new(TempDir::new()?); // Use the normal metrics sink to verify core's model attribution. let telemetry = OtelProvider::try_new(&OtelSettings { + http_client_factory: codex_core::test_support::default_http_client_factory(), environment: "test".to_string(), service_name: "skills-model-switch".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), diff --git a/codex-rs/exec-server/src/server/process_otel_tests.rs b/codex-rs/exec-server/src/server/process_otel_tests.rs index 3460408ea5..2e53cb1c82 100644 --- a/codex-rs/exec-server/src/server/process_otel_tests.rs +++ b/codex-rs/exec-server/src/server/process_otel_tests.rs @@ -622,6 +622,9 @@ fn exported_logs( tls: None, }; let otel = OtelProvider::try_new(&OtelSettings { + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), environment: "test".to_string(), service_name: "codex-exec-server".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), diff --git a/codex-rs/model-provider/src/amazon_bedrock/auth.rs b/codex-rs/model-provider/src/amazon_bedrock/auth.rs index 1ea6cc690d..4fb04ff816 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/auth.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/auth.rs @@ -95,6 +95,7 @@ pub(super) async fn resolve_auth_method( managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, + http_client_factory: &codex_http_client::HttpClientFactory, ) -> Result { match source { BedrockAuthSource::CommandBearerToken => Err(CodexErr::Fatal( @@ -114,9 +115,13 @@ pub(super) async fn resolve_auth_method( .to_string(), ) })?; - let context = AwsAuthContext::load_with_credentials_provider(config, credential_export) - .await - .map_err(aws_auth_error_to_codex_error)?; + let context = AwsAuthContext::load_with_credentials_provider( + config, + credential_export, + http_client_factory.clone(), + ) + .await + .map_err(aws_auth_error_to_codex_error)?; Ok(BedrockAuthMethod::AwsSdkAuth { context }) } BedrockAuthSource::ManagedBearerToken => { @@ -147,7 +152,7 @@ pub(super) async fn resolve_auth_method( BedrockEndpoint::Mantle => aws_auth_config(aws), BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), }; - let context = AwsAuthContext::load_profile(config) + let context = AwsAuthContext::load_profile(config, http_client_factory.clone()) .await .map_err(aws_auth_error_to_codex_error)?; Ok(BedrockAuthMethod::AwsSdkAuth { context }) @@ -168,9 +173,13 @@ pub(super) async fn resolve_auth_method( BedrockEndpoint::Mantle => aws_auth_config(aws), BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), }; - let context = AwsAuthContext::load_with_access_keys(config, access_keys) - .await - .map_err(aws_auth_error_to_codex_error)?; + let context = AwsAuthContext::load_with_access_keys( + config, + access_keys, + http_client_factory.clone(), + ) + .await + .map_err(aws_auth_error_to_codex_error)?; Ok(BedrockAuthMethod::AwsSdkAuth { context }) } BedrockAuthSource::EnvAwsCredentials | BedrockAuthSource::AwsSdk => { @@ -178,7 +187,7 @@ pub(super) async fn resolve_auth_method( BedrockEndpoint::Mantle => aws_auth_config(aws), BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), }; - let context = AwsAuthContext::load(config) + let context = AwsAuthContext::load(config, http_client_factory.clone()) .await .map_err(aws_auth_error_to_codex_error)?; Ok(BedrockAuthMethod::AwsSdkAuth { context }) @@ -191,8 +200,9 @@ pub(super) async fn resolve_provider_auth( managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, + http_client_factory: &codex_http_client::HttpClientFactory, ) -> Result { - match resolve_auth_method(source, managed_auth, aws, endpoint).await? { + match resolve_auth_method(source, managed_auth, aws, endpoint, http_client_factory).await? { BedrockAuthMethod::ManagedBearerToken { token, .. } | BedrockAuthMethod::EnvBearerToken { token, .. } => Ok(Arc::new(BearerAuthProvider { token: Some(token), @@ -210,19 +220,20 @@ pub(super) async fn resolve_region( managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, + http_client_factory: &codex_http_client::HttpClientFactory, ) -> Result { if source == BedrockAuthSource::CommandBearerToken { let config = match endpoint { BedrockEndpoint::Mantle => aws_auth_config(aws), BedrockEndpoint::Runtime => runtime::aws_auth_config(aws), }; - let context = AwsAuthContext::load(config) + let context = AwsAuthContext::load(config, http_client_factory.clone()) .await .map_err(aws_auth_error_to_codex_error)?; return Ok(context.region().to_string()); } - match resolve_auth_method(source, managed_auth, aws, endpoint).await? { + match resolve_auth_method(source, managed_auth, aws, endpoint, http_client_factory).await? { BedrockAuthMethod::ManagedBearerToken { region, .. } | BedrockAuthMethod::EnvBearerToken { region, .. } => Ok(region), BedrockAuthMethod::AwsSdkAuth { context } => Ok(context.region().to_string()), diff --git a/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs b/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs index fa5db51a8e..1f69d86283 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/credential_export_tests.rs @@ -277,6 +277,37 @@ async fn credential_export_reuses_cache_and_coalesces_refreshes() -> io::Result< ); let exporter = Arc::new(exporter); provider.credential_export = Some(exporter.clone()); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + provider.http_client_factory = provider + .http_client_factory + .clone() + .with_network_policy(policy.clone()); + controller.publish( + policy.revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + let denied = provider + .recover_from_unauthorized() + .await + .expect_err("restricted recovery must not launch the exporter"); + assert_eq!(denied.retry_delay(/*retry_count*/ 1), None); + assert_eq!( + denied.to_string(), + "Fatal error: this SDK transport is disabled by application network restrictions" + ); + assert_eq!( + std::fs::read_to_string(directory.path().join("invocations"))? + .lines() + .collect::>(), + vec!["invoked", "invoked"] + ); + controller.publish( + policy.revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); { let first = provider.recover_from_unauthorized(); let second = provider.recover_from_unauthorized(); diff --git a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs index 1bd2ddee65..a4148a8395 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mantle.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mantle.rs @@ -59,8 +59,16 @@ pub(super) async fn bedrock_mantle_runtime_base_url( source: BedrockAuthSource, managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, + http_client_factory: &codex_http_client::HttpClientFactory, ) -> Result { - let region = resolve_region(source, managed_auth, aws, BedrockEndpoint::Mantle).await?; + let region = resolve_region( + source, + managed_auth, + aws, + BedrockEndpoint::Mantle, + http_client_factory, + ) + .await?; base_url(®ion) } diff --git a/codex-rs/model-provider/src/amazon_bedrock/mod.rs b/codex-rs/model-provider/src/amazon_bedrock/mod.rs index 3a38eb2639..e0e630def1 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/mod.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/mod.rs @@ -59,6 +59,7 @@ pub(super) enum BedrockEndpoint { /// Runtime provider for Amazon Bedrock's OpenAI-compatible endpoints. #[derive(Clone, Debug)] pub(crate) struct AmazonBedrockModelProvider { + http_client_factory: codex_http_client::HttpClientFactory, pub(crate) info: ModelProviderInfo, aws: ModelProviderAwsAuthInfo, endpoint: BedrockEndpoint, @@ -103,8 +104,21 @@ impl AmazonBedrockModelProvider { } else { None }; + let http_client_factory = auth_manager + .as_ref() + .map(|manager| { + manager + .http_client_factory() + .with_network_policy(manager.application_network_policy()) + }) + .unwrap_or_else(|| { + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ) + }); let auth_manager = auth_manager_for_provider(auth_manager, &provider_info); Self { + http_client_factory, info: provider_info, aws, endpoint, @@ -173,21 +187,44 @@ impl AmazonBedrockModelProvider { if let Some(base_url) = self.info.base_url.clone() { return Ok(Some(base_url)); } + let http_client_factory = self.http_client_factory.clone().with_network_policy( + self.http_client_factory + .network_policy() + .clone() + .for_current_account(), + ); let auth_source = self.auth_source(); let managed_auth = self.managed_auth(); let base_url = match self.endpoint { BedrockEndpoint::Mantle => { - bedrock_mantle_runtime_base_url(auth_source, managed_auth.as_ref(), &self.aws) - .await? + bedrock_mantle_runtime_base_url( + auth_source, + managed_auth.as_ref(), + &self.aws, + &http_client_factory, + ) + .await? } BedrockEndpoint::Runtime => { - bedrock_runtime_base_url(auth_source, managed_auth.as_ref(), &self.aws).await? + bedrock_runtime_base_url( + auth_source, + managed_auth.as_ref(), + &self.aws, + &http_client_factory, + ) + .await? } }; Ok(Some(base_url)) } async fn api_auth(&self) -> Result { + let http_client_factory = self.http_client_factory.clone().with_network_policy( + self.http_client_factory + .network_policy() + .clone() + .for_current_account(), + ); let source = self.auth_source(); if source == auth::BedrockAuthSource::CommandBearerToken { let auth = self.auth().await; @@ -195,7 +232,14 @@ impl AmazonBedrockModelProvider { } let managed_auth = self.managed_auth(); - resolve_bedrock_provider_auth(source, managed_auth.as_ref(), &self.aws, self.endpoint).await + resolve_bedrock_provider_auth( + source, + managed_auth.as_ref(), + &self.aws, + self.endpoint, + &http_client_factory, + ) + .await } fn default_model_catalog(&self) -> ModelsResponse { @@ -289,16 +333,26 @@ impl ModelProvider for AmazonBedrockModelProvider { } else { None }; - let result: std::io::Result<()> = async { - if let Some(recovery) = &self.auth_recovery { - recovery.refresh().await?; - } - if let Some(exporter) = export_refresh { - exporter.refresh().await?; - } - Ok(()) - } - .await; + let network_policy = self + .http_client_factory + .network_policy() + .clone() + .for_current_account(); + let permit = network_policy + .acquire_for_unsupported_sdk() + .map_err(|error| CodexErr::Fatal(error.to_string()))?; + let result: std::io::Result<()> = permit + .run(async { + if let Some(recovery) = &self.auth_recovery { + recovery.refresh().await?; + } + if let Some(exporter) = export_refresh { + exporter.refresh().await?; + } + Ok(()) + }) + .await + .map_err(|error| CodexErr::Fatal(error.to_string()))?; result.map_err(|error| { if matches!( error.kind(), diff --git a/codex-rs/model-provider/src/amazon_bedrock/runtime.rs b/codex-rs/model-provider/src/amazon_bedrock/runtime.rs index 7a4d7f410e..f0af3e8de9 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/runtime.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/runtime.rs @@ -26,8 +26,16 @@ pub(super) async fn bedrock_runtime_base_url( source: BedrockAuthSource, managed_auth: Option<&CodexAuth>, aws: &ModelProviderAwsAuthInfo, + http_client_factory: &codex_http_client::HttpClientFactory, ) -> Result { - let region = resolve_region(source, managed_auth, aws, BedrockEndpoint::Runtime).await?; + let region = resolve_region( + source, + managed_auth, + aws, + BedrockEndpoint::Runtime, + http_client_factory, + ) + .await?; Ok(base_url(®ion)) } diff --git a/codex-rs/otel/Cargo.toml b/codex-rs/otel/Cargo.toml index 6efa93043d..362a788972 100644 --- a/codex-rs/otel/Cargo.toml +++ b/codex-rs/otel/Cargo.toml @@ -13,6 +13,9 @@ path = "src/lib.rs" workspace = true [dependencies] +bytes = { workspace = true } +codex-http-client = { workspace = true } +opentelemetry-http = "0.31.0" chrono = { workspace = true } codex-utils-absolute-path = { workspace = true } codex-utils-string = { workspace = true } diff --git a/codex-rs/otel/src/config.rs b/codex-rs/otel/src/config.rs index ab60ea601e..69ad3b62d1 100644 --- a/codex-rs/otel/src/config.rs +++ b/codex-rs/otel/src/config.rs @@ -49,6 +49,7 @@ pub fn validate_span_attributes(attributes: &BTreeMap) -> std::i #[derive(Clone, Debug)] pub struct OtelSettings { + pub http_client_factory: codex_http_client::HttpClientFactory, pub environment: String, pub service_name: String, pub service_version: String, diff --git a/codex-rs/otel/src/lib.rs b/codex-rs/otel/src/lib.rs index 39ebdfddd6..ca4eeab2c7 100644 --- a/codex-rs/otel/src/lib.rs +++ b/codex-rs/otel/src/lib.rs @@ -4,6 +4,7 @@ pub(crate) mod metrics; pub(crate) mod provider; pub(crate) mod trace_context; +mod network_policy; mod otlp; mod targets; mod tool_result; diff --git a/codex-rs/otel/src/metrics/client.rs b/codex-rs/otel/src/metrics/client.rs index 4f48c22693..c8894c8b6e 100644 --- a/codex-rs/otel/src/metrics/client.rs +++ b/codex-rs/otel/src/metrics/client.rs @@ -101,6 +101,7 @@ impl MetricReader for SharedManualReader { #[derive(Debug)] pub(super) struct MetricsClientInner { + pub(super) network_policy: codex_http_client::NetworkPolicy, meter_provider: SdkMeterProvider, meter: Meter, counters: Mutex>>, @@ -323,6 +324,7 @@ impl MetricsClient { /// Build a metrics client from configuration and validate defaults. pub fn new(config: MetricsConfig) -> Result { let MetricsConfig { + http_client_factory, environment, service_name, service_version, @@ -361,13 +363,21 @@ impl MetricsClient { build_provider(resource, exporter, export_interval, runtime_reader.clone()) } MetricsExporter::Otlp(exporter) => { - let exporter = build_otlp_metric_exporter(exporter, Temporality::Delta)?; + let exporter = crate::network_policy::PolicyExporter { + exporter: build_otlp_metric_exporter( + exporter, + Temporality::Delta, + &http_client_factory, + )?, + policy: http_client_factory.network_policy().clone(), + }; build_provider(resource, exporter, export_interval, runtime_reader.clone()) } }; Ok(Self { inner: Arc::new(MetricsClientInner { + network_policy: http_client_factory.network_policy().clone(), meter_provider, meter, counters: Mutex::new(HashMap::new()), @@ -581,12 +591,14 @@ where fn build_otlp_metric_exporter( exporter: OtelExporter, temporality: Temporality, + factory: &codex_http_client::HttpClientFactory, ) -> Result { match exporter { OtelExporter::None => Err(MetricsError::ExporterDisabled), OtelExporter::Statsig => build_otlp_metric_exporter( crate::config::resolve_exporter(&OtelExporter::Statsig), temporality, + factory, ), OtelExporter::OtlpGrpc { endpoint, @@ -638,7 +650,17 @@ fn build_otlp_metric_exporter( .with_protocol(protocol) .with_headers(headers); - if let Some(tls) = tls.as_ref() { + if factory.network_policy().is_managed() { + let client = crate::otlp::build_async_http_client( + factory, + tls.as_ref(), + OTEL_EXPORTER_OTLP_METRICS_TIMEOUT, + ) + .map_err(|err| MetricsError::InvalidConfig { + message: err.to_string(), + })?; + exporter_builder = exporter_builder.with_http_client(client); + } else if let Some(tls) = tls.as_ref() { let client = crate::otlp::build_http_client(tls, OTEL_EXPORTER_OTLP_METRICS_TIMEOUT) .map_err(|err| MetricsError::InvalidConfig { diff --git a/codex-rs/otel/src/metrics/config.rs b/codex-rs/otel/src/metrics/config.rs index 212b63d346..df2203a213 100644 --- a/codex-rs/otel/src/metrics/config.rs +++ b/codex-rs/otel/src/metrics/config.rs @@ -43,6 +43,7 @@ pub enum MetricsExporter { #[derive(Clone, Debug)] pub struct MetricsConfig { + pub(crate) http_client_factory: codex_http_client::HttpClientFactory, pub(crate) environment: String, pub(crate) service_name: String, pub(crate) service_version: String, @@ -54,6 +55,14 @@ pub struct MetricsConfig { } impl MetricsConfig { + pub fn with_http_client_factory( + mut self, + factory: codex_http_client::HttpClientFactory, + ) -> Self { + self.http_client_factory = factory; + self + } + pub fn otlp( environment: impl Into, service_name: impl Into, @@ -66,6 +75,9 @@ impl MetricsConfig { &[] }; Self { + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), environment: environment.into(), service_name: service_name.into(), service_version: service_version.into(), @@ -85,6 +97,9 @@ impl MetricsConfig { exporter: InMemoryMetricExporter, ) -> Self { Self { + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), environment: environment.into(), service_name: service_name.into(), service_version: service_version.into(), diff --git a/codex-rs/otel/src/metrics/mod.rs b/codex-rs/otel/src/metrics/mod.rs index 2367cbd276..32660753ac 100644 --- a/codex-rs/otel/src/metrics/mod.rs +++ b/codex-rs/otel/src/metrics/mod.rs @@ -48,5 +48,11 @@ pub(crate) fn install_global_statsig_settings(settings: StatsigMetricsSettings) } pub(crate) fn global_statsig_settings() -> Option { + if GLOBAL_METRICS + .get() + .is_some_and(|metrics| metrics.active_inner().network_policy.is_managed()) + { + return None; + } GLOBAL_STATSIG_METRICS_SETTINGS.get().cloned() } diff --git a/codex-rs/otel/src/network_policy.rs b/codex-rs/otel/src/network_policy.rs new file mode 100644 index 0000000000..9b9455a1a5 --- /dev/null +++ b/codex-rs/otel/src/network_policy.rs @@ -0,0 +1,127 @@ +//! OTLP SDK transports remain disabled under application destination restrictions. +//! Export futures are revocable, including HTTP exporters driven from SDK worker threads. + +use bytes::Bytes; +use codex_http_client::NetworkPolicy; +use opentelemetry_http::HttpError; +use opentelemetry_sdk::Resource; +use opentelemetry_sdk::error::OTelSdkError; +use opentelemetry_sdk::error::OTelSdkResult; +use opentelemetry_sdk::logs::LogBatch; +use opentelemetry_sdk::logs::LogExporter; +use opentelemetry_sdk::metrics::Temporality; +use opentelemetry_sdk::metrics::data::ResourceMetrics; +use opentelemetry_sdk::metrics::exporter::PushMetricExporter; +use opentelemetry_sdk::trace::SpanData; +use opentelemetry_sdk::trace::SpanExporter; +use std::future::Future; +use std::pin::Pin; +use std::time::Duration; + +#[derive(Debug)] +pub(crate) struct PolicyExporter { + pub(crate) exporter: E, + pub(crate) policy: NetworkPolicy, +} + +fn denied(error: codex_http_client::NetworkPolicyDenied) -> OTelSdkError { + OTelSdkError::InternalFailure(error.to_string()) +} + +impl SpanExporter for PolicyExporter { + async fn export(&self, batch: Vec) -> OTelSdkResult { + let permit = self.policy.acquire_for_unsupported_sdk().map_err(denied)?; + let export = self.exporter.export(batch); + permit.run(export).await.map_err(denied)? + } + fn shutdown_with_timeout(&mut self, timeout: Duration) -> OTelSdkResult { + self.exporter.shutdown_with_timeout(timeout) + } + fn force_flush(&mut self) -> OTelSdkResult { + self.exporter.force_flush() + } + fn set_resource(&mut self, resource: &Resource) { + self.exporter.set_resource(resource); + } +} + +impl LogExporter for PolicyExporter { + async fn export(&self, batch: LogBatch<'_>) -> OTelSdkResult { + let permit = self.policy.acquire_for_unsupported_sdk().map_err(denied)?; + let export = self.exporter.export(batch); + permit.run(export).await.map_err(denied)? + } + fn shutdown_with_timeout(&self, timeout: Duration) -> OTelSdkResult { + self.exporter.shutdown_with_timeout(timeout) + } + fn set_resource(&mut self, resource: &Resource) { + self.exporter.set_resource(resource); + } +} + +impl PushMetricExporter for PolicyExporter { + async fn export(&self, metrics: &ResourceMetrics) -> OTelSdkResult { + let permit = self.policy.acquire_for_unsupported_sdk().map_err(denied)?; + let export = self.exporter.export(metrics); + permit.run(export).await.map_err(denied)? + } + fn force_flush(&self) -> OTelSdkResult { + self.exporter.force_flush() + } + fn shutdown_with_timeout(&self, timeout: Duration) -> OTelSdkResult { + self.exporter.shutdown_with_timeout(timeout) + } + fn temporality(&self) -> Temporality { + self.exporter.temporality() + } +} + +/// Keeps blocking SDK worker threads from performing uncancellable HTTP I/O. +#[derive(Debug)] +pub(crate) struct RuntimeHttpClient { + pub(crate) client: codex_http_client::HttpClient, + pub(crate) timeout: Duration, + pub(crate) runtime: tokio::runtime::Handle, +} + +impl opentelemetry_http::HttpClient for RuntimeHttpClient { + fn send_bytes<'client, 'request>( + &'client self, + request: http::Request, + ) -> Pin, HttpError>> + Send + 'request>> + where + 'client: 'request, + Self: 'request, + { + Box::pin(async move { + let client = self.client.clone(); + let timeout = self.timeout; + let mut requests = tokio::task::JoinSet::new(); + requests.spawn_on( + async move { + let (parts, body) = request.into_parts(); + let response = client + .request(parts.method, parts.uri.to_string()) + .version(parts.version) + .headers(parts.headers) + .body(body) + .timeout(timeout) + .send() + .await? + .error_for_status()?; + let mut builder = http::Response::builder().status(response.status()); + if let Some(headers) = builder.headers_mut() { + *headers = response.headers().clone(); + } + Ok::<_, HttpError>(builder.body(response.bytes().await?)?) + }, + &self.runtime, + ); + // Dropping the SDK export future drops this set and aborts the HTTP request. + requests + .join_next() + .await + .ok_or("telemetry request task disappeared")?? + }) + } +} diff --git a/codex-rs/otel/src/otlp.rs b/codex-rs/otel/src/otlp.rs index f098542d53..09b9b01a5a 100644 --- a/codex-rs/otel/src/otlp.rs +++ b/codex-rs/otel/src/otlp.rs @@ -146,23 +146,23 @@ fn build_http_client_inner( } pub(crate) fn build_async_http_client( + factory: &codex_http_client::HttpClientFactory, tls: Option<&OtelTlsConfig>, timeout_var: &str, -) -> Result> { - let mut builder = reqwest::Client::builder().timeout(resolve_otlp_timeout(timeout_var)); +) -> Result> { + let mut config = codex_http_client::HttpClientTlsConfig::default(); if let Some(tls) = tls { if let Some(path) = tls.ca_certificate.as_ref() { let (pem, location) = read_bytes(path)?; - let certificate = ReqwestCertificate::from_pem(pem.as_slice()).map_err(|error| { - config_error(format!( - "failed to parse certificate {}: {error}", - location.display() - )) - })?; - builder = builder - .tls_built_in_root_certs(false) - .add_root_certificate(certificate); + config = config + .with_root_certificate_pem(pem.as_slice()) + .map_err(|error| { + config_error(format!( + "failed to parse certificate {}: {error}", + location.display() + )) + })?; } match (&tls.client_certificate, &tls.client_private_key) { @@ -170,14 +170,15 @@ pub(crate) fn build_async_http_client( let (mut cert_pem, cert_location) = read_bytes(cert_path)?; let (key_pem, key_location) = read_bytes(key_path)?; cert_pem.extend_from_slice(key_pem.as_slice()); - let identity = ReqwestIdentity::from_pem(cert_pem.as_slice()).map_err(|error| { - config_error(format!( - "failed to parse client identity using {} and {}: {error}", - cert_location.display(), - key_location.display() - )) - })?; - builder = builder.identity(identity).https_only(true); + config = config + .with_client_identity_pem(cert_pem.as_slice()) + .map_err(|error| { + config_error(format!( + "failed to parse client identity using {} and {}: {error}", + cert_location.display(), + key_location.display() + )) + })?; } (Some(_), None) | (None, Some(_)) => { return Err(config_error( @@ -188,9 +189,13 @@ pub(crate) fn build_async_http_client( } } - builder - .build() - .map_err(|error| Box::new(error) as Box) + Ok(crate::network_policy::RuntimeHttpClient { + client: codex_http_client::HttpClientBuilder::new() + .without_request_logging() + .build_with_tls(factory, codex_http_client::ClientRouteClass::Other, config), + timeout: resolve_otlp_timeout(timeout_var), + runtime: tokio::runtime::Handle::try_current()?, + }) } pub(crate) fn resolve_otlp_timeout(signal_var: &str) -> Duration { diff --git a/codex-rs/otel/src/provider.rs b/codex-rs/otel/src/provider.rs index 514453d32b..6d35b555a4 100644 --- a/codex-rs/otel/src/provider.rs +++ b/codex-rs/otel/src/provider.rs @@ -222,7 +222,8 @@ impl OtelProvider { settings.service_name.clone(), settings.service_version.clone(), settings.metrics_exporter.clone(), - ); + ) + .with_http_client_factory(settings.http_client_factory.clone()); if settings.runtime_metrics { config = config.with_runtime_reader(); } @@ -232,7 +233,13 @@ impl OtelProvider { let log_resource = make_resource(settings, ResourceKind::Logs); let trace_resource = make_resource(settings, ResourceKind::Traces); let logger = log_enabled - .then(|| build_logger(&log_resource, &settings.exporter)) + .then(|| { + build_logger( + &log_resource, + &settings.exporter, + &settings.http_client_factory, + ) + }) .transpose()?; let tracer_provider = trace_enabled @@ -241,6 +248,7 @@ impl OtelProvider { &trace_resource, &settings.trace_exporter, settings.span_attributes.clone(), + &settings.http_client_factory, ) }) .transpose()?; @@ -434,6 +442,7 @@ impl SpanProcessor for SpanAttributesProcessor { fn build_logger( resource: &Resource, exporter: &OtelExporter, + factory: &codex_http_client::HttpClientFactory, ) -> Result> { let mut builder = SdkLoggerProvider::builder().with_resource(resource.clone()); @@ -465,7 +474,10 @@ fn build_logger( .with_tls_config(tls_config) .build()?; - builder = builder.with_batch_exporter(exporter); + builder = builder.with_batch_exporter(crate::network_policy::PolicyExporter { + exporter, + policy: factory.network_policy().clone(), + }); } OtelExporter::OtlpHttp { endpoint, @@ -486,14 +498,24 @@ fn build_logger( .with_protocol(protocol) .with_headers(headers); - if let Some(tls) = tls.as_ref() { + if factory.network_policy().is_managed() { + let client = crate::otlp::build_async_http_client( + factory, + tls.as_ref(), + OTEL_EXPORTER_OTLP_LOGS_TIMEOUT, + )?; + exporter_builder = exporter_builder.with_http_client(client); + } else if let Some(tls) = tls.as_ref() { let client = crate::otlp::build_http_client(tls, OTEL_EXPORTER_OTLP_LOGS_TIMEOUT)?; exporter_builder = exporter_builder.with_http_client(client); } let exporter = exporter_builder.build()?; - builder = builder.with_batch_exporter(exporter); + builder = builder.with_batch_exporter(crate::network_policy::PolicyExporter { + exporter, + policy: factory.network_policy().clone(), + }); } } @@ -504,6 +526,7 @@ fn build_tracer_provider( resource: &Resource, exporter: &OtelExporter, span_attributes: BTreeMap, + factory: &codex_http_client::HttpClientFactory, ) -> Result> { let span_exporter = match crate::config::resolve_exporter(exporter) { OtelExporter::None => return Ok(tracer_provider_builder(resource, span_attributes).build()), @@ -554,14 +577,20 @@ fn build_tracer_provider( .with_headers(headers); let client = crate::otlp::build_async_http_client( + factory, tls.as_ref(), OTEL_EXPORTER_OTLP_TRACES_TIMEOUT, )?; exporter_builder = exporter_builder.with_http_client(client); - let processor = - TokioBatchSpanProcessor::builder(exporter_builder.build()?, runtime::Tokio) - .build(); + let processor = TokioBatchSpanProcessor::builder( + crate::network_policy::PolicyExporter { + exporter: exporter_builder.build()?, + policy: factory.network_policy().clone(), + }, + runtime::Tokio, + ) + .build(); return Ok(tracer_provider_builder(resource, span_attributes) .with_span_processor(processor) @@ -579,7 +608,14 @@ fn build_tracer_provider( .with_protocol(protocol) .with_headers(headers); - if let Some(tls) = tls.as_ref() { + if factory.network_policy().is_managed() { + let client = crate::otlp::build_async_http_client( + factory, + tls.as_ref(), + OTEL_EXPORTER_OTLP_TRACES_TIMEOUT, + )?; + exporter_builder = exporter_builder.with_http_client(client); + } else if let Some(tls) = tls.as_ref() { let client = crate::otlp::build_http_client(tls, OTEL_EXPORTER_OTLP_TRACES_TIMEOUT)?; exporter_builder = exporter_builder.with_http_client(client); @@ -589,7 +625,11 @@ fn build_tracer_provider( } }; - let processor = BatchSpanProcessor::builder(span_exporter).build(); + let processor = BatchSpanProcessor::builder(crate::network_policy::PolicyExporter { + exporter: span_exporter, + policy: factory.network_policy().clone(), + }) + .build(); Ok(tracer_provider_builder(resource, span_attributes) .with_span_processor(processor) @@ -614,6 +654,8 @@ mod tests { use crate::metrics::TOOL_CALL_DURATION_METRIC; use crate::metrics::TURN_COST_MICROUSD_METRIC; use crate::metrics::TURN_TOKEN_USAGE_METRIC; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; use opentelemetry_sdk::metrics::InMemoryMetricExporter; use pretty_assertions::assert_eq; use std::path::PathBuf; @@ -690,16 +732,32 @@ mod tests { env!("CARGO_PKG_VERSION"), InMemoryMetricExporter::default(), ))?); + crate::metrics::install_global_statsig_settings(StatsigMetricsSettings { + environment: "test".to_string(), + }); + assert!(crate::metrics::global_statsig_settings().is_some()); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish( + policy.revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); let cached = crate::metrics::global().expect("initial global metrics client"); let exporter = InMemoryMetricExporter::default(); - let replacement = - crate::metrics::install_global(MetricsClient::new(MetricsConfig::in_memory( + let replacement = crate::metrics::install_global(MetricsClient::new( + MetricsConfig::in_memory( "test", "codex-test", env!("CARGO_PKG_VERSION"), exporter.clone(), - ))?); + ) + .with_http_client_factory( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy), + ), + )?); + assert!(crate::metrics::global_statsig_settings().is_none()); cached.counter("codex.after_transition", /*inc*/ 1, &[])?; initial.shutdown()?; replacement.shutdown()?; @@ -777,6 +835,7 @@ mod tests { fn test_otel_settings() -> OtelSettings { OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-test".to_string(), service_version: "0.0.0".to_string(), diff --git a/codex-rs/otel/tests/suite/otlp_http_loopback.rs b/codex-rs/otel/tests/suite/otlp_http_loopback.rs index 7bd209b138..d0cedbef52 100644 --- a/codex-rs/otel/tests/suite/otlp_http_loopback.rs +++ b/codex-rs/otel/tests/suite/otlp_http_loopback.rs @@ -1,3 +1,5 @@ +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; use codex_otel::EXEC_SERVER_CLIENT_REQUEST_COUNT_METRIC; use codex_otel::MetricsClient; use codex_otel::MetricsConfig; @@ -355,6 +357,7 @@ fn otlp_http_exporter_sends_logs_to_collector() }); let otel = OtelProvider::try_new(&OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-cli".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), @@ -414,6 +417,7 @@ fn otlp_http_exporter_sends_logs_to_collector() #[test] fn otel_provider_rejects_header_unsafe_configured_tracestate() { let result = OtelProvider::try_new(&OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-cli".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), @@ -479,6 +483,7 @@ fn otlp_http_exporter_sends_traces_to_collector() }); let otel = OtelProvider::try_new(&OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-cli".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), @@ -624,6 +629,7 @@ async fn otlp_http_exporter_sends_traces_to_collector_with_bounded_shutdown_in_t }); let otel = OtelProvider::try_new(&OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-cli".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), @@ -717,6 +723,7 @@ fn otlp_http_exporter_times_out_when_collector_stalls_during_bounded_shutdown() .expect("build tokio runtime"); let (result, elapsed) = runtime.block_on(async move { let otel = OtelProvider::try_new(&OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-cli".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), @@ -816,6 +823,7 @@ fn otlp_http_exporter_sends_traces_to_collector_in_current_thread_tokio_runtime( let result = runtime.block_on(async move { let otel = OtelProvider::try_new(&OtelSettings { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), environment: "test".to_string(), service_name: "codex-cli".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), diff --git a/codex-rs/windows-sandbox-rs/Cargo.toml b/codex-rs/windows-sandbox-rs/Cargo.toml index b9e5ce694c..b182ed8c9e 100644 --- a/codex-rs/windows-sandbox-rs/Cargo.toml +++ b/codex-rs/windows-sandbox-rs/Cargo.toml @@ -26,6 +26,7 @@ path = "src/bin/managed_deny_probe/main.rs" workspace = true [dependencies] +codex-http-client = { workspace = true } anyhow = "1.0" base64 = { workspace = true } chrono = { version = "0.4.42", default-features = false, features = [ diff --git a/codex-rs/windows-sandbox-rs/src/wfp_setup.rs b/codex-rs/windows-sandbox-rs/src/wfp_setup.rs index b3d6c4d339..45e8dc2382 100644 --- a/codex-rs/windows-sandbox-rs/src/wfp_setup.rs +++ b/codex-rs/windows-sandbox-rs/src/wfp_setup.rs @@ -47,6 +47,9 @@ fn build_wfp_metrics_provider( // Statsig environment in the elevation payload. Other exporters are // intentionally omitted from this helper path. OtelProvider::try_new(&OtelSettings { + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), environment: otel.environment.clone(), service_name: WFP_SETUP_SERVICE_NAME.to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(),