Enforce application network policy for AWS auth and telemetry (#47408)

## Why

AWS credential discovery, analytics, and telemetry used clients that did not share the application's network policy. These requests need to honor destination restrictions and permission revocation.

## What changed

- Route AWS credential and region HTTP requests through the shared HTTP client, and check the signing destination before loading credentials. Skip discovery for static access keys with an explicit region.
- Apply account-scoped policy to Bedrock authentication. Require unrestricted policy for credential exporters and reauthentication commands, and cancel active work when permission is revoked. Document that AWS profile `credential_process` network traffic remains outside the application's HTTP policy.
- Send analytics through the authenticated account's HTTP client factory.
- Guard OTLP log, trace, and metric exports with revocable permits, disable them under destination restrictions, and make managed HTTP exports cancellable. Suppress global Statsig settings while managed policy is active.

## Testing

Add coverage for metadata credential request cancellation, allowed and denied SigV4 destinations, AWS credential precedence and endpoint configuration, and blocked credential exporter recovery. Update telemetry tests for account transitions and managed-policy Statsig suppression.

GitOrigin-RevId: d8a2018bfbbe971ec430699b0d3f86a7a9e1e072
This commit is contained in:
acrognale-oai
2026-09-23 00:47:34 +00:00
committed by copyberry
parent 22a3f6d5d8
commit 8f8ace78cb
39 changed files with 1296 additions and 162 deletions
+16
View File
@@ -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",
+2
View File
@@ -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"
+1
View File
@@ -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 }
+19 -4
View File
@@ -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<TrackEventRequest>,
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())
+6 -2
View File
@@ -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::<Vec<_>>();
@@ -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");
+6
View File
@@ -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.
@@ -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 { .. } => {
+35 -3
View File
@@ -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");
+12 -1
View File
@@ -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"] }
+9 -4
View File
@@ -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<SdkConfig, AwsAuthError> {
pub(crate) async fn load_sdk_config(
config: &AwsAuthConfig,
factory: &HttpClientFactory,
) -> Result<SdkConfig, AwsAuthError> {
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)
}
+13 -19
View File
@@ -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<Vec<AwsProfile>, 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()
}
+120 -48
View File
@@ -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<P>(Arc<P>);
struct AwsCredentialsProviderAdapter<P>(Arc<P>, codex_http_client::NetworkPolicy);
#[derive(Debug, Error)]
#[error("{0}")]
@@ -70,7 +73,12 @@ impl<P: AwsCredentialsProvider> 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<Self, AwsAuthError> {
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<Self, AwsAuthError> {
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<Self, AwsAuthError> {
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<impl AwsCredentialsProvider + 'static>,
factory: codex_http_client::HttpClientFactory,
) -> Result<Self, AwsAuthError> {
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<Self, AwsAuthError> {
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<Self, AwsAuthError> {
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<AwsSignedRequest, AwsAuthError> {
let credentials = self.credentials_provider.provide_credentials().await?;
let url = request
.url
.parse::<url::Url>()
.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::<codex_http_client::NetworkPolicyDenied>()
{
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;
@@ -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<dyn Error>> {
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<dyn Error>> {
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<dyn Error>>(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<dyn Error>> {
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::<usize>())
})
.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 = "<AssumeRoleResponse><AssumeRoleResult><Credentials><AccessKeyId>assumed-key</AccessKeyId><SecretAccessKey>assumed-secret</SecretAccessKey><SessionToken>assumed-token</SessionToken><Expiration>2099-01-01T00:00:00Z</Expiration></Credentials></AssumeRoleResult></AssumeRoleResponse>";
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<dyn Error>> {
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<dyn Error>> {
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::<std::collections::BTreeSet<_>>();
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<dyn Error>>(())
})
}
+119
View File
@@ -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<Duration>,
}
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<B>(Mutex<Pin<Box<B>>>);
impl<B: http_body::Body> http_body::Body for SyncBody<B> {
type Data = B::Data;
type Error = B::Error;
fn poll_frame(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<http_body::Frame<Self::Data>, Self::Error>>> {
self.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.as_mut()
.poll_frame(cx)
}
}
+3 -2
View File
@@ -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<SharedAuthProvider, AwsAuthError> {
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 }))
}
+11 -2
View File
@@ -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");
+8 -5
View File
@@ -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(
+1 -1
View File
@@ -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 }
+3
View File
@@ -90,6 +90,9 @@ async fn main() -> anyhow::Result<()> {
fn build_trace_provider(endpoint: &str) -> anyhow::Result<OtelProvider> {
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(),
+1
View File
@@ -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(),
@@ -2020,6 +2020,7 @@ async fn assert_catalog_model_switch(max_context_tokens: Option<usize>) -> 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(),
@@ -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(),
@@ -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<BedrockAuthMethod> {
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<SharedAuthProvider> {
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<String> {
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()),
@@ -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<_>>(),
vec!["invoked", "invoked"]
);
controller.publish(
policy.revision(),
codex_http_client::DestinationPolicy::Unrestricted,
);
{
let first = provider.recover_from_unauthorized();
let second = provider.recover_from_unauthorized();
@@ -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<String> {
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(&region)
}
@@ -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<SharedAuthProvider> {
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(),
@@ -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<String> {
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(&region))
}
+3
View File
@@ -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 }
+1
View File
@@ -49,6 +49,7 @@ pub fn validate_span_attributes(attributes: &BTreeMap<String, String>) -> 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,
+1
View File
@@ -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;
+24 -2
View File
@@ -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<HashMap<InstrumentKey, Counter<u64>>>,
@@ -323,6 +324,7 @@ impl MetricsClient {
/// Build a metrics client from configuration and validate defaults.
pub fn new(config: MetricsConfig) -> Result<Self> {
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<opentelemetry_otlp::MetricExporter> {
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 {
+15
View File
@@ -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<String>,
service_name: impl Into<String>,
@@ -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(),
+6
View File
@@ -48,5 +48,11 @@ pub(crate) fn install_global_statsig_settings(settings: StatsigMetricsSettings)
}
pub(crate) fn global_statsig_settings() -> Option<StatsigMetricsSettings> {
if GLOBAL_METRICS
.get()
.is_some_and(|metrics| metrics.active_inner().network_policy.is_managed())
{
return None;
}
GLOBAL_STATSIG_METRICS_SETTINGS.get().cloned()
}
+127
View File
@@ -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<E> {
pub(crate) exporter: E,
pub(crate) policy: NetworkPolicy,
}
fn denied(error: codex_http_client::NetworkPolicyDenied) -> OTelSdkError {
OTelSdkError::InternalFailure(error.to_string())
}
impl<E: SpanExporter> SpanExporter for PolicyExporter<E> {
async fn export(&self, batch: Vec<SpanData>) -> 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<E: LogExporter> LogExporter for PolicyExporter<E> {
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<E: PushMetricExporter> PushMetricExporter for PolicyExporter<E> {
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<Bytes>,
) -> Pin<Box<dyn Future<Output = Result<http::Response<Bytes>, 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")??
})
}
}
+27 -22
View File
@@ -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<reqwest::Client, Box<dyn Error>> {
let mut builder = reqwest::Client::builder().timeout(resolve_otlp_timeout(timeout_var));
) -> Result<crate::network_policy::RuntimeHttpClient, Box<dyn Error>> {
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<dyn Error>)
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 {
+72 -13
View File
@@ -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<SdkLoggerProvider, Box<dyn Error>> {
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<String, String>,
factory: &codex_http_client::HttpClientFactory,
) -> Result<SdkTracerProvider, Box<dyn Error>> {
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(),
@@ -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(),
+1
View File
@@ -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 = [
@@ -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(),