diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f6879c58e6e..cc9e6bededf 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1944,13 +1944,16 @@ dependencies = [ "futures", "http 1.4.0", "pretty_assertions", + "rcgen", "regex-lite", "reqwest 0.12.28", + "rustls", "schemars 0.8.22", "serde", "serde_json", "thiserror 2.0.18", "tokio", + "tokio-rustls", "tokio-test", "tokio-tungstenite", "tokio-util", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index ad08ce8e024..dcb549acbb1 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -425,6 +425,7 @@ thiserror = "2.0.17" time = "0.3.47" tiny_http = "0.12" tokio = "1" +tokio-rustls = "0.26.4" tokio-stream = "0.1.18" tokio-test = "0.4" tokio-tungstenite = { version = "0.28.0", features = [ diff --git a/codex-rs/cli/src/doctor.rs b/codex-rs/cli/src/doctor.rs index ef587cd1d4e..f695274675b 100644 --- a/codex-rs/cli/src/doctor.rs +++ b/codex-rs/cli/src/doctor.rs @@ -2371,9 +2371,11 @@ async fn websocket_reachability_check( HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), ); let client = ResponsesWebsocketClient::new(api_provider, api_auth); + let http_client_factory = config.http_client_factory(); match tokio::time::timeout( provider.websocket_connect_timeout(), client.probe_handshake( + &http_client_factory, extra_headers, default_headers(), WEBSOCKET_IMMEDIATE_CLOSE_GRACE, diff --git a/codex-rs/codex-api/Cargo.toml b/codex-rs/codex-api/Cargo.toml index 084b058e126..fe1f3864e7b 100644 --- a/codex-rs/codex-api/Cargo.toml +++ b/codex-rs/codex-api/Cargo.toml @@ -16,11 +16,13 @@ codex-utils-rustls-provider = { workspace = true } futures = { workspace = true } http = { workspace = true } reqwest = { workspace = true, features = ["json", "stream"] } +rustls = { workspace = true } schemars = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["fs", "macros", "net", "rt", "sync", "time"] } +tokio-rustls = { workspace = true } tokio-tungstenite = { workspace = true } tungstenite = { workspace = true } tracing = { workspace = true } @@ -33,6 +35,7 @@ url = { workspace = true } anyhow = { workspace = true } assert_matches = { workspace = true } pretty_assertions = { workspace = true } +rcgen = { workspace = true } tokio-test = { workspace = true } wiremock = { workspace = true } reqwest = { workspace = true } diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index cce056f638b..17fd7ffd9a3 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -10,8 +10,11 @@ use crate::safety_buffering::treatment_from_headers; use crate::sse::ResponsesStreamEvent; use crate::sse::process_responses_event; use crate::telemetry::WebsocketTelemetry; +use crate::websocket_connector; +use crate::websocket_connector::ConnectedWebSocket; use codex_client::TransportError; -use codex_http_client::maybe_build_rustls_client_config_with_custom_ca; +use codex_http_client::HttpClientFactory; +use codex_http_client::build_rustls_client_config_with_custom_ca; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use futures::SinkExt; use futures::StreamExt; @@ -25,14 +28,13 @@ use serde_json::map::Map as JsonMap; use std::sync::Arc; use std::sync::OnceLock; use std::time::Duration; -use tokio::net::TcpStream; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::time::Instant; -use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; -use tokio_tungstenite::connect_async_tls_with_config; use tokio_tungstenite::tungstenite::Error as WsError; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; @@ -62,7 +64,17 @@ enum WsCommand { } impl WsStream { - fn new(inner: WebSocketStream>) -> Self { + fn new(inner: ConnectedWebSocket) -> Self { + match inner { + ConnectedWebSocket::TransportDefault(inner) => Self::new_inner(inner), + ConnectedWebSocket::Routed(inner) => Self::new_inner(inner), + } + } + + fn new_inner(inner: WebSocketStream) -> Self + where + S: AsyncRead + AsyncWrite + Send + Unpin + 'static, + { let (tx_command, mut rx_command) = mpsc::channel::(32); let (tx_message, rx_message) = mpsc::unbounded_channel::>(); @@ -334,6 +346,7 @@ impl ResponsesWebsocketClient { )] pub async fn connect( &self, + http_client_factory: &HttpClientFactory, extra_headers: HeaderMap, default_headers: HeaderMap, turn_state: Option>>, @@ -349,7 +362,7 @@ impl ResponsesWebsocketClient { self.auth.add_auth_headers(&mut headers); let (stream, _status, server_reasoning_included, models_etag, server_model) = - connect_websocket(ws_url, headers, turn_state.clone()).await?; + connect_websocket(ws_url, headers, http_client_factory, turn_state.clone()).await?; Ok(ResponsesWebsocketConnection::new( stream, self.provider.stream_idle_timeout, @@ -369,6 +382,7 @@ impl ResponsesWebsocketClient { /// a usable connection from a policy rejection that closes right away. pub async fn probe_handshake( &self, + http_client_factory: &HttpClientFactory, extra_headers: HeaderMap, default_headers: HeaderMap, immediate_close_timeout: Duration, @@ -383,7 +397,13 @@ impl ResponsesWebsocketClient { self.auth.add_auth_headers(&mut headers); let (mut stream, status, reasoning_included, models_etag, server_model) = - connect_websocket(ws_url.clone(), headers, /*turn_state*/ None).await?; + connect_websocket( + ws_url.clone(), + headers, + http_client_factory, + /*turn_state*/ None, + ) + .await?; let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next()) .await .ok() @@ -437,6 +457,7 @@ fn merge_request_headers( async fn connect_websocket( url: Url, headers: HeaderMap, + http_client_factory: &HttpClientFactory, turn_state: Option>>, ) -> Result<(WsStream, StatusCode, bool, Option, Option), ApiError> { ensure_rustls_crypto_provider(); @@ -448,20 +469,15 @@ async fn connect_websocket( .map_err(|err| ApiError::Stream(format!("failed to build websocket request: {err}")))?; request.headers_mut().extend(headers); - // Secure websocket traffic needs the same custom-CA policy as reqwest-based HTTPS traffic. - // If a Codex-specific CA bundle is configured, build an explicit rustls connector so this - // websocket path does not fall back to tungstenite's default native-roots-only behavior. - let connector = maybe_build_rustls_client_config_with_custom_ca() - .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))? - .map(tokio_tungstenite::Connector::Rustls); - - let response = connect_async_tls_with_config( - request, - Some(websocket_config()), - false, // `false` means "do not disable Nagle", which is tungstenite's recommended default. - connector, - ) - .await; + // Secure websocket traffic and TLS-encrypted proxy connections need the same custom-CA policy + // as reqwest-based HTTPS traffic, so build one rustls config for both handshakes. + let tls_config = build_rustls_client_config_with_custom_ca() + .map_err(|err| ApiError::Stream(format!("failed to configure websocket TLS: {err}")))?; + + let proxy_route = http_client_factory.resolve_proxy_route(url.as_str()); + let response = + websocket_connector::connect(request, Some(websocket_config()), tls_config, proxy_route) + .await; let (stream, response) = match response { Ok((stream, response)) => { diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 1d1a4c8c6e5..4a5c63391e4 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -12,6 +12,7 @@ pub(crate) mod safety_buffering; pub(crate) mod search; pub(crate) mod sse; pub(crate) mod telemetry; +mod websocket_connector; pub use crate::requests::headers::build_session_headers; pub use codex_client::RequestTelemetry; diff --git a/codex-rs/codex-api/src/websocket_connector.rs b/codex-rs/codex-api/src/websocket_connector.rs new file mode 100644 index 00000000000..a80c23d53ce --- /dev/null +++ b/codex-rs/codex-api/src/websocket_connector.rs @@ -0,0 +1,251 @@ +//! Route-aware WebSocket connection setup. + +use std::collections::VecDeque; +use std::io; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use codex_http_client::OutboundProxyRoute; +use futures::StreamExt; +use futures::stream::FuturesUnordered; +use rustls::ClientConfig; +use rustls::pki_types::ServerName; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::net::TcpStream; +use tokio::time::Instant; +use tokio::time::sleep_until; +use tokio_rustls::TlsConnector; +use tokio_tungstenite::Connector; +use tokio_tungstenite::MaybeTlsStream; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::client_async_tls_with_config; +use tokio_tungstenite::connect_async_tls_with_config; +use tokio_tungstenite::proxy::connect_via_proxy; +use tokio_tungstenite::tungstenite::Error as WsError; +use tokio_tungstenite::tungstenite::error::TlsError; +use tokio_tungstenite::tungstenite::error::UrlError; +use tokio_tungstenite::tungstenite::handshake::client::Request; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; +use tokio_tungstenite::tungstenite::proxy::ProxyConfig; + +const HAPPY_EYEBALLS_DELAY: Duration = Duration::from_millis(250); + +/// An established WebSocket whose concrete network stream depends on route selection. +pub(crate) enum ConnectedWebSocket { + /// A connection established by `tokio-tungstenite`'s default transport. + TransportDefault(WebSocketStream>), + /// A connection established over a caller-selected direct or proxy stream. + Routed(WebSocketStream>>), +} + +/// Async network I/O that can be carried through a proxy tunnel and target TLS handshake. +pub(crate) trait AsyncIo: AsyncRead + AsyncWrite + Send + Unpin {} + +impl AsyncIo for T where T: AsyncRead + AsyncWrite + Send + Unpin {} + +struct ProxyEndpoint { + config: ProxyConfig, + tls: bool, +} + +impl ProxyEndpoint { + fn parse(url: &str) -> Result { + let mut parsed_url = url::Url::parse(url).map_err(|_| invalid_proxy_config())?; + let tls = parsed_url.scheme() == "https"; + if tls { + if parsed_url.port().is_none() { + parsed_url + .set_port(Some(443)) + .map_err(|_| invalid_proxy_config())?; + } + parsed_url + .set_scheme("http") + .map_err(|_| invalid_proxy_config())?; + } + let config = ProxyConfig::parse(parsed_url.as_str()).map_err(|error| match error { + WsError::Url(UrlError::UnsupportedProxyScheme) => error, + _ => invalid_proxy_config(), + })?; + Ok(Self { config, tls }) + } +} + +/// Connects a WebSocket using the resolved outbound proxy route. +pub(crate) async fn connect( + request: Request, + config: Option, + tls_config: Arc, + proxy_route: OutboundProxyRoute, +) -> Result<(ConnectedWebSocket, Response), WsError> { + let stream: Box = match proxy_route { + OutboundProxyRoute::TransportDefault => { + let (stream, response) = connect_async_tls_with_config( + request, + config, + false, // Preserve tungstenite's recommended Nagle default. + Some(Connector::Rustls(tls_config)), + ) + .await?; + return Ok((ConnectedWebSocket::TransportDefault(stream), response)); + } + OutboundProxyRoute::Direct => { + let host = request + .uri() + .host() + .ok_or(WsError::Url(UrlError::NoHostName))?; + let port = websocket_port(&request)?; + let address = host_port(host, port); + Box::new(connect_tcp(address).await.map_err(WsError::Io)?) + } + OutboundProxyRoute::Proxy { url } => { + let proxy = ProxyEndpoint::parse(&url)?; + let host = request + .uri() + .host() + .ok_or(WsError::Url(UrlError::NoHostName))?; + let port = websocket_port(&request)?; + let stream = connect_tcp(proxy.config.authority()) + .await + .map_err(WsError::Io)?; + let stream: Box = if proxy.tls { + let server_name = ServerName::try_from(proxy.config.host.clone()) + .map_err(|_| WsError::Tls(TlsError::InvalidDnsName))?; + let stream = TlsConnector::from(Arc::clone(&tls_config)) + .connect(server_name, stream) + .await + .map_err(WsError::Io)?; + Box::new(stream) + } else { + Box::new(stream) + }; + connect_via_proxy(stream, &proxy.config, host, port).await? + } + }; + + let (stream, response) = + client_async_tls_with_config(request, stream, config, Some(Connector::Rustls(tls_config))) + .await?; + Ok((ConnectedWebSocket::Routed(stream), response)) +} + +fn invalid_proxy_config() -> WsError { + WsError::Url(UrlError::InvalidProxyConfig("".to_string())) +} + +fn websocket_port(request: &Request) -> Result { + request + .uri() + .port_u16() + .or_else(|| match request.uri().scheme_str() { + Some("ws") => Some(80), + Some("wss") => Some(443), + _ => None, + }) + .ok_or(WsError::Url(UrlError::UnsupportedUrlScheme)) +} + +fn host_port(host: &str, port: u16) -> String { + if host.contains(':') && !host.starts_with('[') { + format!("[{host}]:{port}") + } else { + format!("{host}:{port}") + } +} + +/// Preserves the Happy Eyeballs behavior of `tokio-tungstenite`'s built-in dialer. +/// +/// Explicit routes require a caller-provided TCP stream, so they cannot use that private dialer +/// directly. Keep the same family interleaving and 250 ms attempt delay here to avoid regressing +/// direct or proxy connections when one address family is unreachable. +async fn connect_tcp(address: String) -> io::Result { + let mut resolved_addresses = tokio::net::lookup_host(address).await?; + let Some(first_address) = resolved_addresses.next() else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "could not resolve to any address", + )); + }; + + let first_is_ipv4 = first_address.is_ipv4(); + let mut preferred = VecDeque::new(); + let mut alternate = VecDeque::new(); + for address in resolved_addresses { + if address.is_ipv4() == first_is_ipv4 { + preferred.push_back(address); + } else { + alternate.push_back(address); + } + } + + let mut addresses = VecDeque::::new(); + while !preferred.is_empty() || !alternate.is_empty() { + if let Some(address) = alternate.pop_front() { + addresses.push_back(address); + } + if let Some(address) = preferred.pop_front() { + addresses.push_back(address); + } + } + + let mut attempts = FuturesUnordered::new(); + attempts.push(TcpStream::connect(first_address)); + let mut next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY; + let mut last_error = None; + + loop { + if addresses.is_empty() { + match attempts.next().await { + Some(Ok(stream)) => return Ok(stream), + Some(Err(error)) => { + if attempts.is_empty() { + return Err(error); + } + last_error = Some(error); + } + None => { + return Err(last_error.unwrap_or_else(|| { + io::Error::other("connection attempts ended without an error") + })); + } + } + continue; + } + + tokio::select! { + result = attempts.next() => { + match result { + Some(Ok(stream)) => return Ok(stream), + Some(Err(error)) => { + last_error = Some(error); + let address = take_next_address(&mut addresses)?; + attempts.push(TcpStream::connect(address)); + next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY; + } + None => { + let address = take_next_address(&mut addresses)?; + attempts.push(TcpStream::connect(address)); + next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY; + } + } + } + _ = sleep_until(next_attempt_at) => { + let address = take_next_address(&mut addresses)?; + attempts.push(TcpStream::connect(address)); + next_attempt_at = Instant::now() + HAPPY_EYEBALLS_DELAY; + } + } + } +} + +fn take_next_address(addresses: &mut VecDeque) -> io::Result { + addresses + .pop_front() + .ok_or_else(|| io::Error::other("connection address queue unexpectedly empty")) +} + +#[cfg(test)] +#[path = "websocket_connector_tests.rs"] +mod tests; diff --git a/codex-rs/codex-api/src/websocket_connector_tests.rs b/codex-rs/codex-api/src/websocket_connector_tests.rs new file mode 100644 index 00000000000..f1a778a513b --- /dev/null +++ b/codex-rs/codex-api/src/websocket_connector_tests.rs @@ -0,0 +1,174 @@ +use std::net::SocketAddr; +use std::sync::Arc; + +use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use pretty_assertions::assert_eq; +use rcgen::CertifiedKey; +use rcgen::generate_simple_self_signed; +use rustls::ClientConfig; +use rustls::RootCertStore; +use rustls::ServerConfig; +use rustls::pki_types::PrivateKeyDer; +use rustls::pki_types::PrivatePkcs8KeyDer; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::sync::Mutex; +use tokio::task::JoinHandle; +use tokio_rustls::TlsAcceptor; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; + +use super::*; + +#[tokio::test] +async fn direct_route_connects_secure_websocket() { + let (tls_config, acceptor) = test_tls_configs(); + let (target_addr, target_task) = start_tls_websocket_server(acceptor).await; + let request = format!("wss://localhost:{}/v1/responses", target_addr.port()) + .into_client_request() + .expect("websocket request should build"); + + let (websocket, _) = connect( + request, + /*config*/ None, + tls_config, + OutboundProxyRoute::Direct, + ) + .await + .expect("direct websocket handshake should succeed"); + drop(websocket); + + target_task.await.expect("target task should finish"); +} + +#[tokio::test] +async fn http_proxy_tunnels_secure_websocket_before_handshake() { + assert_proxy_tunnels_secure_websocket(/*proxy_tls*/ false).await; +} + +#[tokio::test] +async fn https_proxy_tunnels_secure_websocket_before_handshake() { + assert_proxy_tunnels_secure_websocket(/*proxy_tls*/ true).await; +} + +async fn assert_proxy_tunnels_secure_websocket(proxy_tls: bool) { + let (tls_config, acceptor) = test_tls_configs(); + let (target_addr, target_task) = start_tls_websocket_server(acceptor.clone()).await; + + let proxy_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("proxy listener should bind"); + let proxy_addr = proxy_listener + .local_addr() + .expect("proxy listener should have an address"); + let connect_request = Arc::new(Mutex::new(None)); + let proxy_connect_request = Arc::clone(&connect_request); + let proxy_task = tokio::spawn(async move { + let (client, _) = proxy_listener.accept().await.expect("proxy should accept"); + let mut client: Box = if proxy_tls { + Box::new( + acceptor + .accept(client) + .await + .expect("proxy TLS handshake should succeed"), + ) + } else { + Box::new(client) + }; + let mut request = Vec::new(); + let mut byte = [0_u8; 1]; + while !request.ends_with(b"\r\n\r\n") { + client + .read_exact(&mut byte) + .await + .expect("proxy should read CONNECT request"); + request.push(byte[0]); + } + *proxy_connect_request.lock().await = + Some(String::from_utf8(request).expect("CONNECT request should contain valid UTF-8")); + + let mut target = tokio::net::TcpStream::connect(target_addr) + .await + .expect("proxy should connect to target"); + client + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .expect("proxy should acknowledge CONNECT"); + let _ = tokio::io::copy_bidirectional(&mut client, &mut target).await; + }); + + let target_authority = format!("localhost:{}", target_addr.port()); + let proxy_scheme = if proxy_tls { "https" } else { "http" }; + let request = format!("wss://{target_authority}/v1/responses") + .into_client_request() + .expect("websocket request should build"); + let (websocket, _) = connect( + request, + /*config*/ None, + tls_config, + OutboundProxyRoute::Proxy { + url: format!("{proxy_scheme}://localhost:{}", proxy_addr.port()), + }, + ) + .await + .expect("proxied websocket handshake should succeed"); + drop(websocket); + + target_task.await.expect("target task should finish"); + proxy_task.await.expect("proxy task should finish"); + let request = connect_request + .lock() + .await + .clone() + .expect("proxy should record CONNECT request"); + let expected_request_line = format!("CONNECT {target_authority} HTTP/1.1"); + assert_eq!(request.lines().next(), Some(expected_request_line.as_str())); +} + +async fn start_tls_websocket_server(acceptor: TlsAcceptor) -> (SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("target listener should bind"); + let address = listener + .local_addr() + .expect("target listener should have an address"); + let task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("target should accept"); + let stream = acceptor + .accept(stream) + .await + .expect("target TLS handshake should succeed"); + let mut websocket = accept_async(stream) + .await + .expect("target websocket handshake should succeed"); + let _ = websocket.close(None).await; + }); + (address, task) +} + +fn test_tls_configs() -> (Arc, TlsAcceptor) { + ensure_rustls_crypto_provider(); + let CertifiedKey { cert, signing_key } = + generate_simple_self_signed(vec!["localhost".to_string()]) + .expect("test certificate should generate"); + let certificate = cert.der().clone(); + let private_key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(signing_key.serialize_der())); + let server_config = ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![certificate.clone()], private_key) + .expect("test server config should build"); + + let mut roots = RootCertStore::empty(); + roots + .add(certificate) + .expect("test certificate should be trusted"); + let client_config = ClientConfig::builder() + .with_root_certificates(roots) + .with_no_client_auth(); + + ( + Arc::new(client_config), + TlsAcceptor::from(Arc::new(server_config)), + ) +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 0bf33b39933..5d29979c37d 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1006,6 +1006,7 @@ impl ModelClient { let result = match tokio::time::timeout( websocket_connect_timeout, ApiWebSocketResponsesClient::new(api_provider, api_auth).connect( + &self.http_client_factory, headers, codex_login::default_client::default_headers(), /*turn_state*/ None,