From 5800d3a4b67055b8be92c9d23457fe27f76d9640 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Tue, 23 Jun 2026 11:56:05 -0700 Subject: [PATCH] refactor(network-proxy): bind attribution at shared ingress Propagate an optional execution identity through proxy policy and audit paths, and bind either fixed or token-backed context once per accepted connection. Unknown tokens are rejected before request handling. Co-authored-by: Codex noreply@openai.com --- .../core/src/network_policy_decision_tests.rs | 2 + .../core/src/tools/network_approval_tests.rs | 1 + codex-rs/network-proxy/src/attribution.rs | 166 +++++++++++++ .../network-proxy/src/attribution_tests.rs | 58 +++++ codex-rs/network-proxy/src/http_proxy.rs | 219 +++++++++++++----- codex-rs/network-proxy/src/lib.rs | 3 + codex-rs/network-proxy/src/mitm.rs | 11 + codex-rs/network-proxy/src/mitm_tests.rs | 1 + codex-rs/network-proxy/src/network_policy.rs | 29 +++ codex-rs/network-proxy/src/proxy.rs | 38 +-- codex-rs/network-proxy/src/runtime.rs | 8 + codex-rs/network-proxy/src/socks5.rs | 119 +++++++--- 12 files changed, 549 insertions(+), 106 deletions(-) create mode 100644 codex-rs/network-proxy/src/attribution.rs create mode 100644 codex-rs/network-proxy/src/attribution_tests.rs diff --git a/codex-rs/core/src/network_policy_decision_tests.rs b/codex-rs/core/src/network_policy_decision_tests.rs index 064004184701..8888bbe9c2d8 100644 --- a/codex-rs/core/src/network_policy_decision_tests.rs +++ b/codex-rs/core/src/network_policy_decision_tests.rs @@ -161,6 +161,7 @@ fn denied_network_policy_message_requires_deny_decision() { method: Some("GET".to_string()), mode: None, protocol: "http".to_string(), + execution_id: None, decision: Some("ask".to_string()), source: Some("decider".to_string()), port: Some(80), @@ -178,6 +179,7 @@ fn denied_network_policy_message_for_denylist_block_is_explicit() { method: Some("GET".to_string()), mode: None, protocol: "http".to_string(), + execution_id: None, decision: Some("deny".to_string()), source: Some("baseline_policy".to_string()), port: Some(80), diff --git a/codex-rs/core/src/tools/network_approval_tests.rs b/codex-rs/core/src/tools/network_approval_tests.rs index 60c7215b1b0b..8ad1f9909a9a 100644 --- a/codex-rs/core/src/tools/network_approval_tests.rs +++ b/codex-rs/core/src/tools/network_approval_tests.rs @@ -280,6 +280,7 @@ fn denied_blocked_request(host: &str) -> BlockedRequest { method: None, mode: None, protocol: "http".to_string(), + execution_id: None, decision: Some("deny".to_string()), source: Some("decider".to_string()), port: Some(80), diff --git a/codex-rs/network-proxy/src/attribution.rs b/codex-rs/network-proxy/src/attribution.rs new file mode 100644 index 000000000000..23c0fd310c9e --- /dev/null +++ b/codex-rs/network-proxy/src/attribution.rs @@ -0,0 +1,166 @@ +use crate::network_policy::NetworkRequestContext; +use rama_core::Service; +use rama_core::error::BoxError; +use rama_core::extensions::ExtensionsMut; +use rama_tcp::TcpStream; +use std::collections::HashMap; +use std::io; +use std::io::Write; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use tokio::io::AsyncReadExt; + +/// Internal handoff to a trusted proxy bridge. +#[doc(hidden)] +pub const PROXY_ATTRIBUTION_TOKEN_ENV_KEY: &str = "CODEX_NETWORK_PROXY_ATTRIBUTION"; + +const ATTRIBUTION_FRAME_MAGIC: &[u8; 8] = b"\0CDXPXY1"; +const MAX_ATTRIBUTION_TOKEN_LEN: usize = 128; +const ATTRIBUTION_FRAME_TIMEOUT: Duration = Duration::from_secs(3); + +#[derive(Clone, Default)] +pub(crate) struct AttributionRegistry { + contexts: Arc>>, +} + +impl AttributionRegistry { + pub(crate) fn context_for_token(&self, token: &str) -> Option { + self.contexts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(token) + .cloned() + } + + pub(crate) fn clear(&self) { + self.contexts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clear(); + } + + #[cfg(test)] + fn register(&self, token: impl Into, context: NetworkRequestContext) { + self.contexts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(token.into(), context); + } +} + +#[derive(Clone)] +pub(crate) enum ConnectionAttribution { + Fixed(NetworkRequestContext), + Registry(AttributionRegistry), +} + +pub(crate) struct BindConnectionAttribution { + inner: S, + attribution: ConnectionAttribution, +} + +impl BindConnectionAttribution { + pub(crate) fn new(inner: S, attribution: ConnectionAttribution) -> Self { + Self { inner, attribution } + } +} + +impl Service for BindConnectionAttribution +where + S: Service, + S::Error: Into, +{ + type Output = S::Output; + type Error = BoxError; + + async fn serve(&self, mut stream: TcpStream) -> Result { + let context = match &self.attribution { + ConnectionAttribution::Fixed(context) => context.clone(), + ConnectionAttribution::Registry(registry) => { + match read_attribution_token(&mut stream).await? { + Some(token) => registry.context_for_token(&token).ok_or_else(|| { + io::Error::new( + io::ErrorKind::PermissionDenied, + "unknown network proxy attribution token", + ) + })?, + None => NetworkRequestContext::default(), + } + } + }; + stream.extensions_mut().insert(context); + self.inner.serve(stream).await.map_err(Into::into) + } +} + +async fn read_attribution_token(stream: &mut TcpStream) -> Result, BoxError> { + let mut marker = [0_u8; 1]; + let read = stream.stream.peek(&mut marker).await?; + if read == 0 { + return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "empty proxy connection").into()); + } + if marker[0] != ATTRIBUTION_FRAME_MAGIC[0] { + return Ok(None); + } + + let token = tokio::time::timeout(ATTRIBUTION_FRAME_TIMEOUT, async { + let mut magic = [0_u8; ATTRIBUTION_FRAME_MAGIC.len()]; + stream.read_exact(&mut magic).await?; + if &magic != ATTRIBUTION_FRAME_MAGIC { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid network proxy attribution frame", + )); + } + + let token_len = stream.read_u16().await? as usize; + if token_len == 0 || token_len > MAX_ATTRIBUTION_TOKEN_LEN { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid network proxy attribution token length", + )); + } + let mut token = vec![0_u8; token_len]; + stream.read_exact(&mut token).await?; + String::from_utf8(token).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + "network proxy attribution token is not UTF-8", + ) + }) + }) + .await + .map_err(|_| { + io::Error::new( + io::ErrorKind::TimedOut, + "network proxy attribution frame timed out", + ) + })??; + + Ok(Some(token)) +} + +/// Writes the trusted bridge preface consumed by the shared proxy ingress. +#[doc(hidden)] +pub fn write_attribution_frame(writer: &mut impl Write, token: &str) -> io::Result<()> { + if token.is_empty() || token.len() > MAX_ATTRIBUTION_TOKEN_LEN { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid network proxy attribution token length", + )); + } + let token_len = u16::try_from(token.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "network proxy attribution token is too long", + ) + })?; + writer.write_all(ATTRIBUTION_FRAME_MAGIC)?; + writer.write_all(&token_len.to_be_bytes())?; + writer.write_all(token.as_bytes()) +} + +#[cfg(test)] +#[path = "attribution_tests.rs"] +mod tests; diff --git a/codex-rs/network-proxy/src/attribution_tests.rs b/codex-rs/network-proxy/src/attribution_tests.rs new file mode 100644 index 000000000000..923d78c1d763 --- /dev/null +++ b/codex-rs/network-proxy/src/attribution_tests.rs @@ -0,0 +1,58 @@ +use super::AttributionRegistry; +use super::BindConnectionAttribution; +use super::ConnectionAttribution; +use super::write_attribution_frame; +use crate::network_policy::NetworkRequestContext; +use pretty_assertions::assert_eq; +use rama_core::Service; +use rama_core::error::BoxError; +use rama_core::extensions::ExtensionsRef; +use rama_core::service::service_fn; +use rama_tcp::TcpStream as RamaTcpStream; +use std::io; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::net::TcpStream; + +#[test] +fn attribution_frame_has_bounded_binary_prefix() -> io::Result<()> { + let mut frame = Vec::new(); + write_attribution_frame(&mut frame, "token-1")?; + + assert_eq!(&frame[..8], b"\0CDXPXY1"); + assert_eq!(u16::from_be_bytes([frame[8], frame[9]]), 7); + assert_eq!(&frame[10..], b"token-1"); + Ok(()) +} + +#[tokio::test] +async fn framed_connection_receives_registered_attribution() -> Result<(), BoxError> { + let registry = AttributionRegistry::default(); + let expected = NetworkRequestContext { + environment_id: Some("local".to_string()), + execution_id: Some("execution-1".to_string()), + }; + registry.register("token-1", expected.clone()); + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let addr = listener.local_addr()?; + let client = tokio::spawn(async move { + let mut stream = TcpStream::connect(addr).await?; + let mut frame = Vec::new(); + write_attribution_frame(&mut frame, "token-1")?; + stream.write_all(&frame).await + }); + + let (stream, _) = listener.accept().await?; + let service = BindConnectionAttribution::new( + service_fn(|stream: RamaTcpStream| async move { + Ok::<_, io::Error>(stream.extensions().get::().cloned()) + }), + ConnectionAttribution::Registry(registry), + ); + let actual = service.serve(RamaTcpStream::new(stream)).await?; + client.await??; + + assert_eq!(actual, Some(expected)); + Ok(()) +} diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index 3f02ad3fc07d..2f522f6590b6 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -1,3 +1,5 @@ +use crate::attribution::BindConnectionAttribution; +use crate::attribution::ConnectionAttribution; use crate::config::NetworkMode; use crate::connect_policy::TargetCheckedTcpConnector; use crate::mitm; @@ -9,6 +11,7 @@ use crate::network_policy::NetworkPolicyDecision; use crate::network_policy::NetworkPolicyRequest; use crate::network_policy::NetworkPolicyRequestArgs; use crate::network_policy::NetworkProtocol; +use crate::network_policy::NetworkRequestContext; use crate::network_policy::emit_allow_decision_audit_event; use crate::network_policy::emit_block_decision_audit_event; use crate::network_policy::evaluate_host_policy; @@ -83,11 +86,11 @@ use tracing::warn; #[derive(Clone, Copy, Debug)] struct ConnectMitmEnabled(bool); -pub async fn run_http_proxy( +pub(crate) async fn run_http_proxy_with_attribution( state: Arc, addr: SocketAddr, policy_decider: Option>, - environment_id: Option, + attribution: ConnectionAttribution, ) -> Result<()> { let listener = TcpListener::build() .bind(addr) @@ -100,25 +103,40 @@ pub async fn run_http_proxy( .map_err(anyhow::Error::from) .with_context(|| format!("bind HTTP proxy: {addr}"))?; - run_http_proxy_with_listener(state, listener, policy_decider, environment_id).await + run_http_proxy_with_listener(state, listener, policy_decider, attribution).await } pub async fn run_http_proxy_with_std_listener( state: Arc, listener: StdTcpListener, policy_decider: Option>, - environment_id: Option, + request_context: NetworkRequestContext, +) -> Result<()> { + run_http_proxy_with_std_listener_and_attribution( + state, + listener, + policy_decider, + ConnectionAttribution::Fixed(request_context), + ) + .await +} + +pub(crate) async fn run_http_proxy_with_std_listener_and_attribution( + state: Arc, + listener: StdTcpListener, + policy_decider: Option>, + attribution: ConnectionAttribution, ) -> Result<()> { let listener = TcpListener::try_from(listener).context("convert std listener to HTTP proxy listener")?; - run_http_proxy_with_listener(state, listener, policy_decider, environment_id).await + run_http_proxy_with_listener(state, listener, policy_decider, attribution).await } async fn run_http_proxy_with_listener( state: Arc, listener: TcpListener, policy_decider: Option>, - environment_id: Option, + attribution: ConnectionAttribution, ) -> Result<()> { ensure_rustls_crypto_provider(); @@ -136,9 +154,21 @@ async fn run_http_proxy_with_listener( MethodMatcher::CONNECT, service_fn({ let policy_decider = policy_decider.clone(); - let environment_id = environment_id.clone(); - move |req| { - http_connect_accept(policy_decider.clone(), environment_id.clone(), req) + move |req: Request| { + let request_context = + req.extensions().get::().cloned(); + let policy_decider = policy_decider.clone(); + async move { + match request_context { + Some(request_context) => { + http_connect_accept(policy_decider, request_context, req).await + } + None => Err(text_response( + StatusCode::INTERNAL_SERVER_ERROR, + "missing request context", + )), + } + } } }), service_fn(http_connect_proxy), @@ -147,22 +177,38 @@ async fn run_http_proxy_with_listener( ) .into_layer(service_fn({ let policy_decider = policy_decider.clone(); - let environment_id = environment_id.clone(); - move |req| http_plain_proxy(policy_decider.clone(), environment_id.clone(), req) + move |req: Request| { + let request_context = req.extensions().get::().cloned(); + let policy_decider = policy_decider.clone(); + async move { + match request_context { + Some(request_context) => { + http_plain_proxy(policy_decider, request_context, req).await + } + None => Ok(text_response( + StatusCode::INTERNAL_SERVER_ERROR, + "missing request context", + )), + } + } + } })), ); info!("HTTP proxy listening on {addr}"); listener - .serve(AddInputExtensionLayer::new(state).into_layer(http_service)) + .serve(BindConnectionAttribution::new( + AddInputExtensionLayer::new(state).into_layer(http_service), + attribution, + )) .await; Ok(()) } async fn http_connect_accept( policy_decider: Option>, - environment_id: Option, + request_context: NetworkRequestContext, mut req: Request, ) -> Result<(Response, Request), Response> { let app_state = req @@ -194,12 +240,15 @@ async fn http_connect_accept( warn!("CONNECT blocked; proxy disabled (client={client}, host={host})"); return Err(proxy_disabled_response( &app_state, - host, - authority.port, - client_addr(&req), - Some("CONNECT".to_string()), - NetworkProtocol::HttpsConnect, - /*audit_endpoint_override*/ None, + ProxyDisabledResponseArgs { + host, + port: authority.port, + client: client_addr(&req), + method: Some("CONNECT".to_string()), + protocol: NetworkProtocol::HttpsConnect, + execution_id: request_context.execution_id.clone(), + audit_endpoint_override: None, + }, ) .await); } @@ -208,7 +257,8 @@ async fn http_connect_accept( protocol: NetworkProtocol::HttpsConnect, host: host.clone(), port: authority.port, - environment_id, + environment_id: request_context.environment_id.clone(), + execution_id: request_context.execution_id.clone(), client_addr: client.clone(), method: Some("CONNECT".to_string()), command: None, @@ -237,6 +287,7 @@ async fn http_connect_accept( method: Some("CONNECT".to_string()), mode: None, protocol: "http-connect".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(authority.port), @@ -308,6 +359,7 @@ async fn http_connect_accept( method: Some("CONNECT".to_string()), mode: Some(mode), protocol: "http-connect".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(authority.port), @@ -324,6 +376,7 @@ async fn http_connect_accept( req.extensions_mut() .insert(ConnectMitmEnabled(connect_needs_mitm)); req.extensions_mut().insert(mode); + req.extensions_mut().insert(request_context); if connect_needs_mitm && let Some(mitm_state) = mitm_state { req.extensions_mut().insert(mitm_state); } @@ -487,7 +540,7 @@ async fn forward_connect_tunnel( async fn http_plain_proxy( policy_decider: Option>, - environment_id: Option, + request_context: NetworkRequestContext, mut req: Request, ) -> Result { let app_state = match req.extensions().get::>().cloned() { @@ -534,12 +587,15 @@ async fn http_plain_proxy( warn!("unix socket blocked; proxy disabled (client={client}, path={socket_path})"); return Ok(proxy_disabled_response( &app_state, - socket_path, - /*port*/ 0, - client_addr(&req), - Some(req.method().as_str().to_string()), - NetworkProtocol::Http, - Some(("unix-socket", 0)), + ProxyDisabledResponseArgs { + host: socket_path, + port: 0, + client: client_addr(&req), + method: Some(req.method().as_str().to_string()), + protocol: NetworkProtocol::Http, + execution_id: request_context.execution_id.clone(), + audit_endpoint_override: Some(("unix-socket", 0)), + }, ) .await); } @@ -679,12 +735,15 @@ async fn http_plain_proxy( warn!("request blocked; proxy disabled (client={client}, host={host}, method={method})"); return Ok(proxy_disabled_response( &app_state, - host, - port, - client_addr(&req), - Some(req.method().as_str().to_string()), - NetworkProtocol::Http, - /*audit_endpoint_override*/ None, + ProxyDisabledResponseArgs { + host, + port, + client: client_addr(&req), + method: Some(req.method().as_str().to_string()), + protocol: NetworkProtocol::Http, + execution_id: request_context.execution_id.clone(), + audit_endpoint_override: None, + }, ) .await); } @@ -693,7 +752,8 @@ async fn http_plain_proxy( protocol: NetworkProtocol::Http, host: host.clone(), port, - environment_id, + environment_id: request_context.environment_id.clone(), + execution_id: request_context.execution_id.clone(), client_addr: client.clone(), method: Some(req.method().as_str().to_string()), command: None, @@ -722,6 +782,7 @@ async fn http_plain_proxy( method: Some(req.method().as_str().to_string()), mode: None, protocol: "http".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -767,6 +828,7 @@ async fn http_plain_proxy( method: Some(req.method().as_str().to_string()), mode: Some(NetworkMode::Limited), protocol: "http".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -951,15 +1013,29 @@ fn blocked_text_with_details(reason: &str, details: &PolicyDecisionDetails<'_>) blocked_text_response_with_policy(reason, details) } -async fn proxy_disabled_response( - app_state: &NetworkProxyState, +struct ProxyDisabledResponseArgs<'a> { host: String, port: u16, client: Option, method: Option, protocol: NetworkProtocol, - audit_endpoint_override: Option<(&str, u16)>, + execution_id: Option, + audit_endpoint_override: Option<(&'a str, u16)>, +} + +async fn proxy_disabled_response( + app_state: &NetworkProxyState, + args: ProxyDisabledResponseArgs<'_>, ) -> Response { + let ProxyDisabledResponseArgs { + host, + port, + client, + method, + protocol, + execution_id, + audit_endpoint_override, + } = args; let (audit_server_address, audit_server_port) = audit_endpoint_override.unwrap_or((host.as_str(), port)); emit_http_block_decision_audit_event( @@ -984,6 +1060,7 @@ async fn proxy_disabled_response( method, mode: None, protocol: protocol.as_policy_protocol().to_string(), + execution_id, decision: Some("deny".to_string()), source: Some("proxy_state".to_string()), port: Some(port), @@ -1087,7 +1164,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_connect_accept( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap_err(); @@ -1119,7 +1198,9 @@ mod tests { req.extensions_mut().insert(state); let (response, _request) = http_connect_accept( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap(); @@ -1127,17 +1208,18 @@ mod tests { } #[tokio::test] - async fn http_connect_accept_passes_environment_id_to_decider() { + async fn http_connect_accept_passes_request_context_to_decider() { let state = Arc::new(network_proxy_state_for_policy( NetworkProxySettings::default(), )); - let seen_environment_id = Arc::new(Mutex::new(None)); + let seen_request_context = Arc::new(Mutex::new(None)); let decider: Arc = Arc::new({ - let seen_environment_id = seen_environment_id.clone(); + let seen_request_context = seen_request_context.clone(); move |request: NetworkPolicyRequest| { - *seen_environment_id + *seen_request_context .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = request.environment_id; + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some((request.environment_id, request.execution_id)); async { NetworkDecision::Allow } } }); @@ -1150,18 +1232,24 @@ mod tests { .unwrap(); req.extensions_mut().insert(state); - let (response, _request) = - http_connect_accept(Some(decider), Some("remote".to_string()), req) - .await - .unwrap(); + let (response, _request) = http_connect_accept( + Some(decider), + NetworkRequestContext { + environment_id: Some("remote".to_string()), + execution_id: Some("execution-1".to_string()), + }, + req, + ) + .await + .unwrap(); assert_eq!(response.status(), StatusCode::OK); assert_eq!( - seen_environment_id + seen_request_context .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .as_deref(), - Some("remote") + .as_ref(), + Some(&(Some("remote".to_string()), Some("execution-1".to_string()))) ); } @@ -1192,7 +1280,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_connect_accept( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap_err(); @@ -1232,7 +1322,10 @@ mod tests { .local_addr() .expect("proxy listener should expose local addr"); let proxy_task = tokio::spawn(run_http_proxy_with_std_listener( - state, listener, /*policy_decider*/ None, /*environment_id*/ None, + state, + listener, + /*policy_decider*/ None, + NetworkRequestContext::default(), )); let mut stream = tokio::net::TcpStream::connect(proxy_addr) @@ -1284,7 +1377,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_plain_proxy( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap(); @@ -1311,7 +1406,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_plain_proxy( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap(); @@ -1345,7 +1442,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_plain_proxy( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap(); @@ -1371,7 +1470,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_connect_accept( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await .unwrap_err(); @@ -1396,7 +1497,9 @@ mod tests { req.extensions_mut().insert(state); let response = http_plain_proxy( - /*policy_decider*/ None, /*environment_id*/ None, req, + /*policy_decider*/ None, + NetworkRequestContext::default(), + req, ) .await; assert_eq!(response.unwrap().status(), StatusCode::BAD_REQUEST); diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index c8950f1f5592..fbacf7029033 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -1,5 +1,6 @@ #![deny(clippy::print_stdout, clippy::print_stderr)] +mod attribution; mod certs; mod config; mod connect_policy; @@ -17,6 +18,8 @@ mod socks5; mod state; mod upstream; +pub use attribution::PROXY_ATTRIBUTION_TOKEN_ENV_KEY; +pub use attribution::write_attribution_frame; pub use certs::CUSTOM_CA_ENV_KEYS; pub use certs::is_managed_mitm_ca_trust_bundle_path; pub use config::NetworkDomainPermission; diff --git a/codex-rs/network-proxy/src/mitm.rs b/codex-rs/network-proxy/src/mitm.rs index 70175ca70602..80e756e36cf0 100644 --- a/codex-rs/network-proxy/src/mitm.rs +++ b/codex-rs/network-proxy/src/mitm.rs @@ -2,6 +2,7 @@ use crate::certs::ManagedMitmCa; use crate::config::NetworkMode; use crate::mitm_hook::HookEvaluation; use crate::mitm_hook::MitmHookActions; +use crate::network_policy::NetworkRequestContext; use crate::policy::normalize_host; use crate::reasons::REASON_METHOD_NOT_ALLOWED; use crate::reasons::REASON_MITM_HOOK_DENIED; @@ -70,6 +71,7 @@ struct MitmPolicyContext { target_port: u16, mode: NetworkMode, app_state: Arc, + request_context: NetworkRequestContext, } #[derive(Clone)] @@ -177,12 +179,18 @@ where .get::() .copied() .unwrap_or(NetworkMode::Full); + let network_request_context = stream + .extensions() + .get::() + .cloned() + .unwrap_or_default(); let request_ctx = Arc::new(MitmRequestContext { policy: MitmPolicyContext { target_host, target_port, mode, app_state, + request_context: network_request_context, }, mitm, }); @@ -344,6 +352,7 @@ async fn evaluate_mitm_policy( method: Some(method.clone()), mode: Some(policy.mode), protocol: "https".to_string(), + execution_id: policy.request_context.execution_id.clone(), decision: None, source: None, port: Some(policy.target_port), @@ -372,6 +381,7 @@ async fn evaluate_mitm_policy( method: Some(method.clone()), mode: Some(policy.mode), protocol: "https".to_string(), + execution_id: policy.request_context.execution_id.clone(), decision: None, source: None, port: Some(policy.target_port), @@ -398,6 +408,7 @@ async fn evaluate_mitm_policy( method: Some(method.clone()), mode: Some(policy.mode), protocol: "https".to_string(), + execution_id: policy.request_context.execution_id.clone(), decision: None, source: None, port: Some(policy.target_port), diff --git a/codex-rs/network-proxy/src/mitm_tests.rs b/codex-rs/network-proxy/src/mitm_tests.rs index 0823be517c9d..2c130aac172e 100644 --- a/codex-rs/network-proxy/src/mitm_tests.rs +++ b/codex-rs/network-proxy/src/mitm_tests.rs @@ -47,6 +47,7 @@ fn policy_ctx( target_port, mode, app_state, + request_context: crate::network_policy::NetworkRequestContext::default(), } } diff --git a/codex-rs/network-proxy/src/network_policy.rs b/codex-rs/network-proxy/src/network_policy.rs index e5792a07c0ee..30cb990087aa 100644 --- a/codex-rs/network-proxy/src/network_policy.rs +++ b/codex-rs/network-proxy/src/network_policy.rs @@ -19,6 +19,21 @@ const POLICY_REASON_ALLOW: &str = "allow"; const DEFAULT_METHOD: &str = "none"; const DEFAULT_CLIENT_ADDRESS: &str = "unknown"; +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub(crate) struct NetworkRequestContext { + pub environment_id: Option, + pub execution_id: Option, +} + +impl NetworkRequestContext { + pub(crate) fn for_environment(environment_id: impl Into) -> Self { + Self { + environment_id: Some(environment_id.into()), + execution_id: None, + } + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum NetworkProtocol { Http, @@ -80,6 +95,7 @@ pub struct NetworkPolicyRequest { pub host: String, pub port: u16, pub environment_id: Option, + pub execution_id: Option, pub client_addr: Option, pub method: Option, pub command: Option, @@ -91,6 +107,7 @@ pub struct NetworkPolicyRequestArgs { pub host: String, pub port: u16, pub environment_id: Option, + pub execution_id: Option, pub client_addr: Option, pub method: Option, pub command: Option, @@ -104,6 +121,7 @@ impl NetworkPolicyRequest { host, port, environment_id, + execution_id, client_addr, method, command, @@ -114,6 +132,7 @@ impl NetworkPolicyRequest { host, port, environment_id, + execution_id, client_addr, method, command, @@ -211,6 +230,7 @@ fn emit_non_domain_policy_decision_audit_event( server_port: args.server_port, method: args.method, client_addr: args.client_addr, + execution_id: None, policy_override: false, }, ); @@ -226,6 +246,7 @@ struct PolicyAuditEventArgs<'a> { server_port: u16, method: Option<&'a str>, client_addr: Option<&'a str>, + execution_id: Option<&'a str>, policy_override: bool, } @@ -254,6 +275,7 @@ fn emit_policy_audit_event(state: &NetworkProxyState, args: PolicyAuditEventArgs server.port = args.server_port, http.request.method = args.method.unwrap_or(DEFAULT_METHOD), client.address = args.client_addr.unwrap_or(DEFAULT_CLIENT_ADDRESS), + execution.id = args.execution_id, network.policy.override = args.policy_override, ); } @@ -355,6 +377,7 @@ pub(crate) async fn evaluate_host_policy( server_port: request.port, method: request.method.as_deref(), client_addr: request.client_addr.as_deref(), + execution_id: request.execution_id.as_deref(), policy_override, }, ); @@ -630,6 +653,7 @@ mod tests { host: "example.com".to_string(), port: 80, environment_id: None, + execution_id: Some("execution-1".to_string()), client_addr: None, method: None, command: None, @@ -664,6 +688,7 @@ mod tests { assert_eq!(event.field("server.port"), Some("80")); assert_eq!(event.field("http.request.method"), Some(DEFAULT_METHOD)); assert_eq!(event.field("client.address"), Some(DEFAULT_CLIENT_ADDRESS)); + assert_eq!(event.field("execution.id"), Some("execution-1")); assert_eq!(event.field("network.policy.override"), Some("true")); let timestamp = event .field("event.timestamp") @@ -692,6 +717,7 @@ mod tests { host: "blocked.com".to_string(), port: 80, environment_id: None, + execution_id: None, client_addr: Some("127.0.0.1:1234".to_string()), method: Some("GET".to_string()), command: None, @@ -736,6 +762,7 @@ mod tests { host: "example.com".to_string(), port: 80, environment_id: None, + execution_id: None, client_addr: None, method: Some("GET".to_string()), command: None, @@ -787,6 +814,7 @@ mod tests { host: "example.com".to_string(), port: 80, environment_id: None, + execution_id: None, client_addr: None, method: Some("GET".to_string()), command: None, @@ -874,6 +902,7 @@ mod tests { host: "127.0.0.1".to_string(), port: 80, environment_id: None, + execution_id: None, client_addr: None, method: Some("GET".to_string()), command: None, diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index 13998dd447ee..8ef18ff6afb4 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -1,6 +1,9 @@ +use crate::attribution::AttributionRegistry; +use crate::attribution::ConnectionAttribution; use crate::config; use crate::http_proxy; use crate::network_policy::NetworkPolicyDecider; +use crate::network_policy::NetworkRequestContext; use crate::runtime::BlockedRequestObserver; use crate::runtime::ConfigState; use crate::runtime::unix_socket_permissions_supported; @@ -229,6 +232,7 @@ impl NetworkProxyBuilder { reserved_listeners, policy_decider: self.policy_decider, environment_proxies: Arc::new(Mutex::new(HashMap::new())), + execution_proxies: AttributionRegistry::default(), }) } } @@ -345,6 +349,7 @@ pub struct NetworkProxy { reserved_listeners: Option>, policy_decider: Option>, environment_proxies: Arc>>, + execution_proxies: AttributionRegistry, } impl std::fmt::Debug for NetworkProxy { @@ -739,13 +744,13 @@ impl NetworkProxy { let environment_id = environment_id.to_string(); let http_state = self.state.clone(); let http_decider = self.policy_decider.clone(); - let http_environment_id = Some(environment_id.clone()); + let http_request_context = NetworkRequestContext::for_environment(&environment_id); let http_task = runtime.spawn(async move { http_proxy::run_http_proxy_with_std_listener( http_state, http_listener, http_decider, - http_environment_id, + http_request_context, ) .await }); @@ -753,7 +758,7 @@ impl NetworkProxy { let socks_task = if self.socks_enabled { let socks_state = self.state.clone(); let socks_decider = self.policy_decider.clone(); - let socks_environment_id = Some(environment_id.clone()); + let socks_request_context = NetworkRequestContext::for_environment(&environment_id); let socks5_udp_enabled = self.socks5_udp_enabled; socks_listener.map(|listener| { runtime.spawn(async move { @@ -761,7 +766,7 @@ impl NetworkProxy { socks_state, listener, socks_decider, - socks_environment_id, + socks_request_context, socks5_udp_enabled, ) .await @@ -841,24 +846,25 @@ impl NetworkProxy { let http_state = self.state.clone(); let http_decider = self.policy_decider.clone(); + let http_attribution = ConnectionAttribution::Registry(self.execution_proxies.clone()); let http_addr = self.http_addr; let http_task = tokio::spawn(async move { match http_listener { Some(listener) => { - http_proxy::run_http_proxy_with_std_listener( + http_proxy::run_http_proxy_with_std_listener_and_attribution( http_state, listener, http_decider, - /*environment_id*/ None, + http_attribution, ) .await } None => { - http_proxy::run_http_proxy( + http_proxy::run_http_proxy_with_attribution( http_state, http_addr, http_decider, - /*environment_id*/ None, + http_attribution, ) .await } @@ -870,24 +876,25 @@ impl NetworkProxy { let socks_decider = self.policy_decider.clone(); let socks_addr = self.socks_addr; let enable_socks5_udp = current_cfg.network.enable_socks5_udp; + let socks_attribution = ConnectionAttribution::Registry(self.execution_proxies.clone()); Some(tokio::spawn(async move { match socks_listener { Some(listener) => { - socks5::run_socks5_with_std_listener( + socks5::run_socks5_with_std_listener_and_attribution( socks_state, listener, socks_decider, - /*environment_id*/ None, + socks_attribution, enable_socks5_udp, ) .await } None => { - socks5::run_socks5( + socks5::run_socks5_with_attribution( socks_state, socks_addr, socks_decider, - /*environment_id*/ None, + socks_attribution, enable_socks5_udp, ) .await @@ -902,6 +909,7 @@ impl NetworkProxy { http_task: Some(http_task), socks_task, environment_proxies: self.environment_proxies.clone(), + execution_proxies: self.execution_proxies.clone(), completed: false, }) } @@ -911,6 +919,7 @@ pub struct NetworkProxyHandle { http_task: Option>>, socks_task: Option>>, environment_proxies: Arc>>, + execution_proxies: AttributionRegistry, completed: bool, } @@ -920,6 +929,7 @@ impl NetworkProxyHandle { http_task: Some(tokio::spawn(async { Ok(()) })), socks_task: None, environment_proxies: Arc::new(Mutex::new(HashMap::new())), + execution_proxies: AttributionRegistry::default(), completed: true, } } @@ -934,6 +944,7 @@ impl NetworkProxyHandle { }; self.completed = true; abort_environment_proxies(self.environment_proxies.clone()).await; + self.execution_proxies.clear(); http_result??; if let Some(socks_result) = socks_result { socks_result??; @@ -944,6 +955,7 @@ impl NetworkProxyHandle { pub async fn shutdown(mut self) -> Result<()> { abort_tasks(self.http_task.take(), self.socks_task.take()).await; abort_environment_proxies(self.environment_proxies.clone()).await; + self.execution_proxies.clear(); self.completed = true; Ok(()) } @@ -987,6 +999,7 @@ impl Drop for NetworkProxyHandle { let http_task = self.http_task.take(); let socks_task = self.socks_task.take(); let environment_proxies = self.environment_proxies.clone(); + self.execution_proxies.clear(); tokio::spawn(async move { abort_tasks(http_task, socks_task).await; abort_environment_proxies(environment_proxies).await; @@ -1077,7 +1090,6 @@ mod tests { )); let proxy = NetworkProxy::builder().state(state).build().await?; let handle = proxy.run().await?; - let mut local_env = HashMap::new(); proxy.apply_to_env_for_environment(&mut local_env, "local")?; let mut remote_env = HashMap::new(); diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index e721d0de708e..e9b3af4c7304 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -95,6 +95,8 @@ pub struct BlockedRequest { pub mode: Option, pub protocol: String, #[serde(skip_serializing_if = "Option::is_none")] + pub execution_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub decision: Option, #[serde(skip_serializing_if = "Option::is_none")] pub source: Option, @@ -110,6 +112,7 @@ pub struct BlockedRequestArgs { pub method: Option, pub mode: Option, pub protocol: String, + pub execution_id: Option, pub decision: Option, pub source: Option, pub port: Option, @@ -124,6 +127,7 @@ impl BlockedRequest { method, mode, protocol, + execution_id, decision, source, port, @@ -135,6 +139,7 @@ impl BlockedRequest { method, mode, protocol, + execution_id, decision, source, port, @@ -1139,6 +1144,7 @@ mod tests { method: Some("GET".to_string()), mode: None, protocol: "http".to_string(), + execution_id: None, decision: Some("ask".to_string()), source: Some("decider".to_string()), port: Some(80), @@ -1179,6 +1185,7 @@ mod tests { method: Some("GET".to_string()), mode: None, protocol: "http".to_string(), + execution_id: None, decision: Some("ask".to_string()), source: Some("decider".to_string()), port: Some(80), @@ -1201,6 +1208,7 @@ mod tests { method: Some("GET".to_string()), mode: Some(NetworkMode::Full), protocol: "http".to_string(), + execution_id: None, decision: Some("ask".to_string()), source: Some("decider".to_string()), port: Some(80), diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index 5cc7b2eed39f..8387a7a455a8 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -1,3 +1,5 @@ +use crate::attribution::BindConnectionAttribution; +use crate::attribution::ConnectionAttribution; use crate::config::NetworkMode; use crate::connect_policy::TargetCheckedTcpConnector; use crate::mitm; @@ -9,6 +11,7 @@ use crate::network_policy::NetworkPolicyDecision; use crate::network_policy::NetworkPolicyRequest; use crate::network_policy::NetworkPolicyRequestArgs; use crate::network_policy::NetworkProtocol; +use crate::network_policy::NetworkRequestContext; use crate::network_policy::emit_block_decision_audit_event; use crate::network_policy::evaluate_host_policy; use crate::policy::normalize_host; @@ -60,11 +63,11 @@ use tracing::error; use tracing::info; use tracing::warn; -pub async fn run_socks5( +pub(crate) async fn run_socks5_with_attribution( state: Arc, addr: SocketAddr, policy_decider: Option>, - environment_id: Option, + attribution: ConnectionAttribution, enable_socks5_udp: bool, ) -> Result<()> { let listener = TcpListener::build() @@ -79,7 +82,7 @@ pub async fn run_socks5( state, listener, policy_decider, - environment_id, + attribution, enable_socks5_udp, ) .await @@ -89,7 +92,24 @@ pub async fn run_socks5_with_std_listener( state: Arc, listener: StdTcpListener, policy_decider: Option>, - environment_id: Option, + request_context: NetworkRequestContext, + enable_socks5_udp: bool, +) -> Result<()> { + run_socks5_with_std_listener_and_attribution( + state, + listener, + policy_decider, + ConnectionAttribution::Fixed(request_context), + enable_socks5_udp, + ) + .await +} + +pub(crate) async fn run_socks5_with_std_listener_and_attribution( + state: Arc, + listener: StdTcpListener, + policy_decider: Option>, + attribution: ConnectionAttribution, enable_socks5_udp: bool, ) -> Result<()> { let listener = @@ -98,7 +118,7 @@ pub async fn run_socks5_with_std_listener( state, listener, policy_decider, - environment_id, + attribution, enable_socks5_udp, ) .await @@ -108,7 +128,7 @@ async fn run_socks5_with_listener( state: Arc, listener: TcpListener, policy_decider: Option>, - environment_id: Option, + attribution: ConnectionAttribution, enable_socks5_udp: bool, ) -> Result<()> { let addr = listener @@ -132,12 +152,15 @@ async fn run_socks5_with_listener( let tcp_connector = TargetCheckedTcpConnector::new(state.clone()); let policy_tcp_connector = service_fn({ let policy_decider = policy_decider.clone(); - let environment_id = environment_id.clone(); move |req: TcpRequest| { let tcp_connector = tcp_connector.clone(); let policy_decider = policy_decider.clone(); - let environment_id = environment_id.clone(); - async move { handle_socks5_tcp(req, tcp_connector, policy_decider, environment_id).await } + let request_context = req.extensions().get::().cloned(); + async move { + let request_context = request_context + .ok_or_else(|| io::Error::other("missing network request context"))?; + handle_socks5_tcp(req, tcp_connector, policy_decider, request_context).await + } } }); @@ -150,25 +173,31 @@ async fn run_socks5_with_listener( if enable_socks5_udp { let udp_state = state.clone(); let udp_decider = policy_decider.clone(); - let udp_relay = - DefaultUdpRelay::default().with_async_inspector(service_fn({ - let environment_id = environment_id.clone(); - move |request: RelayRequest| { - let udp_state = udp_state.clone(); - let udp_decider = udp_decider.clone(); - let environment_id = environment_id.clone(); - async move { - inspect_socks5_udp(request, udp_state, udp_decider, environment_id).await - } + let udp_relay = DefaultUdpRelay::default().with_async_inspector(service_fn({ + move |request: RelayRequest| { + let udp_state = udp_state.clone(); + let udp_decider = udp_decider.clone(); + let request_context = request.extensions.get::().cloned(); + async move { + let request_context = request_context + .ok_or_else(|| io::Error::other("missing network request context"))?; + inspect_socks5_udp(request, udp_state, udp_decider, request_context).await } - })); + } + })); let socks_acceptor = base.with_udp_associator(udp_relay); listener - .serve(AddInputExtensionLayer::new(state).into_layer(socks_acceptor)) + .serve(BindConnectionAttribution::new( + AddInputExtensionLayer::new(state).into_layer(socks_acceptor), + attribution, + )) .await; } else { listener - .serve(AddInputExtensionLayer::new(state).into_layer(base)) + .serve(BindConnectionAttribution::new( + AddInputExtensionLayer::new(state).into_layer(base), + attribution, + )) .await; } Ok(()) @@ -178,7 +207,7 @@ async fn handle_socks5_tcp( req: TcpRequest, tcp_connector: TargetCheckedTcpConnector, policy_decider: Option>, - environment_id: Option, + request_context: NetworkRequestContext, ) -> Result, BoxError> { let app_state = req .extensions() @@ -226,6 +255,7 @@ async fn handle_socks5_tcp( method: None, mode: None, protocol: "socks5".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -277,6 +307,7 @@ async fn handle_socks5_tcp( method: None, mode: Some(NetworkMode::Limited), protocol: "socks5".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -293,7 +324,8 @@ async fn handle_socks5_tcp( protocol: NetworkProtocol::Socks5Tcp, host: host.clone(), port, - environment_id, + environment_id: request_context.environment_id.clone(), + execution_id: request_context.execution_id.clone(), client_addr: client.clone(), method: None, command: None, @@ -322,6 +354,7 @@ async fn handle_socks5_tcp( method: None, mode: None, protocol: "socks5".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -385,6 +418,7 @@ async fn handle_socks5_tcp( method: None, mode: Some(mode), protocol: "socks5".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -400,13 +434,15 @@ async fn handle_socks5_tcp( if socks_needs_mitm && let Some(mitm_state) = mitm_state { let client = client.as_deref().unwrap_or_default(); info!("SOCKS MITM enabled (client={client}, host={host}, port={port}, mode={mode:?})"); + let mut extensions = Extensions::new(); + extensions.insert(request_context); return Ok(EstablishedClientConnection { input: req, conn: Socks5TcpConnection::Mitm { target, mode, mitm: mitm_state, - extensions: Extensions::new(), + extensions, }, }); } @@ -530,11 +566,17 @@ async fn proxy_socks5_tcp( .await .map_err(Into::into), Socks5TcpConnection::Mitm { - target, mode, mitm, .. + target, + mode, + mitm, + extensions, } => { source.extensions_mut().insert(ProxyTarget(target)); source.extensions_mut().insert(mode); source.extensions_mut().insert(mitm); + if let Some(request_context) = extensions.get::().cloned() { + source.extensions_mut().insert(request_context); + } mitm::mitm_stream(source).await.map_err(Into::into) } } @@ -544,7 +586,7 @@ async fn inspect_socks5_udp( request: RelayRequest, state: Arc, policy_decider: Option>, - environment_id: Option, + request_context: NetworkRequestContext, ) -> io::Result { let RelayRequest { server_address, @@ -591,6 +633,7 @@ async fn inspect_socks5_udp( method: None, mode: None, protocol: "socks5-udp".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -633,6 +676,7 @@ async fn inspect_socks5_udp( method: None, mode: Some(NetworkMode::Limited), protocol: "socks5-udp".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -651,7 +695,8 @@ async fn inspect_socks5_udp( protocol: NetworkProtocol::Socks5Udp, host: host.clone(), port, - environment_id, + environment_id: request_context.environment_id.clone(), + execution_id: request_context.execution_id.clone(), client_addr: client.clone(), method: None, command: None, @@ -680,6 +725,7 @@ async fn inspect_socks5_udp( method: None, mode: None, protocol: "socks5-udp".to_string(), + execution_id: request_context.execution_id.clone(), decision: Some(details.decision.as_str().to_string()), source: Some(details.source.as_str().to_string()), port: Some(port), @@ -809,7 +855,7 @@ mod tests { request, TargetCheckedTcpConnector::new(state.clone()), /*policy_decider*/ None, - /*environment_id*/ None, + NetworkRequestContext::default(), ) .await }) @@ -853,7 +899,7 @@ mod tests { request, TargetCheckedTcpConnector::new(state), /*policy_decider*/ None, - /*environment_id*/ None, + NetworkRequestContext::default(), ) .await .expect("limited-mode HTTPS should use MITM"); @@ -879,7 +925,7 @@ mod tests { request, TargetCheckedTcpConnector::new(state), /*policy_decider*/ None, - /*environment_id*/ None, + NetworkRequestContext::default(), ) .await }) @@ -925,7 +971,7 @@ mod tests { request, TargetCheckedTcpConnector::new(state), /*policy_decider*/ None, - /*environment_id*/ None, + NetworkRequestContext::default(), ) .await .expect_err("limited-mode HTTPS requires MITM"); @@ -963,7 +1009,7 @@ mod tests { request, TargetCheckedTcpConnector::new(state), /*policy_decider*/ None, - /*environment_id*/ None, + NetworkRequestContext::default(), ) .await .expect("hooked HTTPS should use MITM"); @@ -998,7 +1044,7 @@ mod tests { request, TargetCheckedTcpConnector::new(state), /*policy_decider*/ None, - /*environment_id*/ None, + NetworkRequestContext::default(), ) .await .expect_err("hooked non-HTTPS SOCKS should require MITM"); @@ -1025,7 +1071,10 @@ mod tests { let (result, events) = capture_events(|| async { inspect_socks5_udp( - request, state, /*policy_decider*/ None, /*environment_id*/ None, + request, + state, + /*policy_decider*/ None, + NetworkRequestContext::default(), ) .await })