diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 023c5dfaa366..dc109aeabadb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3620,6 +3620,7 @@ dependencies = [ "codex-utils-home-dir", "codex-utils-rustls-provider", "globset", + "hickory-proto", "pretty_assertions", "rama-core", "rama-http", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 7b153dfaa0b7..bc4330731d17 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -322,6 +322,7 @@ gethostname = "1.1.0" gix = { version = "0.81.0", default-features = false, features = ["sha1"] } glob = "0.3" globset = "0.4" +hickory-proto = { version = "0.25.2", default-features = false, features = ["std"] } hmac = "0.12.1" http = "1.3.1" httpdate = "1.0.3" diff --git a/codex-rs/network-proxy/Cargo.toml b/codex-rs/network-proxy/Cargo.toml index d1d2a72ce323..1e7152fde99f 100644 --- a/codex-rs/network-proxy/Cargo.toml +++ b/codex-rs/network-proxy/Cargo.toml @@ -21,6 +21,7 @@ codex-utils-absolute-path = { workspace = true } codex-utils-home-dir = { workspace = true } codex-utils-rustls-provider = { workspace = true } globset = { workspace = true } +hickory-proto = { workspace = true } rand = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } diff --git a/codex-rs/network-proxy/src/attribution.rs b/codex-rs/network-proxy/src/attribution.rs index f526460fd622..f9260473f6da 100644 --- a/codex-rs/network-proxy/src/attribution.rs +++ b/codex-rs/network-proxy/src/attribution.rs @@ -72,16 +72,19 @@ where } 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); - } + Ok(tokio::time::timeout(ATTRIBUTION_FRAME_TIMEOUT, async { + 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", + )); + } + 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 { @@ -100,12 +103,13 @@ async fn read_attribution_token(stream: &mut TcpStream) -> Result } let mut token = vec![0_u8; token_len]; stream.read_exact(&mut token).await?; - String::from_utf8(token).map_err(|_| { + let token = String::from_utf8(token).map_err(|_| { io::Error::new( io::ErrorKind::InvalidData, "network proxy attribution token is not UTF-8", ) - }) + })?; + Ok(Some(token)) }) .await .map_err(|_| { @@ -113,9 +117,7 @@ async fn read_attribution_token(stream: &mut TcpStream) -> Result io::ErrorKind::TimedOut, "network proxy attribution frame timed out", ) - })??; - - Ok(Some(token)) + })??) } /// Writes the trusted bridge preface consumed by the shared proxy ingress. diff --git a/codex-rs/network-proxy/src/dns.rs b/codex-rs/network-proxy/src/dns.rs new file mode 100644 index 000000000000..7f424e02a2b4 --- /dev/null +++ b/codex-rs/network-proxy/src/dns.rs @@ -0,0 +1,364 @@ +use crate::attribution::BindConnectionAttribution; +use crate::network_policy::NetworkDecision; +use crate::network_policy::NetworkDecisionSource; +use crate::network_policy::NetworkPolicyDecision; +use crate::network_policy::emit_dns_policy_decision_audit_event; +use crate::policy::is_non_public_ip; +use crate::policy::normalize_host; +use crate::reasons::REASON_PROXY_DISABLED; +use crate::runtime::HostBlockReason; +use crate::runtime::HostResolutionPolicy; +use crate::state::BlockedRequest; +use crate::state::BlockedRequestArgs; +use crate::state::NetworkProxyState; +use anyhow::Context; +use anyhow::Result; +use hickory_proto::op::Message; +use hickory_proto::op::MessageType; +use hickory_proto::op::OpCode; +use hickory_proto::op::ResponseCode; +use hickory_proto::rr::DNSClass; +use hickory_proto::rr::RData; +use hickory_proto::rr::Record; +use hickory_proto::rr::RecordType; +use hickory_proto::rr::rdata::A; +use hickory_proto::rr::rdata::AAAA; +use rama_core::Service; +use rama_core::error::BoxError; +use rama_core::extensions::ExtensionsRef; +use rama_net::stream::SocketInfo; +use rama_tcp::TcpStream; +use rama_tcp::server::TcpListener; +use std::collections::HashSet; +use std::future::Future; +use std::io; +use std::net::IpAddr; +use std::net::TcpListener as StdTcpListener; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::lookup_host; +use tokio::sync::Semaphore; +use tokio::time::timeout; +use tracing::info; + +/// Internal endpoint handed to the trusted Linux proxy bridge. +#[doc(hidden)] +pub const DNS_PROXY_ENV_KEY: &str = "CODEX_NETWORK_PROXY_DNS"; + +/// Versioned application preface for the private DNS stream. +/// +/// The nonzero first byte distinguishes an unattributed application stream from the attribution +/// frame, whose magic intentionally starts with a NUL byte. Each session exchanges this preface, +/// one length-prefixed query, and one length-prefixed response before closing. +#[doc(hidden)] +pub const DNS_PROXY_SESSION_PREFACE: &[u8; 8] = b"CDXDNS1\n"; + +const DNS_PORT: u16 = 53; +const MAX_DNS_QUERY_BYTES: usize = 4 * 1024; +const MAX_DNS_RESPONSE_BYTES: usize = 512; +const MAX_DNS_ANSWERS: usize = 16; +const DNS_LOOKUP_TIMEOUT: Duration = Duration::from_secs(2); +const DNS_SESSION_PREFACE_TIMEOUT: Duration = Duration::from_secs(3); +const DNS_FRAME_IO_TIMEOUT: Duration = Duration::from_secs(3); +const MAX_CONCURRENT_DNS_SESSIONS: usize = 64; +const RESPONSE_TTL: u32 = 0; + +pub(crate) async fn run_dns_proxy_with_std_listener( + state: Arc, + listener: StdTcpListener, + environment_id: Option, +) -> Result<()> { + let listener = + TcpListener::try_from(listener).context("convert std listener to DNS proxy listener")?; + let addr = listener + .local_addr() + .context("read DNS proxy listener local addr")?; + info!("DNS proxy listening on {addr}"); + listener + .serve(LimitDnsSessions { + inner: BindConnectionAttribution::new(DnsService, state, environment_id), + sessions: Arc::new(Semaphore::new(MAX_CONCURRENT_DNS_SESSIONS)), + }) + .await; + Ok(()) +} + +struct LimitDnsSessions { + inner: S, + sessions: Arc, +} + +impl Service for LimitDnsSessions +where + S: Service, + S::Error: Into, +{ + type Output = S::Output; + type Error = BoxError; + + async fn serve(&self, stream: TcpStream) -> std::result::Result { + let _session = Arc::clone(&self.sessions) + .try_acquire_owned() + .map_err(|_| io::Error::other("DNS proxy session limit reached"))?; + self.inner.serve(stream).await.map_err(Into::into) + } +} + +#[derive(Clone, Copy)] +struct DnsService; + +impl Service for DnsService { + type Output = (); + type Error = BoxError; + + async fn serve(&self, mut stream: TcpStream) -> std::result::Result<(), BoxError> { + let state = stream + .extensions() + .get::>() + .cloned() + .ok_or_else(|| io::Error::other("missing network proxy state"))?; + let client_addr = stream + .extensions() + .get::() + .map(|info| info.peer_addr().to_string()); + + exchange_session_preface(&mut stream).await?; + let query = read_query_frame(&mut stream).await?; + let response = resolve_query(&state, &query, client_addr.as_deref()).await; + write_response_frame(&mut stream, &response).await?; + Ok(()) + } +} + +async fn exchange_session_preface(stream: &mut TcpStream) -> io::Result<()> { + timeout(DNS_SESSION_PREFACE_TIMEOUT, async { + let mut preface = [0_u8; DNS_PROXY_SESSION_PREFACE.len()]; + stream.read_exact(&mut preface).await?; + if &preface != DNS_PROXY_SESSION_PREFACE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid DNS proxy session preface", + )); + } + stream.write_all(DNS_PROXY_SESSION_PREFACE).await + }) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "DNS proxy preface timed out"))? +} + +async fn read_query_frame(stream: &mut TcpStream) -> io::Result> { + timeout(DNS_FRAME_IO_TIMEOUT, async { + let mut len = [0_u8; 2]; + stream.read_exact(&mut len).await?; + let len = usize::from(u16::from_be_bytes(len)); + if len == 0 || len > MAX_DNS_QUERY_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid DNS proxy query length", + )); + } + let mut payload = vec![0_u8; len]; + stream.read_exact(&mut payload).await?; + Ok(payload) + }) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "DNS query frame timed out"))? +} + +async fn write_response_frame(stream: &mut TcpStream, payload: &[u8]) -> io::Result<()> { + let len = u16::try_from(payload.len()) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "DNS response is too large"))?; + timeout(DNS_FRAME_IO_TIMEOUT, async { + stream.write_all(&len.to_be_bytes()).await?; + stream.write_all(payload).await + }) + .await + .map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "DNS response frame timed out"))? +} + +async fn resolve_query( + state: &NetworkProxyState, + wire: &[u8], + client_addr: Option<&str>, +) -> Vec { + resolve_query_with_lookup(state, wire, client_addr, |host| async move { + lookup_host((host.as_str(), 0)) + .await + .map(|addresses| addresses.map(|address| address.ip()).collect()) + }) + .await +} + +async fn resolve_query_with_lookup( + state: &NetworkProxyState, + wire: &[u8], + client_addr: Option<&str>, + lookup: F, +) -> Vec +where + F: FnOnce(String) -> Fut, + Fut: Future>>, +{ + let parsed = Message::from_vec(wire).ok(); + let Some(query) = parsed.as_ref().filter(|query| { + query.message_type() == MessageType::Query + && query.op_code() == OpCode::Query + && query.queries().len() == 1 + && query.queries()[0].query_class() == DNSClass::IN + && matches!( + query.queries()[0].query_type(), + RecordType::A | RecordType::AAAA + ) + }) else { + return bounded_response(refused_response(wire, parsed.as_ref())); + }; + let question = &query.queries()[0]; + let host = normalize_host(&question.name().to_utf8()); + match state.enabled().await { + Ok(true) => {} + Ok(false) => { + record_blocked_dns_request( + state, + &host, + REASON_PROXY_DISABLED, + NetworkDecisionSource::ProxyState, + client_addr, + ) + .await; + return bounded_response(refused_response(wire, Some(query))); + } + Err(_) => return bounded_response(response_message(query, ResponseCode::ServFail)), + } + let policy = match state.host_resolution_policy(&host).await { + Ok(policy) => policy, + Err(_) => return bounded_response(response_message(query, ResponseCode::ServFail)), + }; + let allow_non_public_ips = match policy { + HostResolutionPolicy::Allowed { + allow_non_public_ips, + } => allow_non_public_ips, + HostResolutionPolicy::Blocked(reason) => { + record_blocked_dns_request( + state, + &host, + reason.as_str(), + NetworkDecisionSource::BaselinePolicy, + client_addr, + ) + .await; + return bounded_response(refused_response(wire, Some(query))); + } + }; + + let addresses = match timeout(DNS_LOOKUP_TIMEOUT, lookup(host.clone())).await { + Ok(Ok(addresses)) => addresses, + Ok(Err(_)) | Err(_) => { + return bounded_response(response_message(query, ResponseCode::ServFail)); + } + }; + if !allow_non_public_ips && addresses.iter().copied().any(is_non_public_ip) { + record_blocked_dns_request( + state, + &host, + HostBlockReason::NotAllowedLocal.as_str(), + NetworkDecisionSource::BaselinePolicy, + client_addr, + ) + .await; + return bounded_response(refused_response(wire, Some(query))); + } + + emit_dns_policy_decision_audit_event(state, &host, &NetworkDecision::Allow, client_addr); + bounded_answer_response(query, question.query_type(), addresses) +} + +async fn record_blocked_dns_request( + state: &NetworkProxyState, + host: &str, + reason: &str, + source: NetworkDecisionSource, + client_addr: Option<&str>, +) { + let decision = NetworkDecision::deny_with_source(reason, source); + emit_dns_policy_decision_audit_event(state, host, &decision, client_addr); + let _ = state + .record_blocked(BlockedRequest::new(BlockedRequestArgs { + host: host.to_string(), + reason: reason.to_string(), + client: client_addr.map(str::to_string), + method: None, + mode: None, + protocol: "dns".to_string(), + decision: Some(NetworkPolicyDecision::Deny.as_str().to_string()), + source: Some(source.as_str().to_string()), + port: Some(DNS_PORT), + })) + .await; +} + +fn bounded_answer_response( + query: &Message, + record_type: RecordType, + addresses: Vec, +) -> Vec { + let mut response = response_message(query, ResponseCode::NoError); + let answer_name = query.queries()[0].name().clone(); + let mut seen = HashSet::new(); + for data in addresses + .into_iter() + .filter(|address| seen.insert(*address)) + .filter_map(|address| match (record_type, address) { + (RecordType::A, IpAddr::V4(ip)) => Some(RData::A(A(ip))), + (RecordType::AAAA, IpAddr::V6(ip)) => Some(RData::AAAA(AAAA(ip))), + _ => None, + }) + .take(MAX_DNS_ANSWERS) + { + response.add_answer(Record::from_rdata(answer_name.clone(), RESPONSE_TTL, data)); + if !response + .to_vec() + .is_ok_and(|wire| wire.len() <= MAX_DNS_RESPONSE_BYTES) + { + response.answers_mut().pop(); + break; + } + } + bounded_response(response) +} + +fn response_message(query: &Message, code: ResponseCode) -> Message { + let mut response = Message::error_msg(query.id(), query.op_code(), code); + response + .set_recursion_desired(query.recursion_desired()) + .set_recursion_available(true) + .add_queries(query.queries().iter().take(1).cloned()); + response +} + +fn refused_response(wire: &[u8], query: Option<&Message>) -> Message { + if let Some(query) = query { + return response_message(query, ResponseCode::Refused); + } + let id = wire + .get(..2) + .map(|bytes| u16::from_be_bytes([bytes[0], bytes[1]])) + .unwrap_or_default(); + Message::error_msg(id, OpCode::Query, ResponseCode::Refused) +} + +fn bounded_response(message: Message) -> Vec { + let id = message.id(); + let op_code = message.op_code(); + match message.to_vec() { + Ok(wire) if wire.len() <= MAX_DNS_RESPONSE_BYTES => wire, + _ => Message::error_msg(id, op_code, ResponseCode::ServFail) + .to_vec() + .unwrap_or_default(), + } +} + +#[cfg(test)] +#[path = "dns_tests.rs"] +mod tests; diff --git a/codex-rs/network-proxy/src/dns_tests.rs b/codex-rs/network-proxy/src/dns_tests.rs new file mode 100644 index 000000000000..831a4bb5526b --- /dev/null +++ b/codex-rs/network-proxy/src/dns_tests.rs @@ -0,0 +1,436 @@ +use super::*; +use crate::attribution::write_attribution_frame; +use crate::config::NetworkProxyConfig; +use crate::network_policy::test_support::POLICY_DECISION_EVENT_NAME; +use crate::network_policy::test_support::capture_events; +use crate::network_policy::test_support::find_event_by_name; +use crate::runtime::network_proxy_state_for_policy; +use crate::state::NetworkProxyConstraints; +use crate::state::build_config_state; +use hickory_proto::op::Query; +use hickory_proto::rr::Name; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpStream as TokioTcpStream; + +fn settings(allowed_domains: &[&str], denied_domains: &[&str]) -> NetworkProxyConfig { + let mut settings = NetworkProxyConfig { + enabled: true, + allow_local_binding: false, + ..NetworkProxyConfig::default() + }; + settings.set_allowed_domains( + allowed_domains + .iter() + .map(|domain| (*domain).to_string()) + .collect(), + ); + settings.set_denied_domains( + denied_domains + .iter() + .map(|domain| (*domain).to_string()) + .collect(), + ); + settings +} + +fn query(name: &str, record_type: RecordType, id: u16) -> Vec { + let mut message = Message::new(); + message.set_id(id); + message.add_query(Query::query( + Name::from_ascii(name).expect("valid DNS name"), + record_type, + )); + message.to_vec().expect("serialize DNS query") +} + +fn response(wire: Vec) -> Message { + Message::from_vec(&wire).expect("parse DNS response") +} + +async fn query_listener( + addr: std::net::SocketAddr, + attribution_token: Option<&str>, + name: &str, +) -> io::Result { + let mut stream = TokioTcpStream::connect(addr).await?; + if let Some(token) = attribution_token { + let mut attribution = Vec::new(); + write_attribution_frame(&mut attribution, token)?; + stream.write_all(&attribution).await?; + } + stream.write_all(DNS_PROXY_SESSION_PREFACE).await?; + let mut ack = [0_u8; DNS_PROXY_SESSION_PREFACE.len()]; + stream.read_exact(&mut ack).await?; + if &ack != DNS_PROXY_SESSION_PREFACE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid DNS session acknowledgement", + )); + } + + let query = query(name, RecordType::A, /*id*/ 16); + let query_len = u16::try_from(query.len()).expect("bounded DNS query length"); + stream.write_all(&query_len.to_be_bytes()).await?; + stream.write_all(&query).await?; + let response_len = stream.read_u16().await?; + let mut wire = vec![0_u8; usize::from(response_len)]; + stream.read_exact(&mut wire).await?; + Message::from_vec(&wire).map_err(io::Error::other) +} + +#[test] +fn session_preface_cannot_be_confused_with_attribution_frame() { + assert_ne!(DNS_PROXY_SESSION_PREFACE[0], 0); + assert_eq!(DNS_PROXY_SESSION_PREFACE.len(), 8); +} + +#[tokio::test] +async fn resolves_allowed_a_and_aaaa_queries_with_bounded_answers() { + let state = network_proxy_state_for_policy(settings(&["example.test"], &[])); + let addresses = vec![ + "93.184.216.34".parse().expect("IPv4 address"), + "2001:4860:4860::8888".parse().expect("IPv6 address"), + ]; + + for (record_type, expected) in [ + (RecordType::A, addresses[0]), + (RecordType::AAAA, addresses[1]), + ] { + let wire = resolve_query_with_lookup( + &state, + &query("example.test", record_type, /*id*/ 7), + Some("127.0.0.1:1234"), + |_| { + let addresses = addresses.clone(); + async move { Ok(addresses) } + }, + ) + .await; + assert!(wire.len() <= MAX_DNS_RESPONSE_BYTES); + let response = response(wire); + assert_eq!(response.response_code(), ResponseCode::NoError); + assert!(response.answers().iter().any(|answer| match answer.data() { + RData::A(address) => IpAddr::V4(address.0) == expected, + RData::AAAA(address) => IpAddr::V6(address.0) == expected, + _ => false, + })); + } +} + +#[tokio::test] +async fn synthesized_response_never_exceeds_classic_udp_size() { + let state = network_proxy_state_for_policy(settings(&["many.example.test"], &[])); + let addresses = (1..=64) + .map(|suffix| { + format!("2001:4860::{suffix}") + .parse() + .expect("public IPv6 address") + }) + .collect::>(); + let wire = resolve_query_with_lookup( + &state, + &query("many.example.test", RecordType::AAAA, /*id*/ 7), + /*client_addr*/ None, + |_| async move { Ok(addresses) }, + ) + .await; + + assert!(wire.len() <= MAX_DNS_RESPONSE_BYTES); + let response = response(wire); + assert_eq!(response.response_code(), ResponseCode::NoError); + assert!(!response.answers().is_empty()); + assert!(response.answers().len() <= MAX_DNS_ANSWERS); +} + +#[tokio::test] +async fn denied_name_is_not_sent_to_the_host_resolver() { + let state = network_proxy_state_for_policy(settings(&["allowed.test"], &[])); + let lookups = Arc::new(AtomicUsize::new(0)); + let response = response( + resolve_query_with_lookup( + &state, + &query("denied.test", RecordType::A, /*id*/ 8), + /*client_addr*/ None, + { + let lookups = Arc::clone(&lookups); + move |_| async move { + lookups.fetch_add(1, Ordering::SeqCst); + Ok(Vec::new()) + } + }, + ) + .await, + ); + + assert_eq!(response.response_code(), ResponseCode::Refused); + assert_eq!(lookups.load(Ordering::SeqCst), 0); +} + +#[tokio::test] +async fn refuses_entire_answer_set_when_resolution_contains_a_private_address() { + let state = network_proxy_state_for_policy(settings(&["example.test"], &[])); + let response = response( + resolve_query_with_lookup( + &state, + &query("example.test", RecordType::A, /*id*/ 9), + /*client_addr*/ None, + |_| async { + Ok(vec![ + "93.184.216.34".parse().expect("public address"), + "127.0.0.1".parse().expect("private address"), + ]) + }, + ) + .await, + ); + + assert_eq!(response.response_code(), ResponseCode::Refused); + assert_eq!(response.answers(), &[]); +} + +#[tokio::test] +async fn explicitly_allowlisted_local_literal_preserves_local_policy_semantics() { + let state = network_proxy_state_for_policy(settings(&["127.0.0.1"], &[])); + let response = response( + resolve_query_with_lookup( + &state, + &query("127.0.0.1", RecordType::A, /*id*/ 9), + /*client_addr*/ None, + |_| async { Ok(vec!["127.0.0.1".parse().expect("loopback address")]) }, + ) + .await, + ); + + assert_eq!(response.response_code(), ResponseCode::NoError); + assert_eq!(response.answers().len(), 1); +} + +#[tokio::test] +async fn policy_changes_apply_without_restarting_the_dns_service() { + let state = network_proxy_state_for_policy(settings(&["example.test"], &[])); + let allowed = response( + resolve_query_with_lookup( + &state, + &query("example.test", RecordType::A, /*id*/ 10), + /*client_addr*/ None, + |_| async { Ok(vec!["93.184.216.34".parse().expect("public address")]) }, + ) + .await, + ); + assert_eq!(allowed.response_code(), ResponseCode::NoError); + + let mut replacement_network = settings(&["other.test"], &[]); + replacement_network.enabled = true; + let replacement = build_config_state(replacement_network, NetworkProxyConstraints::default()) + .expect("replacement config state"); + state + .replace_config_state(replacement) + .await + .expect("replace proxy config state"); + let denied = response( + resolve_query_with_lookup( + &state, + &query("example.test", RecordType::A, /*id*/ 11), + /*client_addr*/ None, + |_| async { Ok(vec!["93.184.216.34".parse().expect("public address")]) }, + ) + .await, + ); + assert_eq!(denied.response_code(), ResponseCode::Refused); +} + +#[tokio::test] +async fn live_disable_is_a_proxy_state_denial_without_host_lookup() { + let state = network_proxy_state_for_policy(settings(&["example.test"], &[])); + let mut disabled_network = settings(&["example.test"], &[]); + disabled_network.enabled = false; + let disabled = build_config_state(disabled_network, NetworkProxyConstraints::default()) + .expect("disabled config state"); + state + .replace_config_state(disabled) + .await + .expect("disable proxy state"); + + let lookups = Arc::new(AtomicUsize::new(0)); + let (response, events) = capture_events(|| async { + response( + resolve_query_with_lookup( + &state, + &query("example.test", RecordType::A, /*id*/ 12), + /*client_addr*/ None, + { + let lookups = Arc::clone(&lookups); + move |_| async move { + lookups.fetch_add(1, Ordering::SeqCst); + Ok(Vec::new()) + } + }, + ) + .await, + ) + }) + .await; + + assert_eq!(response.response_code(), ResponseCode::Refused); + assert_eq!(lookups.load(Ordering::SeqCst), 0); + let event = + find_event_by_name(&events, POLICY_DECISION_EVENT_NAME).expect("DNS policy audit event"); + assert_eq!(event.field("network.policy.source"), Some("proxy_state")); + let blocked = state.blocked_snapshot().await.expect("blocked DNS queries"); + assert_eq!(blocked.len(), 1); + assert_eq!(blocked[0].reason, REASON_PROXY_DISABLED); +} + +#[tokio::test] +async fn resolver_errors_are_not_reported_as_policy_denials() { + let state = network_proxy_state_for_policy(settings(&["example.test"], &[])); + let (_, events) = capture_events(|| async { + let response = response( + resolve_query_with_lookup( + &state, + &query("example.test", RecordType::A, /*id*/ 12), + /*client_addr*/ None, + |_| async { Err(io::Error::other("resolver unavailable")) }, + ) + .await, + ); + assert_eq!(response.response_code(), ResponseCode::ServFail); + }) + .await; + + assert!(find_event_by_name(&events, POLICY_DECISION_EVENT_NAME).is_none()); + assert!( + state + .blocked_snapshot() + .await + .expect("blocked request snapshot") + .is_empty() + ); +} + +#[tokio::test] +async fn refuses_unsupported_and_untrusted_query_shapes() { + let state = network_proxy_state_for_policy(settings(&["example.test"], &[])); + let mut multi_question = Message::new(); + multi_question.set_id(12); + multi_question.add_query(Query::query( + Name::from_ascii("example.test").expect("valid DNS name"), + RecordType::A, + )); + multi_question.add_query(Query::query( + Name::from_ascii("example.test").expect("valid DNS name"), + RecordType::AAAA, + )); + let queries = [ + query("example.test", RecordType::CNAME, /*id*/ 12), + query("example.test", RecordType::TXT, /*id*/ 13), + multi_question.to_vec().expect("serialize query"), + vec![0, 14, 0xff], + ]; + + for query in queries { + let response = response( + resolve_query_with_lookup( + &state, + &query, + /*client_addr*/ None, + |_| async { Ok(Vec::new()) }, + ) + .await, + ); + assert_eq!(response.response_code(), ResponseCode::Refused); + assert_eq!(response.answers(), &[]); + } +} + +#[tokio::test] +async fn denied_query_is_audited_and_reported_with_execution_attribution() { + let state = Arc::new(network_proxy_state_for_policy(settings( + &["allowed.test"], + &[], + ))); + let blocked = Arc::new(Mutex::new(Vec::new())); + state + .set_blocked_request_observer(Some(Arc::new({ + let blocked = Arc::clone(&blocked); + move |request: BlockedRequest| { + let blocked = Arc::clone(&blocked); + async move { + blocked + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(request); + } + } + }))) + .await; + state.register_execution("token-1", "local", "execution-1"); + let attributed = state + .for_execution_token("token-1") + .expect("registered execution state"); + + let (_, events) = capture_events(|| async { + resolve_query_with_lookup( + &attributed, + &query("denied.test", RecordType::A, /*id*/ 15), + Some("127.0.0.1:1234"), + |_| async { Ok(Vec::new()) }, + ) + .await + }) + .await; + let event = + find_event_by_name(&events, POLICY_DECISION_EVENT_NAME).expect("DNS policy audit event"); + assert_eq!(event.field("network.transport.protocol"), Some("dns")); + assert_eq!(event.field("execution.id"), Some("execution-1")); + + let blocked = blocked + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!(blocked.len(), 1); + assert_eq!(blocked[0].protocol, "dns"); + assert_eq!(blocked[0].execution_id.as_deref(), Some("execution-1")); +} + +#[tokio::test] +async fn listener_accepts_unscoped_policy_and_binds_attributed_execution() { + let mut network = settings(&["127.0.0.1"], &[]); + network.allow_local_binding = true; + let state = Arc::new(network_proxy_state_for_policy(network)); + state.register_execution("token-1", "local", "execution-1"); + let listener = StdTcpListener::bind("127.0.0.1:0").expect("bind DNS listener"); + let addr = listener.local_addr().expect("DNS listener address"); + let task = tokio::spawn(run_dns_proxy_with_std_listener( + Arc::clone(&state), + listener, + Some("local".to_string()), + )); + + let unscoped = query_listener(addr, /*attribution_token*/ None, "127.0.0.1") + .await + .expect("unscoped DNS query"); + assert_eq!(unscoped.response_code(), ResponseCode::NoError); + + assert!( + query_listener(addr, Some("unknown-token"), "127.0.0.1") + .await + .is_err() + ); + + let attributed = query_listener(addr, Some("token-1"), "denied.test") + .await + .expect("attributed DNS query"); + assert_eq!(attributed.response_code(), ResponseCode::Refused); + let blocked = state.blocked_snapshot().await.expect("blocked DNS queries"); + assert_eq!(blocked.len(), 1); + assert_eq!(blocked[0].execution_id.as_deref(), Some("execution-1")); + + task.abort(); + let _ = task.await; +} diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 753a0b0f4e39..27cd3481fad8 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -5,6 +5,7 @@ mod certs; mod config; mod connect_policy; mod credential_broker; +mod dns; mod http_proxy; mod mitm; mod mitm_hook; @@ -34,6 +35,8 @@ pub use config::host_and_port_from_network_addr; pub use credential_broker::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; pub use credential_broker::brokered_credential_dummy_env_keys; pub use credential_broker::brokered_credential_env_keys; +pub use dns::DNS_PROXY_ENV_KEY; +pub use dns::DNS_PROXY_SESSION_PREFACE; pub use mitm_hook::InjectedHeaderConfig; pub use mitm_hook::MitmHookActionsConfig; pub use mitm_hook::MitmHookBodyConfig; diff --git a/codex-rs/network-proxy/src/network_policy.rs b/codex-rs/network-proxy/src/network_policy.rs index 05d9f06944fe..bf76e29ad992 100644 --- a/codex-rs/network-proxy/src/network_policy.rs +++ b/codex-rs/network-proxy/src/network_policy.rs @@ -209,7 +209,7 @@ fn emit_non_domain_policy_decision_audit_event( decision, source: args.source.as_str(), reason: args.reason, - protocol: args.protocol, + protocol: args.protocol.as_policy_protocol(), server_address: args.server_address, server_port: args.server_port, method: args.method, @@ -225,7 +225,7 @@ struct PolicyAuditEventArgs<'a> { decision: &'a str, source: &'a str, reason: &'a str, - protocol: NetworkProtocol, + protocol: &'a str, server_address: &'a str, server_port: u16, method: Option<&'a str>, @@ -254,7 +254,7 @@ fn emit_policy_audit_event(state: &NetworkProxyState, args: PolicyAuditEventArgs network.policy.decision = args.decision, network.policy.source = args.source, network.policy.reason = args.reason, - network.transport.protocol = args.protocol.as_policy_protocol(), + network.transport.protocol = args.protocol, server.address = args.server_address, server.port = args.server_port, http.request.method = args.method.unwrap_or(DEFAULT_METHOD), @@ -264,6 +264,43 @@ fn emit_policy_audit_event(state: &NetworkProxyState, args: PolicyAuditEventArgs ); } +pub(crate) fn emit_dns_policy_decision_audit_event( + state: &NetworkProxyState, + host: &str, + network_decision: &NetworkDecision, + client_addr: Option<&str>, +) { + let execution_id = state.execution_id(); + let (decision, source, reason) = match network_decision { + NetworkDecision::Allow => ( + POLICY_DECISION_ALLOW, + NetworkDecisionSource::BaselinePolicy, + POLICY_REASON_ALLOW, + ), + NetworkDecision::Deny { + reason, + source, + decision, + } => (decision.as_str(), *source, reason.as_str()), + }; + emit_policy_audit_event( + state, + PolicyAuditEventArgs { + scope: POLICY_SCOPE_DOMAIN, + decision, + source: source.as_str(), + reason, + protocol: "dns", + server_address: host, + server_port: 53, + method: None, + client_addr, + execution_id: execution_id.as_deref(), + policy_override: false, + }, + ); +} + fn audit_timestamp() -> String { Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true) } @@ -364,7 +401,7 @@ pub(crate) async fn evaluate_host_policy( decision: policy_decision, source: source.as_str(), reason, - protocol: request.protocol, + protocol: request.protocol.as_policy_protocol(), server_address: request.host.as_str(), server_port: request.port, method: request.method.as_deref(), diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index 3e6a1dbe5309..5cb38fd8619d 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -4,6 +4,8 @@ use crate::attribution::PROXY_ATTRIBUTION_TOKEN_ENV_KEY; use crate::config; use crate::credential_broker::BROKERED_CREDENTIALS_ENV_KEY; use crate::credential_broker::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; +use crate::dns; +use crate::dns::DNS_PROXY_ENV_KEY; use crate::http_proxy; use crate::network_policy::NetworkPolicyDecider; use crate::runtime::BlockedRequestObserver; @@ -36,13 +38,19 @@ pub struct Args {} struct ReservedListeners { http: Mutex>, socks: Mutex>, + dns: Mutex>, } impl ReservedListeners { - fn new(http: StdTcpListener, socks: Option) -> Self { + fn new( + http: StdTcpListener, + socks: Option, + dns: Option, + ) -> Self { Self { http: Mutex::new(Some(http)), socks: Mutex::new(socks), + dns: Mutex::new(dns), } } @@ -61,18 +69,32 @@ impl ReservedListeners { .unwrap_or_else(std::sync::PoisonError::into_inner); guard.take() } + + fn take_dns(&self) -> Option { + let mut guard = self + .dns + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + guard.take() + } } struct ReservedListenerSet { http_listener: StdTcpListener, socks_listener: Option, + dns_listener: Option, } impl ReservedListenerSet { - fn new(http_listener: StdTcpListener, socks_listener: Option) -> Self { + fn new( + http_listener: StdTcpListener, + socks_listener: Option, + dns_listener: Option, + ) -> Self { Self { http_listener, socks_listener, + dns_listener, } } @@ -92,10 +114,19 @@ impl ReservedListenerSet { }) } + fn dns_addr(&self) -> Result> { + self.dns_listener + .as_ref() + .map(StdTcpListener::local_addr) + .transpose() + .context("failed to read reserved DNS proxy address") + } + fn into_reserved_listeners(self) -> Arc { Arc::new(ReservedListeners::new( self.http_listener, self.socks_listener, + self.dns_listener, )) } } @@ -183,7 +214,7 @@ impl NetworkProxyBuilder { .set_blocked_request_observer(self.blocked_request_observer.clone()) .await; let current_cfg = state.current_cfg().await?; - let (requested_http_addr, requested_socks_addr, reserved_listeners) = + let (requested_http_addr, requested_socks_addr, dns_addr, reserved_listeners) = if self.managed_by_codex { let runtime = config::resolve_runtime(¤t_cfg)?; #[cfg(target_os = "windows")] @@ -197,13 +228,18 @@ impl NetworkProxyBuilder { ) .context("reserve managed loopback proxy listeners")?; #[cfg(not(target_os = "windows"))] - let reserved = reserve_loopback_ephemeral_listeners(current_cfg.enable_socks5) - .context("reserve managed loopback proxy listeners")?; + let reserved = reserve_loopback_ephemeral_listeners( + current_cfg.enable_socks5, + /*reserve_dns_listener*/ true, + ) + .context("reserve managed loopback proxy listeners")?; let http_addr = reserved.http_addr()?; let socks_addr = reserved.socks_addr(runtime.socks_addr)?; + let dns_addr = reserved.dns_addr()?; ( http_addr, socks_addr, + dns_addr, Some(reserved.into_reserved_listeners()), ) } else { @@ -212,6 +248,7 @@ impl NetworkProxyBuilder { self.http_addr.unwrap_or(runtime.http_addr), self.socks_addr.unwrap_or(runtime.socks_addr), None, + None, ) }; @@ -223,6 +260,7 @@ impl NetworkProxyBuilder { state, http_addr, socks_addr, + dns_addr, socks_enabled: current_cfg.enable_socks5, socks5_udp_enabled: current_cfg.enable_socks5_udp, runtime_settings: Arc::new(RwLock::new(NetworkProxyRuntimeSettings::from_config( @@ -238,6 +276,7 @@ impl NetworkProxyBuilder { fn reserve_loopback_ephemeral_listeners( reserve_socks_listener: bool, + reserve_dns_listener: bool, ) -> Result { let http_listener = reserve_loopback_ephemeral_listener().context("reserve HTTP proxy listener")?; @@ -246,7 +285,16 @@ fn reserve_loopback_ephemeral_listeners( } else { None }; - Ok(ReservedListenerSet::new(http_listener, socks_listener)) + let dns_listener = if reserve_dns_listener { + Some(reserve_loopback_ephemeral_listener().context("reserve DNS proxy listener")?) + } else { + None + }; + Ok(ReservedListenerSet::new( + http_listener, + socks_listener, + dns_listener, + )) } #[cfg(target_os = "windows")] @@ -262,8 +310,11 @@ fn reserve_windows_managed_listeners( Ok(listeners) => Ok(listeners), Err(err) if err.kind() == std::io::ErrorKind::AddrInUse => { warn!("managed Windows proxy ports are busy; falling back to ephemeral loopback ports"); - reserve_loopback_ephemeral_listeners(reserve_socks_listener) - .context("reserve fallback loopback proxy listeners") + reserve_loopback_ephemeral_listeners( + reserve_socks_listener, + /*reserve_dns_listener*/ true, + ) + .context("reserve fallback loopback proxy listeners") } Err(err) => Err(err).context("reserve Windows managed proxy listeners"), } @@ -281,7 +332,12 @@ fn try_reserve_windows_managed_listeners( } else { None }; - Ok(ReservedListenerSet::new(http_listener, socks_listener)) + let dns_listener = Some(StdTcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))?); + Ok(ReservedListenerSet::new( + http_listener, + socks_listener, + dns_listener, + )) } #[cfg(target_os = "windows")] @@ -329,6 +385,7 @@ impl NetworkProxyRuntimeSettings { struct EnvironmentProxyAddrs { http_addr: SocketAddr, socks_addr: SocketAddr, + dns_addr: Option, } /// Portable managed-network facts needed by an operating-system sandbox. @@ -356,6 +413,7 @@ struct EnvironmentProxy { addrs: EnvironmentProxyAddrs, http_task: JoinHandle>, socks_task: Option>>, + dns_task: Option>>, } #[derive(Clone)] @@ -363,6 +421,7 @@ pub struct NetworkProxy { state: Arc, http_addr: SocketAddr, socks_addr: SocketAddr, + dns_addr: Option, socks_enabled: bool, socks5_udp_enabled: bool, runtime_settings: Arc>, @@ -379,6 +438,7 @@ impl std::fmt::Debug for NetworkProxy { f.debug_struct("NetworkProxy") .field("http_addr", &self.http_addr) .field("socks_addr", &self.socks_addr) + .field("dns_addr", &self.dns_addr) .finish_non_exhaustive() } } @@ -387,6 +447,7 @@ impl PartialEq for NetworkProxy { fn eq(&self, other: &Self) -> bool { self.http_addr == other.http_addr && self.socks_addr == other.socks_addr + && self.dns_addr == other.dns_addr && self.runtime_settings() == other.runtime_settings() } } @@ -425,6 +486,7 @@ pub const PROXY_ENV_KEYS: &[&str] = &[ BROKERED_CREDENTIALS_ENV_KEY, ALLOW_LOCAL_BINDING_ENV_KEY, PROXY_ATTRIBUTION_TOKEN_ENV_KEY, + DNS_PROXY_ENV_KEY, ELECTRON_GET_USE_PROXY_ENV_KEY, NODE_USE_ENV_PROXY_ENV_KEY, "HTTP_PROXY", @@ -698,6 +760,11 @@ impl NetworkProxy { runtime_settings.mitm_ca_trust_bundle.as_ref(), ); self.state.virtualize_child_credentials(&mut env); + env.remove(PROXY_ATTRIBUTION_TOKEN_ENV_KEY); + env.remove(DNS_PROXY_ENV_KEY); + if let Some(dns_addr) = addrs.dns_addr { + env.insert(DNS_PROXY_ENV_KEY.to_string(), format!("tcp://{dns_addr}")); + } if let Some(execution_scope) = self.execution_scope.as_ref() { env.insert( PROXY_ATTRIBUTION_TOKEN_ENV_KEY.to_string(), @@ -739,6 +806,7 @@ impl NetworkProxy { EnvironmentProxyAddrs { http_addr: self.http_addr, socks_addr: self.socks_addr, + dns_addr: self.dns_addr, }, ); } @@ -779,6 +847,7 @@ impl NetworkProxy { None => EnvironmentProxyAddrs { http_addr: self.http_addr, socks_addr: self.socks_addr, + dns_addr: self.dns_addr, }, }; Ok(self.prepare_for_addrs(env, addrs)) @@ -805,22 +874,28 @@ impl NetworkProxy { format!("failed to create network proxy for environment `{environment_id}`") })?; let listeners = - reserve_loopback_ephemeral_listeners(self.socks_enabled).with_context(|| { - format!("failed to reserve network proxy for environment `{environment_id}`") - })?; + reserve_loopback_ephemeral_listeners(self.socks_enabled, self.dns_addr.is_some()) + .with_context(|| { + format!("failed to reserve network proxy for environment `{environment_id}`") + })?; let http_addr = listeners.http_addr().with_context(|| { format!("failed to read HTTP proxy address for environment `{environment_id}`") })?; let socks_addr = listeners.socks_addr(self.socks_addr).with_context(|| { format!("failed to read SOCKS proxy address for environment `{environment_id}`") })?; + let dns_addr = listeners.dns_addr().with_context(|| { + format!("failed to read DNS proxy address for environment `{environment_id}`") + })?; let addrs = EnvironmentProxyAddrs { http_addr, socks_addr, + dns_addr, }; let ReservedListenerSet { http_listener, socks_listener, + dns_listener, } = listeners; let environment_id = environment_id.to_string(); @@ -858,12 +933,21 @@ impl NetworkProxy { None }; + let dns_task = dns_listener.map(|listener| { + let dns_state = self.state.clone(); + let dns_environment_id = Some(environment_id.clone()); + runtime.spawn(async move { + dns::run_dns_proxy_with_std_listener(dns_state, listener, dns_environment_id).await + }) + }); + proxies.insert( environment_id, EnvironmentProxy { addrs, http_task, socks_task, + dns_task, }, ); Ok(addrs) @@ -929,6 +1013,7 @@ impl NetworkProxy { let reserved_listeners = self.reserved_listeners.as_ref(); let http_listener = reserved_listeners.and_then(|listeners| listeners.take_http()); let socks_listener = reserved_listeners.and_then(|listeners| listeners.take_socks()); + let dns_listener = reserved_listeners.and_then(|listeners| listeners.take_dns()); let http_state = self.state.clone(); let http_decider = self.policy_decider.clone(); @@ -989,9 +1074,20 @@ impl NetworkProxy { None }; + let dns_task = dns_listener.map(|listener| { + let dns_state = self.state.clone(); + tokio::spawn(async move { + dns::run_dns_proxy_with_std_listener( + dns_state, listener, /*environment_id*/ None, + ) + .await + }) + }); + Ok(NetworkProxyHandle { http_task: Some(http_task), socks_task, + dns_task, environment_proxies: self.environment_proxies.clone(), completed: false, }) @@ -1001,6 +1097,7 @@ impl NetworkProxy { pub struct NetworkProxyHandle { http_task: Option>>, socks_task: Option>>, + dns_task: Option>>, environment_proxies: Arc>>, completed: bool, } @@ -1010,6 +1107,7 @@ impl NetworkProxyHandle { Self { http_task: Some(tokio::spawn(async { Ok(()) })), socks_task: None, + dns_task: None, environment_proxies: Arc::new(Mutex::new(HashMap::new())), completed: true, } @@ -1018,22 +1116,35 @@ impl NetworkProxyHandle { pub async fn wait(mut self) -> Result<()> { let http_task = self.http_task.take().context("missing http proxy task")?; let socks_task = self.socks_task.take(); + let dns_task = self.dns_task.take(); let http_result = http_task.await; let socks_result = match socks_task { Some(task) => Some(task.await), None => None, }; + let dns_result = match dns_task { + Some(task) => Some(task.await), + None => None, + }; self.completed = true; abort_environment_proxies(self.environment_proxies.clone()).await; http_result??; if let Some(socks_result) = socks_result { socks_result??; } + if let Some(dns_result) = dns_result { + dns_result??; + } Ok(()) } pub async fn shutdown(mut self) -> Result<()> { - abort_tasks(self.http_task.take(), self.socks_task.take()).await; + abort_tasks( + self.http_task.take(), + self.socks_task.take(), + self.dns_task.take(), + ) + .await; abort_environment_proxies(self.environment_proxies.clone()).await; self.completed = true; Ok(()) @@ -1050,9 +1161,11 @@ async fn abort_task(task: Option>>) { async fn abort_tasks( http_task: Option>>, socks_task: Option>>, + dns_task: Option>>, ) { abort_task(http_task).await; abort_task(socks_task).await; + abort_task(dns_task).await; } async fn abort_environment_proxies( @@ -1067,6 +1180,7 @@ async fn abort_environment_proxies( for proxy in proxies { abort_task(Some(proxy.http_task)).await; abort_task(proxy.socks_task).await; + abort_task(proxy.dns_task).await; } } @@ -1077,9 +1191,10 @@ impl Drop for NetworkProxyHandle { } let http_task = self.http_task.take(); let socks_task = self.socks_task.take(); + let dns_task = self.dns_task.take(); let environment_proxies = self.environment_proxies.clone(); tokio::spawn(async move { - abort_tasks(http_task, socks_task).await; + abort_tasks(http_task, socks_task, dns_task).await; abort_environment_proxies(environment_proxies).await; }); } @@ -1124,6 +1239,9 @@ mod tests { assert!(proxy.http_addr.ip().is_loopback()); assert!(proxy.socks_addr.ip().is_loopback()); + let dns_addr = proxy.dns_addr.expect("managed DNS proxy address"); + assert!(dns_addr.ip().is_loopback()); + assert_ne!(dns_addr.port(), 0); #[cfg(target_os = "windows")] { assert_eq!(proxy.http_addr, http_addr); @@ -1159,6 +1277,7 @@ mod tests { proxy.socks_addr, "127.0.0.1:48081".parse::().unwrap() ); + assert_eq!(proxy.dns_addr, None); } #[tokio::test] @@ -1167,15 +1286,52 @@ mod tests { let proxy = NetworkProxy::builder().state(state).build().await?; let handle = proxy.run().await?; - let base_env = HashMap::from([("PRESERVED".to_string(), "value".to_string())]); + let base_env = HashMap::from([ + ("PRESERVED".to_string(), "value".to_string()), + ( + DNS_PROXY_ENV_KEY.to_string(), + "tcp://127.0.0.1:9".to_string(), + ), + ( + PROXY_ATTRIBUTION_TOKEN_ENV_KEY.to_string(), + "stale-token".to_string(), + ), + ]); + let global = + proxy.prepare_for_optional_environment(HashMap::new(), /*environment_id*/ None)?; let local = proxy.prepare_for_optional_environment(base_env.clone(), Some("local"))?; let remote = proxy.prepare_for_optional_environment(HashMap::new(), Some("remote"))?; + let local_addrs = proxy.environment_proxy_addrs("local")?; + let remote_addrs = proxy.environment_proxy_addrs("remote")?; + let global_dns_addr = global + .env + .get(DNS_PROXY_ENV_KEY) + .and_then(|value| value.strip_prefix("tcp://")) + .and_then(|value| value.parse::().ok()) + .expect("global DNS proxy address"); + let local_dns_addr = local + .env + .get(DNS_PROXY_ENV_KEY) + .and_then(|value| value.strip_prefix("tcp://")) + .and_then(|value| value.parse::().ok()) + .expect("environment DNS proxy address"); + let remote_dns_addr = remote + .env + .get(DNS_PROXY_ENV_KEY) + .and_then(|value| value.strip_prefix("tcp://")) + .and_then(|value| value.parse::().ok()) + .expect("environment DNS proxy address"); assert_eq!( local.env.get("PRESERVED").map(String::as_str), Some("value") ); + assert_eq!(local.env.get(PROXY_ATTRIBUTION_TOKEN_ENV_KEY), None); assert_ne!(local.env.get("HTTP_PROXY"), remote.env.get("HTTP_PROXY")); + assert_eq!(Some(global_dns_addr), proxy.dns_addr); + assert_eq!(Some(local_dns_addr), local_addrs.dns_addr); + assert_eq!(Some(remote_dns_addr), remote_addrs.dns_addr); + assert_ne!(local_dns_addr, remote_dns_addr); assert_ne!( local.env.get("HTTP_PROXY"), Some(&format!("http://{}", proxy.http_addr())) @@ -1214,7 +1370,41 @@ mod tests { proxy.apply_to_env_for_environment(&mut legacy_env, "local")?; assert_eq!(legacy_env, local.env); + let local_execution = + proxy.for_execution("local", "execution-local", "token-local".to_string())?; + let remote_execution = + proxy.for_execution("remote", "execution-remote", "token-remote".to_string())?; + let local_execution = + local_execution.prepare_for_optional_environment(HashMap::new(), Some("local"))?; + let remote_execution = + remote_execution.prepare_for_optional_environment(HashMap::new(), Some("remote"))?; + let local_execution_dns_addr = local_execution + .env + .get(DNS_PROXY_ENV_KEY) + .and_then(|value| value.strip_prefix("tcp://")) + .and_then(|value| value.parse::().ok()) + .expect("execution-scoped DNS proxy address"); + let remote_execution_dns_addr = remote_execution + .env + .get(DNS_PROXY_ENV_KEY) + .and_then(|value| value.strip_prefix("tcp://")) + .and_then(|value| value.parse::().ok()) + .expect("execution-scoped DNS proxy address"); + assert_eq!(local_execution_dns_addr, local_dns_addr); + assert_eq!(remote_execution_dns_addr, remote_dns_addr); + assert!( + !local_execution + .sandbox_context + .loopback_ports + .contains(&local_execution_dns_addr.port()) + ); + handle.shutdown().await?; + assert!( + tokio::net::TcpStream::connect(local_dns_addr) + .await + .is_err() + ); Ok(()) } diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index 25954e83e90b..1b377456bbbe 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -89,6 +89,12 @@ pub enum HostBlockDecision { Blocked(HostBlockReason), } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum HostResolutionPolicy { + Allowed { allow_non_public_ips: bool }, + Blocked(HostBlockReason), +} + #[derive(Clone, Debug, Serialize)] pub struct BlockedRequest { pub host: String, @@ -525,6 +531,60 @@ impl NetworkProxyState { } } + /// Checks the live domain policy without resolving `host`. + /// + /// DNS callers use this before their single host lookup so denied names are not leaked to the + /// host resolver and the addresses returned to the sandbox are the same addresses checked for + /// local/private network access. + pub(crate) async fn host_resolution_policy(&self, host: &str) -> Result { + self.reload_if_needed().await?; + let host = match Host::parse(host) { + Ok(host) => host, + Err(_) => { + return Ok(HostResolutionPolicy::Blocked(HostBlockReason::NotAllowed)); + } + }; + let (deny_set, allow_set, allow_local_binding, allowed_domains) = { + let guard = self.state.read().await; + let allowed_domains = guard.config.allowed_domains(); + ( + guard.deny_set.clone(), + guard.allow_set.clone(), + guard.config.allow_local_binding, + allowed_domains, + ) + }; + let allowed_domains_empty = allowed_domains.is_none(); + let allowed_domains = allowed_domains.unwrap_or_default(); + let host_str = host.as_str(); + + if globset_matches_host_or_unscoped(&deny_set, host_str) { + return Ok(HostResolutionPolicy::Blocked(HostBlockReason::Denied)); + } + + let is_allowlisted = globset_matches_host_or_unscoped(&allow_set, host_str); + let mut explicitly_allowed_local_literal = false; + if !allow_local_binding { + let host_no_scope = unscoped_ip_literal(host_str).unwrap_or(host_str); + let local_literal = is_loopback_host(&host) + || host_no_scope.parse::().is_ok_and(is_non_public_ip); + if local_literal && !is_explicit_local_allowlisted(&allowed_domains, &host) { + return Ok(HostResolutionPolicy::Blocked( + HostBlockReason::NotAllowedLocal, + )); + } + explicitly_allowed_local_literal = local_literal; + } + + if allowed_domains_empty || !is_allowlisted { + Ok(HostResolutionPolicy::Blocked(HostBlockReason::NotAllowed)) + } else { + Ok(HostResolutionPolicy::Allowed { + allow_non_public_ips: allow_local_binding || explicitly_allowed_local_literal, + }) + } + } + pub async fn record_blocked(&self, mut entry: BlockedRequest) -> Result<()> { self.reload_if_needed().await?; entry.execution_id = self.execution_id();