diff --git a/codex-rs/core/tests/suite/network_approval.rs b/codex-rs/core/tests/suite/network_approval.rs index 5651c6f825..4d6f92e966 100644 --- a/codex-rs/core/tests/suite/network_approval.rs +++ b/codex-rs/core/tests/suite/network_approval.rs @@ -3,6 +3,7 @@ use anyhow::Result; use codex_config::types::ApprovalsReviewer; use codex_core::EnvironmentConfig; use codex_core::EnvironmentNetworkPolicy; +use codex_core::StartThreadOptions; use codex_core::TurnInputRequest; use codex_core::config::Constrained; use codex_core::config::NetworkProxySpec; @@ -1506,6 +1507,104 @@ async fn ambiguous_unattributed_network_request_is_not_assigned_to_active_calls( Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn thread_turnover_closes_managed_proxy_tunnels() -> Result<()> { + skip_if_host_windows!(Ok(())); + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + + let server = start_mock_server().await; + let test = managed_network_unified_exec_test(&server).await?; + let mut config = test.config.clone(); + let mut network = NetworkProxyConfig { + enabled: true, + mode: codex_network_proxy::NetworkMode::Full, + allow_local_binding: true, + allow_upstream_proxy: false, + ..NetworkProxyConfig::default() + }; + network.set_allowed_domains(vec!["127.0.0.1".to_string()]); + config.permissions.network = Some(NetworkProxySpec::from_config_and_constraints( + network, + /*requirements*/ None, + config.permissions.permission_profile(), + )?); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + let target = listener.local_addr()?; + + for cycle in 0..3 { + let started = test + .thread_manager + .start_thread(StartThreadOptions::new(config.clone())) + .await?; + let proxy = started + .session_configured + .network_proxy + .as_ref() + .context("expected managed network proxy for new thread")?; + let (mut client, upstream) = tokio::time::timeout(Duration::from_secs(5), async { + let mut client = tokio::net::TcpStream::connect(&proxy.http_addr).await?; + client + .write_all( + format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n").as_bytes(), + ) + .await?; + let (mut upstream, _) = listener.accept().await?; + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(client.read_u8().await?); + } + let headers = String::from_utf8(headers)?; + assert!(headers.starts_with("HTTP/1.1 200 "), "{headers:?}"); + client.write_all(b"request").await?; + client.shutdown().await?; + let mut request = Vec::new(); + upstream.read_to_end(&mut request).await?; + assert_eq!(request, b"request"); + // Preserve valid half-close behavior while the thread is alive; its + // upstream peer deliberately keeps the write half open through unload. + upstream.write_all(b"response").await?; + let mut response = [0_u8; 8]; + client.read_exact(&mut response).await?; + assert_eq!(&response, b"response"); + Ok::<_, anyhow::Error>((client, upstream)) + }) + .await + .with_context(|| { + format!("thread cycle {cycle} did not establish a half-closed tunnel") + })??; + + tokio::time::timeout(Duration::from_secs(5), started.thread.shutdown_and_wait()) + .await + .with_context(|| format!("thread cycle {cycle} did not shut down"))??; + // Unloading releases both the manager's and caller's strong session owners. + drop( + test.thread_manager + .remove_thread(&started.thread_id) + .await + .context("expected stopped thread in manager")?, + ); + drop(started); + + let result = tokio::time::timeout(Duration::from_secs(2), client.read(&mut [0_u8; 1])) + .await + .with_context(|| format!("thread cycle {cycle} retained its tunnel after unload"))?; + match result { + Ok(bytes) => assert_eq!(bytes, 0), + Err(error) => assert!( + matches!( + error.kind(), + std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted + ), + "unexpected tunnel error after thread unload: {error}" + ), + } + drop(upstream); + } + test.codex.shutdown_and_wait().await?; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[cfg_attr( not(target_os = "linux"), diff --git a/codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs b/codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs new file mode 100644 index 0000000000..60be13c30f --- /dev/null +++ b/codex-rs/network-proxy/src/connection_lifecycle/lifecycle_tests.rs @@ -0,0 +1,266 @@ +//! Exercises proxy teardown through live TCP connections, including half-closed tunnels. + +use crate::NetworkProxy; +use crate::NetworkProxyConfig; +use crate::state::network_proxy_state_for_policy; +use anyhow::Context; +use anyhow::Result; +use pretty_assertions::assert_eq; +use std::collections::HashMap; +use std::future::Future; +use std::future::poll_fn; +use std::io::ErrorKind; +use std::net::Ipv4Addr; +use std::net::SocketAddr; +use std::sync::Arc; +use std::task::Poll; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::time::timeout; + +#[derive(Clone, Copy, Debug)] +enum ConnectionKind { + HttpConnect(TunnelState), + Socks5(TunnelState), + HttpKeepAlive, +} + +#[derive(Clone, Copy, Debug)] +enum TunnelState { + Open, + ClientWriteClosed, + UpstreamWriteClosed, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PeerSide { + Client, + Upstream, +} + +struct Connection { + client: TcpStream, + upstream: TcpStream, + eof_received_by: Option, +} + +#[derive(Clone, Copy, Debug)] +enum StopProxy { + Shutdown, + Drop, + CancelWait, +} + +#[tokio::test] +async fn shutdown_closes_proxy_connections() -> Result<()> { + assert_stop_closes_connections(StopProxy::Shutdown).await +} + +#[tokio::test] +async fn dropping_handle_closes_proxy_connections() -> Result<()> { + assert_stop_closes_connections(StopProxy::Drop).await +} + +#[tokio::test] +async fn canceling_wait_closes_proxy_connections() -> Result<()> { + assert_stop_closes_connections(StopProxy::CancelWait).await +} + +async fn assert_stop_closes_connections(stop: StopProxy) -> Result<()> { + let mut config = NetworkProxyConfig { + allow_local_binding: true, + allow_upstream_proxy: false, + enable_socks5_udp: false, + ..NetworkProxyConfig::default() + }; + config.set_allowed_domains(vec![Ipv4Addr::LOCALHOST.to_string()]); + let proxy = NetworkProxy::builder() + .state(Arc::new(network_proxy_state_for_policy(config))) + .managed_by_codex(!cfg!(target_os = "windows")) + .http_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 0))) + .socks_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 0))) + .build() + .await?; + let handle = proxy.run().await?; + // Environment listeners reserve ephemeral ports on every platform. Managed main + // listeners do so off Windows, whose shared ingress has separate route tests. + let prepared = proxy.prepare_for_remote_environment(HashMap::new(), "lifecycle-test")?; + let scopes = [ + ("environment", prepared.env), + #[cfg(not(target_os = "windows"))] + ( + "main", + proxy + .prepare_for_optional_environment(HashMap::new(), /*environment_id*/ None)? + .env, + ), + ]; + let mut connections = Vec::new(); + for (scope, env) in scopes { + for kind in [ + ConnectionKind::HttpConnect(TunnelState::Open), + ConnectionKind::HttpConnect(TunnelState::ClientWriteClosed), + ConnectionKind::HttpConnect(TunnelState::UpstreamWriteClosed), + ConnectionKind::Socks5(TunnelState::Open), + ConnectionKind::Socks5(TunnelState::ClientWriteClosed), + ConnectionKind::Socks5(TunnelState::UpstreamWriteClosed), + ConnectionKind::HttpKeepAlive, + ] { + let (key, prefix) = match kind { + ConnectionKind::HttpConnect(_) | ConnectionKind::HttpKeepAlive => { + ("HTTP_PROXY", "http://") + } + ConnectionKind::Socks5(_) => ("ALL_PROXY", "socks5h://"), + }; + let addr = env[key] + .strip_prefix(prefix) + .context("proxy URL scheme")? + .parse()?; + let connection = timeout(Duration::from_secs(5), open_connection(kind, addr)) + .await + .with_context(|| format!("{scope} {kind:?} connection did not become ready"))??; + connections.push((scope, kind, connection)); + } + } + + match stop { + StopProxy::Shutdown => timeout(Duration::from_secs(2), handle.shutdown()) + .await + .context("proxy shutdown did not finish with active connections")??, + StopProxy::Drop => drop(handle), + StopProxy::CancelWait => { + let mut wait = Box::pin(handle.wait()); + let state = poll_fn(|cx| Poll::Ready(wait.as_mut().poll(cx))).await; + assert!( + state.is_pending(), + "running proxy unexpectedly stopped: {state:?}" + ); + drop(wait); + } + } + + // Retain every endpoint until all checks finish: dropping a test peer could + // otherwise cause the EOF that a later assertion attributes to proxy teardown. + for (scope, kind, connection) in &mut connections { + for (side, stream) in [ + (PeerSide::Client, &mut connection.client), + (PeerSide::Upstream, &mut connection.upstream), + ] { + if connection.eof_received_by == Some(side) { + continue; + } + let mut byte = [0_u8; 1]; + let result = timeout(Duration::from_secs(2), stream.read(&mut byte)) + .await + .with_context(|| { + format!("{scope} {kind:?} {side:?} connection remained open after {stop:?}") + })?; + match result { + Ok(0) => {} + Err(error) + if matches!( + error.kind(), + ErrorKind::ConnectionReset | ErrorKind::ConnectionAborted + ) => {} + result => anyhow::bail!( + "{scope} {kind:?} {side:?} did not close after {stop:?}: {result:?}" + ), + } + } + } + Ok(()) +} + +async fn open_connection(kind: ConnectionKind, proxy_addr: SocketAddr) -> Result { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await?; + let target = listener.local_addr()?; + let mut client = TcpStream::connect(proxy_addr).await?; + let request = match kind { + ConnectionKind::HttpConnect(_) => { + format!("CONNECT {target} HTTP/1.1\r\nHost: {target}\r\n\r\n").into_bytes() + } + ConnectionKind::HttpKeepAlive => { + format!("GET http://{target}/ HTTP/1.1\r\nHost: {target}\r\n\r\n").into_bytes() + } + ConnectionKind::Socks5(_) => { + client.write_all(&[5, 1, 0]).await?; + let mut greeting = [0_u8; 2]; + client.read_exact(&mut greeting).await?; + assert_eq!(greeting, [5, 0]); + let [port_high, port_low] = target.port().to_be_bytes(); + vec![5, 1, 0, 1, 127, 0, 0, 1, port_high, port_low] + } + }; + client.write_all(&request).await?; + let (upstream, _) = listener.accept().await?; + let mut connection = Connection { + client, + upstream, + eof_received_by: None, + }; + let tunnel_state = match kind { + ConnectionKind::HttpConnect(state) => { + let response = read_http_headers(&mut connection.client).await?; + assert!(response.starts_with("HTTP/1.1 200 "), "{response:?}"); + state + } + ConnectionKind::Socks5(state) => { + let mut response = [0_u8; 10]; + connection.client.read_exact(&mut response).await?; + assert_eq!(&response[..4], &[5, 0, 0, 1]); + state + } + ConnectionKind::HttpKeepAlive => { + read_http_headers(&mut connection.upstream).await?; + connection + .upstream + .write_all(b"HTTP/1.1 204 No Content\r\n\r\n") + .await?; + let response = read_http_headers(&mut connection.client).await?; + assert!(response.starts_with("HTTP/1.1 204 "), "{response:?}"); + return Ok(connection); + } + }; + + let (sender, receiver, receiver_side) = match tunnel_state { + TunnelState::Open | TunnelState::ClientWriteClosed => ( + &mut connection.client, + &mut connection.upstream, + PeerSide::Upstream, + ), + TunnelState::UpstreamWriteClosed => ( + &mut connection.upstream, + &mut connection.client, + PeerSide::Client, + ), + }; + sender.write_all(b"request").await?; + let mut request = [0_u8; 7]; + receiver.read_exact(&mut request).await?; + assert_eq!(&request, b"request"); + match tunnel_state { + TunnelState::Open => {} + TunnelState::ClientWriteClosed | TunnelState::UpstreamWriteClosed => { + sender.shutdown().await?; + assert_eq!(receiver.read(&mut [0_u8; 1]).await?, 0); + connection.eof_received_by = Some(receiver_side); + } + } + // Traffic in the opposite direction remains valid after a write half closes. + receiver.write_all(b"response").await?; + let mut response = [0_u8; 8]; + sender.read_exact(&mut response).await?; + assert_eq!(&response, b"response"); + Ok(connection) +} + +async fn read_http_headers(stream: &mut TcpStream) -> Result { + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(stream.read_u8().await?); + } + Ok(String::from_utf8(headers)?) +} diff --git a/codex-rs/network-proxy/src/connection_lifecycle/listeners.rs b/codex-rs/network-proxy/src/connection_lifecycle/listeners.rs new file mode 100644 index 0000000000..5c414bcd58 --- /dev/null +++ b/codex-rs/network-proxy/src/connection_lifecycle/listeners.rs @@ -0,0 +1,47 @@ +//! Keeps listener tasks and their connection scope alive under the same runtime owner. + +use super::scope::ConnectionLifecycle; +use anyhow::Result; +use rama_core::graceful::ShutdownGuard; +use std::future::Future; +use tokio::task::JoinSet; + +pub(crate) struct ProxyListeners { + connections: ConnectionLifecycle, + listeners: JoinSet>, +} + +impl ProxyListeners { + pub(crate) fn new() -> Self { + Self { + connections: ConnectionLifecycle::new(), + listeners: JoinSet::new(), + } + } + + pub(crate) fn spawn(&mut self, listener: F) + where + F: FnOnce(ShutdownGuard) -> Fut, + Fut: Future> + Send + 'static, + { + self.listeners.spawn(listener(self.connections.guard())); + } + + pub(crate) fn cancel(&mut self) { + self.connections.cancel(); + self.listeners.abort_all(); + } + + pub(crate) async fn wait(&mut self) -> Result<()> { + while let Some(result) = self.listeners.join_next().await { + result??; + } + Ok(()) + } + + pub(crate) async fn shutdown(mut self) { + self.cancel(); + self.listeners.shutdown().await; + self.connections.shutdown().await; + } +} diff --git a/codex-rs/network-proxy/src/connection_lifecycle/mod.rs b/codex-rs/network-proxy/src/connection_lifecycle/mod.rs new file mode 100644 index 0000000000..b0219bbcd4 --- /dev/null +++ b/codex-rs/network-proxy/src/connection_lifecycle/mod.rs @@ -0,0 +1,12 @@ +mod listeners; +mod scope; +mod service; + +pub(crate) use listeners::ProxyListeners; +#[cfg(test)] +pub(crate) use scope::ConnectionLifecycle; +pub(crate) use service::CancelOnShutdown; + +#[cfg(test)] +#[path = "lifecycle_tests.rs"] +mod tests; diff --git a/codex-rs/network-proxy/src/connection_lifecycle/scope.rs b/codex-rs/network-proxy/src/connection_lifecycle/scope.rs new file mode 100644 index 0000000000..2af4b500c8 --- /dev/null +++ b/codex-rs/network-proxy/src/connection_lifecycle/scope.rs @@ -0,0 +1,35 @@ +//! Owns cancellation and completion of a proxy's accepted connections and descendant tasks. + +use rama_core::graceful::Shutdown; +use rama_core::graceful::ShutdownGuard; +use tokio::sync::oneshot; + +pub(crate) struct ConnectionLifecycle { + // Only the runtime owner holds the sender. Dropping it also signals shutdown when an + // enclosing wait/shutdown future is cancelled, without spawning another cleanup task. + cancel: Option>, + shutdown: Shutdown, +} + +impl ConnectionLifecycle { + pub(crate) fn new() -> Self { + let (cancel, signal) = oneshot::channel(); + Self { + cancel: Some(cancel), + shutdown: Shutdown::new(signal), + } + } + + pub(crate) fn guard(&self) -> ShutdownGuard { + self.shutdown.guard() + } + + pub(crate) fn cancel(&mut self) { + self.cancel.take(); + } + + pub(crate) async fn shutdown(mut self) { + self.cancel(); + self.shutdown.shutdown().await; + } +} diff --git a/codex-rs/network-proxy/src/connection_lifecycle/service.rs b/codex-rs/network-proxy/src/connection_lifecycle/service.rs new file mode 100644 index 0000000000..2444b9cb67 --- /dev/null +++ b/codex-rs/network-proxy/src/connection_lifecycle/service.rs @@ -0,0 +1,43 @@ +//! Cancels connection work when its executor's proxy scope ends, including HTTP upgrades. + +use rama_core::Service; +use rama_core::extensions::ExtensionsRef; +use rama_core::rt::Executor; + +#[derive(Clone)] +pub(crate) struct CancelOnShutdown { + inner: S, +} + +impl CancelOnShutdown { + pub(crate) fn new(inner: S) -> Self { + Self { inner } + } +} + +impl Service for CancelOnShutdown +where + S: Service, + Request: ExtensionsRef + Send + 'static, +{ + type Output = (); + type Error = S::Error; + + async fn serve(&self, request: Request) -> Result<(), Self::Error> { + let guard = request + .extensions() + .get::() + .and_then(Executor::guard) + .cloned(); + match guard { + Some(guard) => { + tokio::select! { + biased; + _ = guard.cancelled() => Ok(()), + result = self.inner.serve(request) => result, + } + } + None => self.inner.serve(request).await, + } + } +} diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index 8d9670a41d..da9bc494a8 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -1,6 +1,7 @@ use crate::attribution::BindConnectionAttribution; use crate::config::NetworkMode; use crate::connect_policy::TargetCheckedTcpConnector; +use crate::connection_lifecycle::CancelOnShutdown; use crate::mitm; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; @@ -41,6 +42,7 @@ use rama_core::error::ErrorExt as _; use rama_core::error::OpaqueError; use rama_core::extensions::ExtensionsMut; use rama_core::extensions::ExtensionsRef; +use rama_core::graceful::ShutdownGuard; use rama_core::service::BoxService; use rama_core::service::service_fn; use rama_core::stream::Stream; @@ -95,6 +97,7 @@ pub async fn run_http_proxy( addr: SocketAddr, policy_decider: Option>, environment_id: Option, + guard: ShutdownGuard, ) -> Result<()> { let listener = TcpListener::build() .bind(addr) @@ -107,7 +110,7 @@ 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, environment_id, guard).await } pub async fn run_http_proxy_with_std_listener( @@ -115,10 +118,11 @@ pub async fn run_http_proxy_with_std_listener( listener: StdTcpListener, policy_decider: Option>, environment_id: Option, + guard: ShutdownGuard, ) -> 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, environment_id, guard).await } async fn run_http_proxy_with_listener( @@ -126,6 +130,7 @@ async fn run_http_proxy_with_listener( listener: TcpListener, policy_decider: Option>, environment_id: Option, + guard: ShutdownGuard, ) -> Result<()> { let addr = listener .local_addr() @@ -134,7 +139,10 @@ async fn run_http_proxy_with_listener( info!("HTTP proxy listening on {addr}"); listener - .serve(http_proxy_service(state, policy_decider, environment_id)) + .serve_graceful( + guard, + CancelOnShutdown::new(http_proxy_service(state, policy_decider, environment_id)), + ) .await; Ok(()) } @@ -161,7 +169,7 @@ pub(crate) fn http_proxy_service( http_connect_accept(policy_decider.clone(), environment_id.clone(), req) } }), - service_fn(http_connect_proxy), + CancelOnShutdown::new(service_fn(http_connect_proxy)), ), RemoveResponseHeaderLayer::hop_by_hop(), ) @@ -1445,8 +1453,13 @@ mod tests { let proxy_addr = listener .local_addr() .expect("proxy listener should expose local addr"); + let lifecycle = crate::connection_lifecycle::ConnectionLifecycle::new(); let proxy_task = tokio::spawn(run_http_proxy_with_std_listener( - state, listener, /*policy_decider*/ None, /*environment_id*/ None, + state, + listener, + /*policy_decider*/ None, + /*environment_id*/ None, + lifecycle.guard(), )); let mut stream = tokio::net::TcpStream::connect(proxy_addr) @@ -1522,11 +1535,13 @@ mod tests { let proxy_addr = listener .local_addr() .expect("proxy listener should expose local addr"); + let lifecycle = crate::connection_lifecycle::ConnectionLifecycle::new(); let proxy_task = tokio::spawn(run_http_proxy_with_std_listener( state.clone(), listener, /*policy_decider*/ None, /*environment_id*/ None, + lifecycle.guard(), )); let mut stream = tokio::net::TcpStream::connect(proxy_addr) diff --git a/codex-rs/network-proxy/src/lib.rs b/codex-rs/network-proxy/src/lib.rs index b2affe0b8c..38ac654d46 100644 --- a/codex-rs/network-proxy/src/lib.rs +++ b/codex-rs/network-proxy/src/lib.rs @@ -5,6 +5,7 @@ mod authorization_path; mod certs; mod config; mod connect_policy; +mod connection_lifecycle; mod credential_broker; mod environment_policy; mod http_proxy; diff --git a/codex-rs/network-proxy/src/proxy.rs b/codex-rs/network-proxy/src/proxy.rs index 8172645e13..42a2a8135b 100644 --- a/codex-rs/network-proxy/src/proxy.rs +++ b/codex-rs/network-proxy/src/proxy.rs @@ -2,6 +2,7 @@ mod execution_scope; use crate::attribution::PROXY_ATTRIBUTION_TOKEN_ENV_KEY; use crate::config; +use crate::connection_lifecycle::ProxyListeners; use crate::credential_broker::BROKERED_CREDENTIALS_ENV_KEY; use crate::credential_broker::CREDENTIAL_BROKER_ACTIVE_ENV_KEY; use crate::http_proxy; @@ -35,7 +36,6 @@ use std::ops::RangeInclusive; use std::sync::Arc; use std::sync::Mutex; use std::sync::RwLock; -use tokio::task::JoinHandle; use tracing::warn; use self::execution_scope::ExecutionScope; @@ -470,7 +470,7 @@ pub struct PreparedManagedNetwork { struct EnvironmentProxy { addrs: EnvironmentProxyAddrs, - runtime: EnvironmentProxyRuntime, + runtime: ProxyRuntime, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -479,21 +479,40 @@ enum EnvironmentProxyClient { TrustedBridge, } -enum EnvironmentProxyRuntime { - ListenerTasks { - http_task: JoinHandle>, - socks_task: Option>>, - }, +enum ProxyRuntime { + Listeners(ProxyListeners), #[cfg(target_os = "windows")] - SharedIngress { _route: Arc }, + SharedIngress { + route: Arc, + }, } -impl EnvironmentProxyRuntime { +impl ProxyRuntime { #[cfg(target_os = "windows")] fn network_proxy_restricting_sid(&self) -> Option { match self { - Self::ListenerTasks { .. } => None, - Self::SharedIngress { _route: route } => Some(route.sid().to_string()), + Self::Listeners(_) => None, + Self::SharedIngress { route } => Some(route.sid().to_string()), + } + } + + async fn wait(&mut self) -> Result<()> { + match self { + Self::Listeners(listeners) => listeners.wait().await, + #[cfg(target_os = "windows")] + Self::SharedIngress { .. } => std::future::pending().await, + } + } + + fn stop(self) -> Option { + match self { + Self::Listeners(mut listeners) => { + listeners.cancel(); + Some(listeners) + } + // Shared ingress owns its listeners; dropping the route unregisters it. + #[cfg(target_os = "windows")] + Self::SharedIngress { .. } => None, } } } @@ -1201,8 +1220,8 @@ impl NetworkProxy { anyhow::ensure!( matches!( (&proxy.runtime, uses_shared_ingress), - (EnvironmentProxyRuntime::SharedIngress { .. }, true) - | (EnvironmentProxyRuntime::ListenerTasks { .. }, false) + (ProxyRuntime::SharedIngress { .. }, true) + | (ProxyRuntime::Listeners(_), false) ), "network proxy for environment `{environment_id}` was prepared for a different client type" ); @@ -1244,13 +1263,13 @@ impl NetworkProxy { environment_id, EnvironmentProxy { addrs, - runtime: EnvironmentProxyRuntime::SharedIngress { _route: route }, + runtime: ProxyRuntime::SharedIngress { route }, }, ); return Ok(addrs); } - let runtime = tokio::runtime::Handle::try_current().with_context(|| { + tokio::runtime::Handle::try_current().with_context(|| { format!("failed to create network proxy for environment `{environment_id}`") })?; let listeners = @@ -1272,49 +1291,47 @@ impl NetworkProxy { socks_listener, } = listeners; + let mut listeners = ProxyListeners::new(); 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_task = runtime.spawn(async move { + listeners.spawn(move |guard| async move { http_proxy::run_http_proxy_with_std_listener( http_state, http_listener, http_decider, http_environment_id, + guard, ) .await }); - let socks_task = if self.socks_enabled { + if self.socks_enabled + && let Some(listener) = socks_listener + { let socks_state = self.state.clone(); let socks_decider = self.policy_decider.clone(); let socks_environment_id = Some(environment_id.clone()); let socks5_udp_enabled = self.socks5_udp_enabled; - socks_listener.map(|listener| { - runtime.spawn(async move { - socks5::run_socks5_with_std_listener( - socks_state, - listener, - socks_decider, - socks_environment_id, - socks5_udp_enabled, - ) - .await - }) - }) - } else { - None - }; + listeners.spawn(move |guard| async move { + socks5::run_socks5_with_std_listener( + socks_state, + listener, + socks_decider, + socks_environment_id, + socks5_udp_enabled, + guard, + ) + .await + }); + } proxies.insert( environment_id, EnvironmentProxy { addrs, - runtime: EnvironmentProxyRuntime::ListenerTasks { - http_task, - socks_task, - }, + runtime: ProxyRuntime::Listeners(listeners), }, ); Ok(addrs) @@ -1386,12 +1403,14 @@ impl NetworkProxy { active_route.is_none(), "shared managed Windows proxy route is already running" ); - *active_route = Some(Arc::new(windows_runtime.ingress.register_route( + let route = Arc::new(windows_runtime.ingress.register_route( windows_runtime.http_service.clone(), windows_runtime.socks_service.clone(), - ))); + )); + *active_route = Some(Arc::clone(&route)); drop(active_route); return Ok(NetworkProxyHandle::windows_shared( + route, Arc::clone(&windows_runtime.active_route), Arc::clone(&self.environment_proxies), )); @@ -1401,10 +1420,11 @@ impl NetworkProxy { let http_listener = reserved_listeners.and_then(|listeners| listeners.take_http()); let socks_listener = reserved_listeners.and_then(|listeners| listeners.take_socks()); + let mut listeners = ProxyListeners::new(); let http_state = self.state.clone(); let http_decider = self.policy_decider.clone(); let http_addr = self.http_addr; - let http_task = tokio::spawn(async move { + listeners.spawn(move |guard| async move { match http_listener { Some(listener) => { http_proxy::run_http_proxy_with_std_listener( @@ -1412,6 +1432,7 @@ impl NetworkProxy { listener, http_decider, /*environment_id*/ None, + guard, ) .await } @@ -1421,18 +1442,19 @@ impl NetworkProxy { http_addr, http_decider, /*environment_id*/ None, + guard, ) .await } } }); - let socks_task = if current_cfg.enable_socks5 { + if current_cfg.enable_socks5 { let socks_state = self.state.clone(); let socks_decider = self.policy_decider.clone(); let socks_addr = self.socks_addr; let enable_socks5_udp = current_cfg.enable_socks5_udp; - Some(tokio::spawn(async move { + listeners.spawn(move |guard| async move { match socks_listener { Some(listener) => { socks5::run_socks5_with_std_listener( @@ -1441,6 +1463,7 @@ impl NetworkProxy { socks_decider, /*environment_id*/ None, enable_socks5_udp, + guard, ) .await } @@ -1451,20 +1474,17 @@ impl NetworkProxy { socks_decider, /*environment_id*/ None, enable_socks5_udp, + guard, ) .await } } - })) - } else { - None - }; + }); + } Ok(NetworkProxyHandle { - http_task: Some(http_task), - socks_task, - environment_proxies: self.environment_proxies.clone(), - completed: false, + runtime: Some(ProxyRuntime::Listeners(listeners)), + environment_proxies: Some(Arc::clone(&self.environment_proxies)), #[cfg(target_os = "windows")] windows_active_route: None, }) @@ -1472,10 +1492,8 @@ impl NetworkProxy { } pub struct NetworkProxyHandle { - http_task: Option>>, - socks_task: Option>>, - environment_proxies: Arc>>, - completed: bool, + runtime: Option, + environment_proxies: Option>>>, #[cfg(target_os = "windows")] windows_active_route: Option>>>>, } @@ -1483,10 +1501,8 @@ pub struct NetworkProxyHandle { impl NetworkProxyHandle { fn noop() -> Self { Self { - http_task: Some(tokio::spawn(async { Ok(()) })), - socks_task: None, - environment_proxies: Arc::new(Mutex::new(HashMap::new())), - completed: true, + runtime: None, + environment_proxies: None, #[cfg(target_os = "windows")] windows_active_route: None, } @@ -1494,131 +1510,67 @@ impl NetworkProxyHandle { #[cfg(target_os = "windows")] fn windows_shared( + route: Arc, active_route: Arc>>>, environment_proxies: Arc>>, ) -> Self { Self { - http_task: Some(tokio::spawn(async { - std::future::pending::<()>().await; - Ok(()) - })), - socks_task: None, - environment_proxies, - completed: false, + runtime: Some(ProxyRuntime::SharedIngress { route }), + environment_proxies: Some(environment_proxies), windows_active_route: Some(active_route), } } - #[cfg(target_os = "windows")] - fn deactivate_windows_route(&mut self) { + pub async fn wait(mut self) -> Result<()> { + let result = match self.runtime.as_mut() { + Some(runtime) => runtime.wait().await, + None => Ok(()), + }; + self.shutdown().await?; + result + } + + pub async fn shutdown(mut self) -> Result<()> { + for listeners in self.stop_runtimes() { + listeners.shutdown().await; + } + Ok(()) + } + + fn stop_runtimes(&mut self) -> Vec { + #[cfg(target_os = "windows")] if let Some(active_route) = self.windows_active_route.take() { active_route .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .take(); } + + let environments = self + .environment_proxies + .take() + .map(|proxies| { + std::mem::take( + &mut *proxies + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) + }) + .unwrap_or_default(); + // Stop every runtime before awaiting any one of them. Cancellation of this + // cleanup future must not leave later environments accepting connections. + self.runtime + .take() + .into_iter() + .chain(environments.into_values().map(|proxy| proxy.runtime)) + .filter_map(ProxyRuntime::stop) + .collect() } - - 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 http_result = http_task.await; - let socks_result = match socks_task { - Some(task) => Some(task.await), - None => None, - }; - #[cfg(target_os = "windows")] - self.deactivate_windows_route(); - self.completed = true; - abort_environment_proxies(self.environment_proxies.clone()).await; - http_result??; - if let Some(socks_result) = socks_result { - socks_result??; - } - Ok(()) - } - - pub async fn shutdown(mut self) -> Result<()> { - #[cfg(target_os = "windows")] - self.deactivate_windows_route(); - abort_tasks(self.http_task.take(), self.socks_task.take()).await; - abort_environment_proxies(self.environment_proxies.clone()).await; - self.completed = true; - Ok(()) - } -} - -async fn abort_task(task: Option>>) { - if let Some(task) = task { - task.abort(); - let _ = task.await; - } -} - -async fn abort_tasks( - http_task: Option>>, - socks_task: Option>>, -) { - abort_task(http_task).await; - abort_task(socks_task).await; -} - -async fn abort_environment_proxies( - environment_proxies: Arc>>, -) { - let proxies = { - let mut guard = environment_proxies - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - guard.drain().map(|(_, proxy)| proxy).collect::>() - }; - for proxy in proxies { - match proxy.runtime { - EnvironmentProxyRuntime::ListenerTasks { - http_task, - socks_task, - } => { - abort_task(Some(http_task)).await; - abort_task(socks_task).await; - } - #[cfg(target_os = "windows")] - EnvironmentProxyRuntime::SharedIngress { .. } => {} - } - } -} - -#[cfg(target_os = "windows")] -fn unregister_windows_ingress_environment_routes( - environment_proxies: &Arc>>, -) { - environment_proxies - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .retain(|_, proxy| { - matches!( - &proxy.runtime, - EnvironmentProxyRuntime::ListenerTasks { .. } - ) - }); } impl Drop for NetworkProxyHandle { fn drop(&mut self) { - if self.completed { - return; - } - let http_task = self.http_task.take(); - let socks_task = self.socks_task.take(); - let environment_proxies = self.environment_proxies.clone(); - #[cfg(target_os = "windows")] - { - self.deactivate_windows_route(); - unregister_windows_ingress_environment_routes(&environment_proxies); - } - tokio::spawn(async move { - abort_tasks(http_task, socks_task).await; - abort_environment_proxies(environment_proxies).await; - }); + drop(self.stop_runtimes()); } } diff --git a/codex-rs/network-proxy/src/socks5.rs b/codex-rs/network-proxy/src/socks5.rs index 9817e880d9..098d3feb27 100644 --- a/codex-rs/network-proxy/src/socks5.rs +++ b/codex-rs/network-proxy/src/socks5.rs @@ -1,6 +1,7 @@ use crate::attribution::BindConnectionAttribution; use crate::config::NetworkMode; use crate::connect_policy::TargetCheckedTcpConnector; +use crate::connection_lifecycle::CancelOnShutdown; use crate::mitm; use crate::network_policy::BlockDecisionAuditEventArgs; use crate::network_policy::NetworkDecision; @@ -29,6 +30,7 @@ use rama_core::error::BoxError; use rama_core::extensions::Extensions; use rama_core::extensions::ExtensionsMut; use rama_core::extensions::ExtensionsRef; +use rama_core::graceful::ShutdownGuard; use rama_core::service::BoxService; use rama_core::service::service_fn; use rama_net::address::HostWithPort; @@ -67,6 +69,7 @@ pub async fn run_socks5( policy_decider: Option>, environment_id: Option, enable_socks5_udp: bool, + guard: ShutdownGuard, ) -> Result<()> { let listener = TcpListener::build() .bind(addr) @@ -82,6 +85,7 @@ pub async fn run_socks5( policy_decider, environment_id, enable_socks5_udp, + guard, ) .await } @@ -92,6 +96,7 @@ pub async fn run_socks5_with_std_listener( policy_decider: Option>, environment_id: Option, enable_socks5_udp: bool, + guard: ShutdownGuard, ) -> Result<()> { let listener = TcpListener::try_from(listener).context("convert std listener to SOCKS5 proxy listener")?; @@ -101,6 +106,7 @@ pub async fn run_socks5_with_std_listener( policy_decider, environment_id, enable_socks5_udp, + guard, ) .await } @@ -111,6 +117,7 @@ async fn run_socks5_with_listener( policy_decider: Option>, environment_id: Option, enable_socks5_udp: bool, + guard: ShutdownGuard, ) -> Result<()> { let addr = listener .local_addr() @@ -131,12 +138,15 @@ async fn run_socks5_with_listener( } listener - .serve(socks5_proxy_service( - state, - policy_decider, - environment_id, - enable_socks5_udp, - )) + .serve_graceful( + guard, + CancelOnShutdown::new(socks5_proxy_service( + state, + policy_decider, + environment_id, + enable_socks5_udp, + )), + ) .await; Ok(()) }