diff --git a/codex-rs/network-proxy/src/connect_policy.rs b/codex-rs/network-proxy/src/connect_policy.rs new file mode 100644 index 000000000..b9425db79 --- /dev/null +++ b/codex-rs/network-proxy/src/connect_policy.rs @@ -0,0 +1,157 @@ +use crate::policy::is_non_public_ip; +use crate::state::NetworkProxyState; +use rama_core::Service; +use rama_core::error::BoxError; +use rama_core::error::ErrorExt as _; +use rama_core::error::OpaqueError; +use rama_core::extensions::ExtensionsMut; +use rama_net::address::ProxyAddress; +use rama_net::client::EstablishedClientConnection; +use rama_net::transport::TryRefIntoTransportContext; +use rama_tcp::TcpStream; +use rama_tcp::client::TcpStreamConnector; +use rama_tcp::client::service::TcpConnector; +use std::io; +use std::net::SocketAddr; +use std::sync::Arc; + +#[derive(Clone)] +pub(crate) struct TargetCheckedTcpConnector { + policy: TargetPolicy, +} + +impl TargetCheckedTcpConnector { + pub(crate) fn new(state: Arc) -> Self { + Self { + policy: TargetPolicy::State(state), + } + } + + pub(crate) fn from_allow_local_binding(allow_local_binding: bool) -> Self { + Self { + policy: TargetPolicy::Config { + allow_local_binding, + }, + } + } +} + +impl Service for TargetCheckedTcpConnector +where + Input: TryRefIntoTransportContext + Send + ExtensionsMut + 'static, + Input::Error: Into + Send + Sync + 'static, +{ + type Output = EstablishedClientConnection; + type Error = BoxError; + + async fn serve(&self, input: Input) -> Result { + if input.extensions().get::().is_some() { + return TcpConnector::new().serve(input).await; + } + + TcpConnector::new() + .with_connector(TargetCheckedStreamConnector { + policy: self.policy.clone(), + }) + .serve(input) + .await + } +} + +#[derive(Clone)] +struct TargetCheckedStreamConnector { + policy: TargetPolicy, +} + +impl TcpStreamConnector for TargetCheckedStreamConnector { + type Error = BoxError; + + async fn connect(&self, addr: SocketAddr) -> Result { + if !self.policy.allow_local_binding().await? && is_non_public_ip(addr.ip()) { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "network target rejected by policy", + ) + .into()); + } + + tokio::net::TcpStream::connect(addr) + .await + .map(TcpStream::from) + .map_err(Into::into) + } +} + +#[derive(Clone)] +enum TargetPolicy { + Config { allow_local_binding: bool }, + State(Arc), +} + +impl TargetPolicy { + async fn allow_local_binding(&self) -> Result { + match self { + Self::Config { + allow_local_binding, + } => Ok(*allow_local_binding), + Self::State(state) => state.allow_local_binding().await.map_err(|err| { + let err: BoxError = err.into(); + OpaqueError::from_boxed(err) + .context("read network proxy config") + .into_boxed() + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::NetworkProxySettings; + use crate::state::network_proxy_state_for_policy; + use rama_net::address::HostWithPort; + use std::net::Ipv4Addr; + use tokio::net::TcpListener; + + #[tokio::test(flavor = "current_thread")] + async fn direct_connector_rejects_non_public_target_when_local_binding_disabled() { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .expect("bind local listener"); + let target = listener.local_addr().expect("local addr"); + let connector = TargetCheckedTcpConnector::new(Arc::new(network_proxy_state_for_policy( + NetworkProxySettings::default(), + ))); + + let request: rama_tcp::client::Request = + rama_tcp::client::Request::new(HostWithPort::from(target)); + let err = Service::serve(&connector, request) + .await + .expect_err("local target should be rejected"); + + assert!( + format!("{err:?}").contains("network target rejected by policy"), + "unexpected error: {err:?}" + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn direct_connector_allows_non_public_target_when_local_binding_enabled() { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .expect("bind local listener"); + let target = listener.local_addr().expect("local addr"); + let connector = TargetCheckedTcpConnector::new(Arc::new(network_proxy_state_for_policy( + NetworkProxySettings { + allow_local_binding: true, + ..NetworkProxySettings::default() + }, + ))); + + let request: rama_tcp::client::Request = + rama_tcp::client::Request::new(HostWithPort::from(target)); + let result = Service::serve(&connector, request).await; + + assert!(result.is_ok(), "local target should be allowed: {result:?}"); + } +} diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index a29af08b1..fa3616844 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -1,4 +1,5 @@ use crate::config::NetworkMode; +use crate::connect_policy::TargetCheckedTcpConnector; use crate::mitm; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; @@ -66,7 +67,6 @@ use rama_net::proxy::ProxyTarget; use rama_net::proxy::StreamForwardService; use rama_net::stream::SocketInfo; use rama_tcp::client::Request as TcpRequest; -use rama_tcp::client::service::TcpConnector; use rama_tcp::server::TcpListener; use rama_tls_rustls::client::TlsConnectorDataBuilder; use rama_tls_rustls::client::TlsConnectorLayer; @@ -345,20 +345,22 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { return Ok(()); } - let allow_upstream_proxy = match upgraded + let app_state = match upgraded .extensions() .get::>() .cloned() { - Some(state) => match state.allow_upstream_proxy().await { - Ok(allowed) => allowed, - Err(err) => { - error!("failed to read upstream proxy setting: {err}"); - false - } - }, + Some(state) => state, None => { error!("missing app state"); + return Ok(()); + } + }; + + let allow_upstream_proxy = match app_state.allow_upstream_proxy().await { + Ok(allowed) => allowed, + Err(err) => { + error!("failed to read upstream proxy setting: {err}"); false } }; @@ -369,7 +371,7 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { None }; - if let Err(err) = forward_connect_tunnel(upgraded, proxy).await { + if let Err(err) = forward_connect_tunnel(upgraded, proxy, app_state).await { warn!("tunnel error: {err}"); } Ok(()) @@ -378,6 +380,7 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { async fn forward_connect_tunnel( upgraded: Upgraded, proxy: Option, + app_state: Arc, ) -> Result<(), BoxError> { let authority = upgraded .extensions() @@ -392,7 +395,7 @@ async fn forward_connect_tunnel( let req = TcpRequest::new_with_extensions(authority.clone(), extensions) .with_protocol(Protocol::HTTPS); - let proxy_connector = HttpProxyConnector::optional(TcpConnector::new()); + let proxy_connector = HttpProxyConnector::optional(TargetCheckedTcpConnector::new(app_state)); let tls_config = TlsConnectorDataBuilder::new() .with_alpn_protocols_http_auto() .build(); @@ -730,9 +733,9 @@ async fn http_plain_proxy( Err(resp) => return Ok(resp), }; let client = if allow_upstream_proxy { - UpstreamClient::from_env_proxy() + UpstreamClient::from_env_proxy(app_state.clone()) } else { - UpstreamClient::direct() + UpstreamClient::direct(app_state.clone()) }; // Strip hop-by-hop headers only after extracting metadata used for policy correlation. diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index 760b09e21..416ebb0a4 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -2,6 +2,7 @@ mod certs; mod config; +mod connect_policy; mod http_proxy; mod mitm; mod network_policy; diff --git a/codex-rs/network-proxy/src/mitm.rs b/codex-rs/network-proxy/src/mitm.rs index 8f20afaa3..7be700b1d 100644 --- a/codex-rs/network-proxy/src/mitm.rs +++ b/codex-rs/network-proxy/src/mitm.rs @@ -52,6 +52,11 @@ pub struct MitmState { max_body_bytes: usize, } +pub(crate) struct MitmUpstreamConfig { + pub(crate) allow_upstream_proxy: bool, + pub(crate) allow_local_binding: bool, +} + #[derive(Clone)] struct MitmPolicyContext { target_host: String, @@ -80,16 +85,16 @@ impl std::fmt::Debug for MitmState { } impl MitmState { - pub(crate) fn new(allow_upstream_proxy: bool) -> Result { + pub(crate) fn new(config: MitmUpstreamConfig) -> Result { // MITM exists to make limited-mode HTTPS enforceable: once CONNECT is established, plain // proxying would lose visibility into the inner HTTP request. We generate/load a local CA // and issue per-host leaf certs so we can terminate TLS and apply policy. let ca = ManagedMitmCa::load_or_create()?; - let upstream = if allow_upstream_proxy { - UpstreamClient::from_env_proxy() + let upstream = if config.allow_upstream_proxy { + UpstreamClient::from_env_proxy_with_allow_local_binding(config.allow_local_binding) } else { - UpstreamClient::direct() + UpstreamClient::direct_with_allow_local_binding(config.allow_local_binding) }; Ok(Self { diff --git a/codex-rs/network-proxy/src/runtime.rs b/codex-rs/network-proxy/src/runtime.rs index daafeec73..bddf81441 100644 --- a/codex-rs/network-proxy/src/runtime.rs +++ b/codex-rs/network-proxy/src/runtime.rs @@ -532,6 +532,12 @@ impl NetworkProxyState { Ok(guard.config.network.allow_upstream_proxy) } + pub async fn allow_local_binding(&self) -> Result { + self.reload_if_needed().await?; + let guard = self.state.read().await; + Ok(guard.config.network.allow_local_binding) + } + pub async fn network_mode(&self) -> Result { self.reload_if_needed().await?; let guard = self.state.read().await; diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index b360d6853..2d4c05f95 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -1,4 +1,5 @@ use crate::config::NetworkMode; +use crate::connect_policy::TargetCheckedTcpConnector; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; use crate::network_policy::NetworkDecisionSource; @@ -34,7 +35,6 @@ use rama_socks5::server::udp::RelayRequest; use rama_socks5::server::udp::RelayResponse; use rama_tcp::TcpStream; use rama_tcp::client::Request as TcpRequest; -use rama_tcp::client::service::TcpConnector; use rama_tcp::server::TcpListener; use std::io; use std::net::SocketAddr; @@ -94,7 +94,7 @@ async fn run_socks5_with_listener( } } - let tcp_connector = TcpConnector::default(); + let tcp_connector = TargetCheckedTcpConnector::new(state.clone()); let policy_tcp_connector = service_fn({ let policy_decider = policy_decider.clone(); move |req: TcpRequest| { @@ -131,7 +131,7 @@ async fn run_socks5_with_listener( async fn handle_socks5_tcp( req: TcpRequest, - tcp_connector: TcpConnector, + tcp_connector: TargetCheckedTcpConnector, policy_decider: Option>, ) -> Result, BoxError> { let app_state = req @@ -548,7 +548,7 @@ mod tests { let (result, events) = capture_events(|| async { handle_socks5_tcp( request, - TcpConnector::default(), + TargetCheckedTcpConnector::new(state.clone()), /*policy_decider*/ None, ) .await diff --git a/codex-rs/network-proxy/src/state.rs b/codex-rs/network-proxy/src/state.rs index 01e4966b5..b14968e9a 100644 --- a/codex-rs/network-proxy/src/state.rs +++ b/codex-rs/network-proxy/src/state.rs @@ -3,6 +3,7 @@ use crate::config::NetworkMode; use crate::config::NetworkProxyConfig; use crate::config::NetworkUnixSocketPermissions; use crate::mitm::MitmState; +use crate::mitm::MitmUpstreamConfig; use crate::policy::DomainPattern; use crate::policy::compile_allowlist_globset; use crate::policy::compile_denylist_globset; @@ -66,9 +67,10 @@ pub fn build_config_state( let deny_set = compile_denylist_globset(&denied_domains)?; let allow_set = compile_allowlist_globset(&allowed_domains)?; let mitm = if config.network.mitm { - Some(Arc::new(MitmState::new( - config.network.allow_upstream_proxy, - )?)) + Some(Arc::new(MitmState::new(MitmUpstreamConfig { + allow_upstream_proxy: config.network.allow_upstream_proxy, + allow_local_binding: config.network.allow_local_binding, + })?)) } else { None }; diff --git a/codex-rs/network-proxy/src/upstream.rs b/codex-rs/network-proxy/src/upstream.rs index 97e78f303..c7b67cc18 100644 --- a/codex-rs/network-proxy/src/upstream.rs +++ b/codex-rs/network-proxy/src/upstream.rs @@ -1,3 +1,5 @@ +use crate::connect_policy::TargetCheckedTcpConnector; +use crate::state::NetworkProxyState; use rama_core::Layer; use rama_core::Service; use rama_core::error::BoxError; @@ -16,9 +18,9 @@ use rama_http_backend::client::proxy::layer::HttpProxyConnectorLayer; use rama_net::address::ProxyAddress; use rama_net::client::EstablishedClientConnection; use rama_net::http::RequestContext; -use rama_tcp::client::service::TcpConnector; use rama_tls_rustls::client::TlsConnectorDataBuilder; use rama_tls_rustls::client::TlsConnectorLayer; +use std::sync::Arc; use tracing::warn; #[cfg(target_os = "macos")] @@ -102,12 +104,32 @@ pub(crate) struct UpstreamClient { } impl UpstreamClient { - pub(crate) fn direct() -> Self { - Self::new(ProxyConfig::default()) + pub(crate) fn direct(state: Arc) -> Self { + Self::new( + ProxyConfig::default(), + TargetCheckedTcpConnector::new(state), + ) } - pub(crate) fn from_env_proxy() -> Self { - Self::new(ProxyConfig::from_env()) + pub(crate) fn from_env_proxy(state: Arc) -> Self { + Self::new( + ProxyConfig::from_env(), + TargetCheckedTcpConnector::new(state), + ) + } + + pub(crate) fn direct_with_allow_local_binding(allow_local_binding: bool) -> Self { + Self::new( + ProxyConfig::default(), + TargetCheckedTcpConnector::from_allow_local_binding(allow_local_binding), + ) + } + + pub(crate) fn from_env_proxy_with_allow_local_binding(allow_local_binding: bool) -> Self { + Self::new( + ProxyConfig::from_env(), + TargetCheckedTcpConnector::from_allow_local_binding(allow_local_binding), + ) } #[cfg(target_os = "macos")] @@ -119,8 +141,8 @@ impl UpstreamClient { } } - fn new(proxy_config: ProxyConfig) -> Self { - let connector = build_http_connector(); + fn new(proxy_config: ProxyConfig, transport: TargetCheckedTcpConnector) -> Self { + let connector = build_http_connector(transport); Self { connector, proxy_config, @@ -158,12 +180,13 @@ impl Service> for UpstreamClient { } } -fn build_http_connector() -> BoxService< +fn build_http_connector( + transport: TargetCheckedTcpConnector, +) -> BoxService< Request, EstablishedClientConnection, Request>, BoxError, > { - let transport = TcpConnector::default(); let proxy = HttpProxyConnectorLayer::optional().into_layer(transport); let tls_config = TlsConnectorDataBuilder::new() .with_alpn_protocols_http_auto()