diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ae86e1ac43..06a13bee36 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3605,6 +3605,7 @@ dependencies = [ "rustls-platform-verifier", "serde", "serde_json", + "serde_urlencoded", "sha2 0.10.9", "system-configuration", "tempfile", diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index 06d57377b1..3c9116d338 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -13,8 +13,8 @@ use codex_api::SharedAuthProvider; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; +use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; use codex_http_client::RouteAwareRequestError; use codex_login::CodexAuth; use codex_login::default_client::get_codex_user_agent; @@ -271,13 +271,13 @@ impl Client { h } - fn request(&self, method: Method, url: &str) -> RouteAwareRequestBuilder { + fn request(&self, method: Method, url: &str) -> RequestBuilder { self.http.request(method, url) } async fn exec_request( &self, - req: RouteAwareRequestBuilder, + req: RequestBuilder, method: &str, url: &str, ) -> Result<(String, String)> { @@ -289,7 +289,7 @@ impl Client { .and_then(|v| v.to_str().ok()) .unwrap_or("") .to_string(); - let body = res.text().await.unwrap_or_default(); + let body = res.text().await.map_err(anyhow::Error::from)?; if !status.is_success() { anyhow::bail!("{method} {url} failed: {status}; content-type={ct}; body={body}"); } @@ -298,7 +298,7 @@ impl Client { async fn exec_request_detailed( &self, - req: RouteAwareRequestBuilder, + req: RequestBuilder, method: &str, url: &str, ) -> std::result::Result<(String, String), RequestError> { @@ -310,7 +310,7 @@ impl Client { .and_then(|v| v.to_str().ok()) .unwrap_or("") .to_string(); - let body = res.text().await.unwrap_or_default(); + let body = res.text().await.map_err(anyhow::Error::from)?; if !status.is_success() { return Err(RequestError::UnexpectedStatus { method: method.to_string(), diff --git a/codex-rs/cli/src/doctor/network.rs b/codex-rs/cli/src/doctor/network.rs index d6cff63662..a413fc8dba 100644 --- a/codex-rs/cli/src/doctor/network.rs +++ b/codex-rs/cli/src/doctor/network.rs @@ -128,7 +128,6 @@ pub(super) async fn probe_status( .timeout(Duration::from_secs(8)) .send() .await - .map_err(RouteAwareRequestError::from) } else { client .request(method, url) diff --git a/codex-rs/codex-api/src/api_bridge.rs b/codex-rs/codex-api/src/api_bridge.rs index 15c9cb9245..a2cda1566e 100644 --- a/codex-rs/codex-api/src/api_bridge.rs +++ b/codex-rs/codex-api/src/api_bridge.rs @@ -226,6 +226,7 @@ pub fn map_api_error(err: ApiError) -> CodexErr { request_id: None, }), TransportError::Timeout => CodexErr::RequestTimeout, + TransportError::Policy(denied) => CodexErr::Fatal(denied.to_string()), TransportError::Connection(source) => { CodexErr::ConnectionFailed(ConnectionFailedError { source }) } diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index 7a2d391e60..58ade3b75f 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -23,6 +23,7 @@ use crate::error::ApiError; use crate::provider::Provider; use codex_client::backoff; use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyDenied; use codex_http_client::maybe_build_rustls_client_config_with_custom_ca; use codex_protocol::protocol::ConversationTextParams; use codex_protocol::protocol::ConversationTextRole; @@ -65,6 +66,7 @@ const TRUNCATED_TRANSCRIPT_PREFIX: &str = "…"; struct WsStream { tx_command: mpsc::Sender, pump_task: tokio::task::JoinHandle<()>, + terminal_policy_error: tokio::sync::watch::Receiver>, } enum WsCommand { @@ -84,8 +86,15 @@ impl WsStream { let (tx_command, mut rx_command) = mpsc::channel::(32); let (tx_message, rx_message) = async_channel::unbounded::>(); + let (terminal_policy_tx, terminal_policy_error) = tokio::sync::watch::channel(None); let pump_task = tokio::spawn(async move { let mut inner = inner; + let report_write_policy_error = |error: &WsError| { + if let Some(denied) = codex_websocket_client::network_policy_denial(error) { + terminal_policy_tx.send_replace(Some(denied)); + let _ = tx_message.try_send(Err(Self::policy_error(denied))); + } + }; loop { tokio::select! { command = rx_command.recv() => { @@ -99,6 +108,7 @@ impl WsStream { let should_break = result.is_err(); if let Err(err) = &result { error!("realtime websocket send failed: {err}"); + report_write_policy_error(err); } let _ = tx_result.send(result); if should_break { @@ -110,6 +120,7 @@ impl WsStream { let result = inner.send(Message::Close(None)).await; if let Err(err) = &result { error!("realtime websocket close failed: {err}"); + report_write_policy_error(err); } let _ = tx_result.send(result); break; @@ -125,6 +136,7 @@ impl WsStream { trace!(payload_len = payload.len(), "realtime websocket received ping"); if let Err(err) = inner.send(Message::Pong(payload)).await { error!("realtime websocket failed to send pong: {err}"); + terminal_policy_tx.send_replace(codex_websocket_client::network_policy_denial(&err)); let _ = tx_message.send(Err(err)).await; break; } @@ -162,6 +174,7 @@ impl WsStream { } Err(err) => { error!("realtime websocket receive failed: {err}"); + terminal_policy_tx.send_replace(codex_websocket_client::network_policy_denial(&err)); let _ = tx_message.send(Err(err)).await; break; } @@ -176,6 +189,7 @@ impl WsStream { Self { tx_command, pump_task, + terminal_policy_error, }, rx_message, ) @@ -187,9 +201,23 @@ impl WsStream { ) -> Result<(), WsError> { let (tx_result, rx_result) = oneshot::channel(); if self.tx_command.send(make_command(tx_result)).await.is_err() { - return Err(WsError::ConnectionClosed); + return Err(self.closed_error()); } - rx_result.await.unwrap_or(Err(WsError::ConnectionClosed)) + rx_result.await.unwrap_or_else(|_| Err(self.closed_error())) + } + + fn closed_error(&self) -> WsError { + match *self.terminal_policy_error.borrow() { + Some(denied) => Self::policy_error(denied), + None => WsError::ConnectionClosed, + } + } + + fn policy_error(denied: NetworkPolicyDenied) -> WsError { + WsError::Io(std::io::Error::new( + std::io::ErrorKind::PermissionDenied, + denied, + )) } async fn send(&self, message: Message) -> Result<(), WsError> { @@ -481,8 +509,8 @@ impl RealtimeWebsocketWriter { pub async fn send_payload(&self, payload: String) -> Result<(), ApiError> { if self.is_closed.load(Ordering::SeqCst) { - return Err(ApiError::Stream( - "realtime websocket connection is closed".to_string(), + return Err(map_realtime_websocket_connect_error( + self.stream.closed_error(), )); } @@ -490,7 +518,7 @@ impl RealtimeWebsocketWriter { self.stream .send(Message::Text(payload.into())) .await - .map_err(|err| ApiError::Stream(format!("failed to send realtime request: {err}")))?; + .map_err(map_realtime_websocket_connect_error)?; Ok(()) } } @@ -519,9 +547,7 @@ impl RealtimeWebsocketEvents { Ok(Err(err)) => { self.is_closed.store(true, Ordering::SeqCst); error!("realtime websocket read failed: {err}"); - return Err(ApiError::Stream(format!( - "failed to read websocket message: {err}" - ))); + return Err(map_realtime_websocket_connect_error(err)); } Err(_) => { self.is_closed.store(true, Ordering::SeqCst); @@ -918,6 +944,9 @@ impl RealtimeWebsocketClient { .await; match result { Ok(connection) => return Ok(connection), + Err(err @ ApiError::Transport(codex_http_client::TransportError::Policy(_))) => { + return Err(err); + } Err(err) if webrtc_sideband_session_ended(&err) => return Err(err), Err(err) if attempt < self.provider.retry.max_attempts => { let delay = backoff(self.provider.retry.base_delay, attempt + 1); @@ -1081,6 +1110,9 @@ fn webrtc_sideband_session_ended(err: &ApiError) -> bool { } fn map_realtime_websocket_connect_error(err: WsError) -> ApiError { + if let Some(denied) = codex_websocket_client::network_policy_denial(&err) { + return ApiError::Transport(codex_http_client::TransportError::Policy(denied)); + } match err { WsError::Http(response) => ApiError::Api { status: response.status(), @@ -1261,7 +1293,7 @@ mod tests { use super::*; use crate::endpoint::realtime_websocket::protocol::RealtimeTranscriptEntry; use crate::provider::RetryConfig; - use codex_http_client::OutboundProxyPolicy; + use codex_http_client::OutboundProxyPolicy::ReqwestDefault; use codex_protocol::protocol::RealtimeHandoffRequested; use codex_protocol::protocol::RealtimeInputAudioSpeechStarted; use codex_protocol::protocol::RealtimeNoopRequested; @@ -1283,6 +1315,10 @@ mod tests { use tokio_tungstenite::tungstenite::Message; use tungstenite::protocol::CloseFrame; + fn test_client(provider: Provider) -> RealtimeWebsocketClient { + RealtimeWebsocketClient::new(provider, HttpClientFactory::new(ReqwestDefault)) + } + #[test] fn parse_session_updated_event() { let payload = json!({ @@ -1531,6 +1567,59 @@ mod tests { events.next_event().await.is_err() } + #[tokio::test] + async fn writer_policy_failure_reaches_realtime_events() { + use codex_http_client::DestinationPolicy::Unrestricted; + use codex_http_client::NetworkPolicyController; + use codex_http_client::NetworkPolicyDenied::Revoked; + use codex_http_client::TransportError::Policy; + use codex_websocket_client::network_policy_denial; + use tokio::time::timeout; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/", listener.local_addr().unwrap()); + let (writing, started) = oneshot::channel(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let socket = accept_async(socket).await.unwrap(); + socket.get_ref().peek(&mut [0]).await.unwrap(); + writing.send(()).unwrap(); + std::future::pending::<()>().await; + }); + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish(policy.revision(), Unrestricted); + let factory = HttpClientFactory::new(ReqwestDefault).with_network_policy(policy.clone()); + let connector = WebSocketConnector::new(&factory).unwrap(); + let request = url.into_client_request().unwrap(); + let connected = connector.connect(request, websocket_config()); + let (socket, _) = connected.await.unwrap(); + let (stream, rx_message) = WsStream::new(socket); + let parser = RealtimeEventParser::V1; + let transcript = RealtimeTranscriptState::default(); + let connection = RealtimeWebsocketConnection::new(stream, rx_message, parser, transcript); + let stream = Arc::clone(&connection.writer.stream); + let payload = Message::Binary(vec![0; 8 * 1024 * 1024].into()); + let writer = tokio::spawn(async move { + loop { + if let Err(error) = stream.send(payload.clone()).await { + return error; + } + } + }); + let deadline = Duration::from_secs(/*secs*/ 5); + timeout(deadline, started).await.unwrap().unwrap(); + policy.invalidate(); + let writer = timeout(deadline, writer).await.unwrap(); + let error = writer.unwrap(); + assert_eq!(network_policy_denial(&error), Some(Revoked)); + let next = timeout(deadline, connection.next_event()); + let error = next.await.unwrap().unwrap_err(); + assert!(matches!(error, ApiError::Transport(Policy(Revoked)))); + assert!(connection.next_event().await.unwrap().is_none()); + server.abort(); + } + #[tokio::test] async fn websocket_termination_classification_preserves_protocol_behavior() { use CloseCode::Away; @@ -2220,23 +2309,20 @@ mod tests { #[test] fn webrtc_frameless_sideband_ignores_provider_base_url() { - let client = RealtimeWebsocketClient::new( - Provider { - name: "chatgpt".to_string(), - base_url: "https://chatgpt.com/backend-api/codex".to_string(), - query_params: None, - headers: HeaderMap::new(), - retry: RetryConfig { - max_attempts: 0, - base_delay: Duration::ZERO, - retry_429: false, - retry_5xx: false, - retry_transport: false, - }, - stream_idle_timeout: Duration::from_secs(/*secs*/ 5), + let client = test_client(Provider { + name: "chatgpt".to_string(), + base_url: "https://chatgpt.com/backend-api/codex".to_string(), + query_params: None, + headers: HeaderMap::new(), + retry: RetryConfig { + max_attempts: 0, + base_delay: Duration::ZERO, + retry_429: false, + retry_5xx: false, + retry_transport: false, }, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + stream_idle_timeout: Duration::from_secs(5), + }); let url = client .webrtc_sideband_url( @@ -2433,10 +2519,7 @@ mod tests { }, stream_idle_timeout: Duration::from_secs(5), }; - let client = RealtimeWebsocketClient::new( - provider, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(provider); let connection = client .connect( RealtimeSessionConfig { @@ -2762,10 +2845,7 @@ mod tests { }, stream_idle_timeout: Duration::from_secs(5), }; - let client = RealtimeWebsocketClient::new( - provider, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(provider); let connection = client .connect( RealtimeSessionConfig { @@ -2892,10 +2972,7 @@ mod tests { }, stream_idle_timeout: Duration::from_secs(5), }; - let client = RealtimeWebsocketClient::new( - provider, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(provider); let connection = client .connect( RealtimeSessionConfig { @@ -3001,10 +3078,7 @@ mod tests { }, stream_idle_timeout: Duration::from_secs(5), }; - let client = RealtimeWebsocketClient::new( - provider, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(provider); let connection = client .connect( RealtimeSessionConfig { @@ -3096,10 +3170,7 @@ mod tests { }, stream_idle_timeout: Duration::from_secs(5), }; - let client = RealtimeWebsocketClient::new( - provider, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(provider); let connection = client .connect( RealtimeSessionConfig { diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 48e99d0106..4fe0071ac9 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -125,17 +125,22 @@ impl WsStream { } async fn request( - &self, + &mut self, make_command: impl FnOnce(oneshot::Sender>) -> WsCommand, ) -> Result<(), WsError> { let (tx_result, rx_result) = oneshot::channel(); - if self.tx_command.send(make_command(tx_result)).await.is_err() { - return Err(WsError::ConnectionClosed); + if self.tx_command.send(make_command(tx_result)).await.is_ok() + && let Ok(result) = rx_result.await + { + return result; } - rx_result.await.unwrap_or(Err(WsError::ConnectionClosed)) + while let Some(message) = self.rx_message.recv().await { + message?; + } + Err(WsError::ConnectionClosed) } - async fn send(&self, message: Message) -> Result<(), WsError> { + async fn send(&mut self, message: Message) -> Result<(), WsError> { self.request(|tx_result| WsCommand::Send { message, tx_result }) .await } @@ -441,9 +446,7 @@ impl ResponsesWebsocketClient { .ok() .flatten() .transpose() - .map_err(|err| { - ApiError::Stream(format!("failed to read websocket probe event: {err}")) - })? + .map_err(map_ws_stream_error)? .and_then(immediate_close_from_message); Ok(ResponsesWebsocketProbe { @@ -548,7 +551,17 @@ fn websocket_config() -> WebSocketConfig { config } +fn map_ws_stream_error(error: WsError) -> ApiError { + match codex_websocket_client::network_policy_denial(&error) { + Some(denied) => ApiError::Transport(TransportError::Policy(denied)), + None => ApiError::Stream(error.to_string()), + } +} + fn map_ws_error(err: WsError, url: &Url) -> ApiError { + if let Some(denied) = codex_websocket_client::network_policy_denial(&err) { + return ApiError::Transport(TransportError::Policy(denied)); + } match err { WsError::Http(response) => { let status = response.status(); @@ -696,7 +709,7 @@ async fn run_websocket_response_stream( let message = match response { Ok(Some(Ok(msg))) => msg, Ok(Some(Err(err))) => { - return Err(ApiError::Stream(err.to_string())); + return Err(map_ws_stream_error(err)); } Ok(None) => { return Err(ApiError::Stream( @@ -870,7 +883,7 @@ fn safety_buffering_for_event( } async fn send_websocket_request( - ws_stream: &WsStream, + ws_stream: &mut WsStream, request_text: String, idle_timeout: Duration, telemetry: Option<&Arc>, @@ -883,9 +896,7 @@ async fn send_websocket_request( ) .await .map_err(|_| ApiError::Stream("idle timeout sending websocket request".into())) - .and_then(|result| { - result.map_err(|err| ApiError::Stream(format!("failed to send websocket request: {err}"))) - }); + .and_then(|result| result.map_err(map_ws_stream_error)); if let Some(t) = telemetry.as_ref() { t.on_ws_request( diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index e0981b8c1d..fab92ed28a 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -3,8 +3,8 @@ use std::time::Duration; use crate::AuthProvider; use bytes::Bytes; use codex_http_client::HttpResponse; +use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; use codex_http_client::RouteAwareRequestError; use futures::Stream; use http::Method; @@ -162,7 +162,13 @@ pub async fn upload_openai_file( source, })?; let create_status = create_response.status(); - let create_body = create_response.text().await.unwrap_or_default(); + let create_body = create_response + .text() + .await + .map_err(|source| OpenAiFileError::Request { + url: create_url.clone(), + source, + })?; if !create_status.is_success() { return Err(OpenAiFileError::UnexpectedStatus { url: create_url, @@ -275,7 +281,14 @@ pub async fn upload_openai_file( source, })?; let finalize_status = finalize_response.status(); - let finalize_body = finalize_response.text().await.unwrap_or_default(); + let finalize_body = + finalize_response + .text() + .await + .map_err(|source| OpenAiFileError::Request { + url: finalize_url.clone(), + source, + })?; if !finalize_status.is_success() { return Err(OpenAiFileError::UnexpectedStatus { url: finalize_url.clone(), @@ -331,7 +344,7 @@ fn authorized_request( auth: &dyn AuthProvider, method: Method, url: &str, -) -> RouteAwareRequestBuilder { +) -> RequestBuilder { let mut headers = http::HeaderMap::new(); auth.add_auth_headers(&mut headers); diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index 3baf9346e9..9b55b00848 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -597,7 +597,13 @@ async fn process_sse_with_treatment( Ok(Some(Ok(sse))) => sse, Ok(Some(Err(e))) => { debug!("SSE Error: {e:#}"); - let _ = tx_event.send(Err(ApiError::Stream(e.to_string()))).await; + let error = match e { + eventsource_stream::EventStreamError::Transport( + error @ codex_client::TransportError::Policy(_), + ) => ApiError::Transport(error), + error => ApiError::Stream(error.to_string()), + }; + let _ = tx_event.send(Err(error)).await; return; } Ok(None) => { diff --git a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs index 5eabd5a4e2..bd6f463852 100644 --- a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs +++ b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs @@ -14,7 +14,7 @@ use codex_api::RealtimeTranscriptState; use codex_api::RealtimeWebsocketClient; use codex_api::RetryConfig; use codex_http_client::HttpClientFactory; -use codex_http_client::OutboundProxyPolicy; +use codex_http_client::OutboundProxyPolicy::ReqwestDefault; use codex_protocol::protocol::RealtimeHandoffRequested; use codex_protocol::protocol::RealtimeTranscriptDelta; use codex_protocol::protocol::RealtimeTranscriptDone; @@ -142,10 +142,7 @@ async fn realtime_ws_e2e_session_create_and_event_flow() { }) .await; - let client = RealtimeWebsocketClient::new( - test_provider(format!("http://{addr}")), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(test_provider(format!("http://{addr}"))); let connection = client .connect( RealtimeSessionConfig { @@ -251,11 +248,7 @@ async fn realtime_ws_connect_webrtc_sideband_retries_join_until_server_is_availa provider.retry.max_attempts = 1; provider.retry.base_delay = Duration::from_millis(100); - let client = RealtimeWebsocketClient::new( - provider, - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ) - .with_webrtc_sideband_base_url(format!("http://{addr}")); + let client = test_client(provider).with_webrtc_sideband_base_url(format!("http://{addr}")); let connection = client .connect_webrtc_sideband( RealtimeSessionConfig { @@ -330,10 +323,7 @@ async fn realtime_ws_e2e_send_while_next_event_waits() { }) .await; - let client = RealtimeWebsocketClient::new( - test_provider(format!("http://{addr}")), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(test_provider(format!("http://{addr}"))); let connection = client .connect( RealtimeSessionConfig { @@ -403,10 +393,7 @@ async fn realtime_ws_e2e_disconnected_emitted_once() { }) .await; - let client = RealtimeWebsocketClient::new( - test_provider(format!("http://{addr}")), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(test_provider(format!("http://{addr}"))); let connection = client .connect( RealtimeSessionConfig { @@ -472,10 +459,7 @@ async fn realtime_ws_e2e_ignores_unknown_text_events() { }) .await; - let client = RealtimeWebsocketClient::new( - test_provider(format!("http://{addr}")), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(test_provider(format!("http://{addr}"))); let connection = client .connect( RealtimeSessionConfig { @@ -584,10 +568,7 @@ async fn realtime_ws_e2e_realtime_v2_parser_emits_handoff_requested() { }) .await; - let client = RealtimeWebsocketClient::new( - test_provider(format!("http://{addr}")), - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ); + let client = test_client(test_provider(format!("http://{addr}"))); let connection = client .connect( RealtimeSessionConfig { @@ -665,3 +646,7 @@ async fn realtime_ws_e2e_realtime_v2_parser_emits_handoff_requested() { connection.close().await.expect("close"); server.await.expect("server task"); } + +fn test_client(provider: Provider) -> RealtimeWebsocketClient { + RealtimeWebsocketClient::new(provider, HttpClientFactory::new(ReqwestDefault)) +} diff --git a/codex-rs/codex-client/src/retry.rs b/codex-rs/codex-client/src/retry.rs index fd6bf007d8..071658272c 100644 --- a/codex-rs/codex-client/src/retry.rs +++ b/codex-rs/codex-client/src/retry.rs @@ -33,6 +33,7 @@ impl RetryOn { | TransportError::Network(_) => self.retry_transport, TransportError::Build(_) | TransportError::RetryLimit + | TransportError::Policy(_) | TransportError::ResponseTooLarge { .. } => false, } } diff --git a/codex-rs/core-plugins/src/http_client_selector.rs b/codex-rs/core-plugins/src/http_client_selector.rs index cd1515bde5..f42abea4cc 100644 --- a/codex-rs/core-plugins/src/http_client_selector.rs +++ b/codex-rs/core-plugins/src/http_client_selector.rs @@ -1,6 +1,6 @@ use codex_http_client::OutboundProxyPolicy; +use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; use http::Method; use std::fmt::Debug; @@ -9,12 +9,12 @@ use std::fmt::Debug; /// Implementations must keep route selection coupled to the request URL. Returning a transport /// client would let callers send a different URL than the one used for route selection. pub(crate) trait HttpClientSelector: Debug + Send + Sync { - fn request(&self, method: Method, url: &str) -> RouteAwareRequestBuilder; + fn request(&self, method: Method, url: &str) -> RequestBuilder; fn outbound_proxy_policy(&self) -> OutboundProxyPolicy; } impl HttpClientSelector for RouteAwareClientPool { - fn request(&self, method: Method, url: &str) -> RouteAwareRequestBuilder { + fn request(&self, method: Method, url: &str) -> RequestBuilder { RouteAwareClientPool::request(self, method, url) } diff --git a/codex-rs/core-plugins/src/remote.rs b/codex-rs/core-plugins/src/remote.rs index 79130e322a..15a639a1c4 100644 --- a/codex-rs/core-plugins/src/remote.rs +++ b/codex-rs/core-plugins/src/remote.rs @@ -22,8 +22,8 @@ use codex_app_server_protocol::ScheduledTaskSummary; use codex_app_server_protocol::SkillInterface; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; +use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; use codex_http_client::RouteAwareRequestError; use codex_login::CodexAuth; use codex_login::default_client::default_headers; @@ -175,7 +175,7 @@ impl RemotePluginServiceConfig { } } - pub(crate) fn http_request(&self, method: Method, url: &str) -> RouteAwareRequestBuilder { + pub(crate) fn http_request(&self, method: Method, url: &str) -> RequestBuilder { self.http_clients .request(method, url) .headers(default_headers()) @@ -2369,10 +2369,10 @@ fn ensure_chatgpt_auth(auth: Option<&CodexAuth>) -> Result<&CodexAuth, RemotePlu } fn authenticated_request( - request: RouteAwareRequestBuilder, + request: RequestBuilder, auth: &CodexAuth, product_sku: &str, -) -> RouteAwareRequestBuilder { +) -> RequestBuilder { request .timeout(REMOTE_PLUGIN_CATALOG_TIMEOUT) .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()) @@ -2380,7 +2380,7 @@ fn authenticated_request( } async fn send_and_decode Deserialize<'de>>( - request: RouteAwareRequestBuilder, + request: RequestBuilder, url: &str, ) -> Result { let response = request @@ -2391,7 +2391,13 @@ async fn send_and_decode Deserialize<'de>>( source, })?; let status = response.status(); - let body = response.text().await.unwrap_or_default(); + let body = response + .text() + .await + .map_err(|source| RemotePluginCatalogError::Request { + url: url.to_string(), + source, + })?; if !status.is_success() { return Err(RemotePluginCatalogError::UnexpectedStatus { url: url.to_string(), diff --git a/codex-rs/core-plugins/src/remote/share.rs b/codex-rs/core-plugins/src/remote/share.rs index 713792d11e..1e499aa9bb 100644 --- a/codex-rs/core-plugins/src/remote/share.rs +++ b/codex-rs/core-plugins/src/remote/share.rs @@ -1,7 +1,7 @@ use super::*; use crate::plugin_bundle_archive::PluginBundlePackError; use crate::plugin_bundle_archive::pack_plugin_bundle_tar_gz; -use codex_http_client::RouteAwareRequestBuilder; +use codex_http_client::RequestBuilder; use codex_login::CodexAuth; use codex_utils_absolute_path::AbsolutePathBuf; use http::Method; @@ -447,7 +447,13 @@ async fn put_workspace_plugin_upload( source, })?; let status = response.status(); - let body = response.text().await.unwrap_or_default(); + let body = response + .text() + .await + .map_err(|source| RemotePluginCatalogError::Request { + url: "workspace plugin upload URL".to_string(), + source, + })?; if ![StatusCode::OK, StatusCode::CREATED].contains(&status) { return Err(RemotePluginCatalogError::UnexpectedStatus { url: "workspace plugin upload URL".to_string(), @@ -513,7 +519,7 @@ fn archive_plugin_for_upload_with_limit( } async fn send_and_expect_status( - request: RouteAwareRequestBuilder, + request: RequestBuilder, url_for_error: &str, expected_statuses: &[StatusCode], ) -> Result<(), RemotePluginCatalogError> { @@ -525,7 +531,13 @@ async fn send_and_expect_status( source, })?; let status = response.status(); - let body = response.text().await.unwrap_or_default(); + let body = response + .text() + .await + .map_err(|source| RemotePluginCatalogError::Request { + url: url_for_error.to_string(), + source, + })?; if !expected_statuses.contains(&status) { return Err(RemotePluginCatalogError::UnexpectedStatus { url: url_for_error.to_string(), diff --git a/codex-rs/core-plugins/src/remote/share/tests.rs b/codex-rs/core-plugins/src/remote/share/tests.rs index 65c0e6adb1..5b04a9bfff 100644 --- a/codex-rs/core-plugins/src/remote/share/tests.rs +++ b/codex-rs/core-plugins/src/remote/share/tests.rs @@ -758,3 +758,69 @@ async fn delete_remote_plugin_share_deletes_workspace_plugin() { BTreeMap::new() ); } + +#[tokio::test] +async fn revoked_success_body_is_a_request_error() { + for method in [Method::PUT, Method::POST] { + use tokio::io::AsyncReadExt; + use tokio::io::AsyncWriteExt; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/upload", listener.local_addr().unwrap()); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish( + policy.revision(), + codex_http_client::DestinationPolicy::Unrestricted, + ); + let config = RemotePluginServiceConfig::new( + url.clone(), + crate::test_support::test_http_client_factory().with_network_policy(policy), + /*product_sku*/ None, + ); + let server = async { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut headers = Vec::new(); + while !headers.ends_with(b"\r\n\r\n") { + headers.push(socket.read_u8().await.unwrap()); + } + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 4\r\nConnection: close\r\n\r\na") + .await + .unwrap(); + // Let the client consume the successful headers and wait for the remaining body. + tokio::time::sleep(std::time::Duration::from_millis(/*millis*/ 50)).await; + controller.publish( + controller.policy().revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: Default::default(), + }, + ); + }; + let operation = async { + if method == Method::PUT { + put_workspace_plugin_upload(&config, &url, Vec::new()).await + } else { + send_and_expect_status(config.http_request(method, &url), &url, &[StatusCode::OK]) + .await + } + }; + let (result, ()) = + tokio::time::timeout(std::time::Duration::from_secs(/*secs*/ 5), async { + tokio::join!(operation, server) + }) + .await + .unwrap(); + assert!( + matches!( + result, + Err(RemotePluginCatalogError::Request { + source: codex_http_client::RouteAwareRequestError::Policy( + codex_http_client::NetworkPolicyDenied::Revoked + ), + .. + }) + ), + "{result:?}" + ); + } +} diff --git a/codex-rs/core-plugins/src/startup_sync.rs b/codex-rs/core-plugins/src/startup_sync.rs index dd7eff45d9..b366cf9e4c 100644 --- a/codex-rs/core-plugins/src/startup_sync.rs +++ b/codex-rs/core-plugins/src/startup_sync.rs @@ -7,8 +7,8 @@ use std::process::Stdio; use std::time::Duration; use self::http_client::StartupSyncHttpClient; -use self::http_client::StartupSyncRequestBuilder; use codex_http_client::HttpClientFactory; +use codex_http_client::RequestBuilder; use codex_login::default_client::default_headers; use codex_otel::CURATED_PLUGINS_STARTUP_SYNC_FINAL_METRIC; use codex_otel::CURATED_PLUGINS_STARTUP_SYNC_METRIC; @@ -928,7 +928,10 @@ async fn fetch_github_text( .await .map_err(|err| format!("failed to {context} from {url}: {err}"))?; let status = response.status(); - let body = response.text().await.unwrap_or_default(); + let body = response + .text() + .await + .map_err(|err| format!("failed to {context} from {url}: {err}"))?; if !status.is_success() { return Err(format!( "{context} from {url} failed with status {status}: {body}" @@ -971,7 +974,10 @@ async fn fetch_public_text( .await .map_err(|err| format!("failed to {context} from {url}: {err}"))?; let status = response.status(); - let body = response.text().await.unwrap_or_default(); + let body = response + .text() + .await + .map_err(|err| format!("failed to {context} from {url}: {err}"))?; if !status.is_success() { return Err(format!( "{context} from {url} failed with status {status}: {body}" @@ -1004,17 +1010,14 @@ async fn fetch_public_bytes( Ok(body.to_vec()) } -fn github_request(http_clients: &StartupSyncHttpClient, url: &str) -> StartupSyncRequestBuilder { +fn github_request(http_clients: &StartupSyncHttpClient, url: &str) -> RequestBuilder { startup_sync_request(http_clients, url) .timeout(CURATED_PLUGINS_HTTP_TIMEOUT) .header("accept", GITHUB_API_ACCEPT_HEADER) .header("x-github-api-version", GITHUB_API_VERSION_HEADER) } -fn startup_sync_request( - http_clients: &StartupSyncHttpClient, - url: &str, -) -> StartupSyncRequestBuilder { +fn startup_sync_request(http_clients: &StartupSyncHttpClient, url: &str) -> RequestBuilder { http_clients .request(Method::GET, url) .headers(default_headers()) diff --git a/codex-rs/core-plugins/src/startup_sync/http_client.rs b/codex-rs/core-plugins/src/startup_sync/http_client.rs index b84813e41e..3debe0cd1e 100644 --- a/codex-rs/core-plugins/src/startup_sync/http_client.rs +++ b/codex-rs/core-plugins/src/startup_sync/http_client.rs @@ -1,106 +1,37 @@ -//! Startup-sync-specific HTTP transport selection. +//! HTTP transport for curated plugin startup sync. //! -//! Curated plugin startup sync normally uses git, so its HTTP path is also a recovery path for -//! machines where git is unavailable or fails. Under `ReqwestDefault`, that recovery path must -//! preserve the legacy `codex_login::default_client::create_client_without_request_logging()` -//! behavior: invalid custom-CA configuration is logged and falls back to a normal client instead -//! of making HTTP sync fail as well. -//! -//! When `RespectSystemProxy` is enabled, however, every concrete request URL—including download -//! URLs returned by another endpoint—must be routed through `RouteAwareClientPool` so PAC and -//! operating-system proxy settings are respected. The route-aware pool intentionally surfaces -//! client-construction errors and therefore cannot provide the legacy custom-CA fallback for free. -//! -//! `StartupSyncHttpClient` keeps those two policies behind one request API without making lenient -//! custom-CA handling a global HTTP-client behavior. This module selects the transport only; -//! startup-sync request helpers remain responsible for applying the standard Codex headers. +//! Both proxy modes retain the supplied application network policy. The default proxy mode also +//! keeps the legacy fallback for invalid custom CA configuration; system-proxy mode fails closed. +//! Startup-sync request helpers apply the standard Codex headers. use std::sync::Arc; -use std::time::Duration; use crate::http_client_selector::HttpClientSelector; use codex_http_client::ClientRouteClass; -use codex_http_client::HttpClient; use codex_http_client::HttpClientFactory; -use codex_http_client::HttpResponse; -use codex_http_client::OutboundProxyPolicy; use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; -use codex_login::default_client::create_client_without_request_logging; -use http::HeaderMap; use http::Method; -pub(super) enum StartupSyncHttpClient { - Default(HttpClient), - RouteAware(Arc), -} +pub(super) struct StartupSyncHttpClient(Arc); impl StartupSyncHttpClient { pub(super) fn new(http_client_factory: &HttpClientFactory) -> Self { - match http_client_factory.outbound_proxy_policy() { - OutboundProxyPolicy::ReqwestDefault => { - Self::Default(create_client_without_request_logging()) - } - OutboundProxyPolicy::RespectSystemProxy => { - let http_clients = - RouteAwareClientPool::with_chatgpt_cloudflare_cookies_without_request_logging( - http_client_factory.clone(), - ClientRouteClass::Api, - ); - Self::RouteAware(Arc::new(http_clients)) - } - } + let http_clients = + RouteAwareClientPool::with_chatgpt_cloudflare_cookies_without_request_logging( + http_client_factory.clone(), + ClientRouteClass::Api, + ) + .with_legacy_custom_ca_fallback(); + Self(Arc::new(http_clients)) } #[cfg(test)] pub(super) fn route_aware(http_clients: Arc) -> Self { - Self::RouteAware(http_clients) + Self(http_clients) } - pub(super) fn request(&self, method: Method, url: &str) -> StartupSyncRequestBuilder { - match self { - Self::Default(client) => { - StartupSyncRequestBuilder::Default(client.request(method, url)) - } - Self::RouteAware(http_clients) => { - StartupSyncRequestBuilder::RouteAware(http_clients.request(method, url)) - } - } - } -} - -pub(super) enum StartupSyncRequestBuilder { - Default(RequestBuilder), - RouteAware(RouteAwareRequestBuilder), -} - -impl StartupSyncRequestBuilder { - pub(super) fn timeout(self, timeout: Duration) -> Self { - match self { - Self::Default(request) => Self::Default(request.timeout(timeout)), - Self::RouteAware(request) => Self::RouteAware(request.timeout(timeout)), - } - } - - pub(super) fn header(self, key: &'static str, value: &'static str) -> Self { - match self { - Self::Default(request) => Self::Default(request.header(key, value)), - Self::RouteAware(request) => Self::RouteAware(request.header(key, value)), - } - } - - pub(super) fn headers(self, headers: HeaderMap) -> Self { - match self { - Self::Default(request) => Self::Default(request.headers(headers)), - Self::RouteAware(request) => Self::RouteAware(request.headers(headers)), - } - } - - pub(super) async fn send(self) -> Result { - match self { - Self::Default(request) => request.send().await.map_err(|err| err.to_string()), - Self::RouteAware(request) => request.send().await.map_err(|err| err.to_string()), - } + pub(super) fn request(&self, method: Method, url: &str) -> RequestBuilder { + self.0.request(method, url) } } diff --git a/codex-rs/core-plugins/src/startup_sync_tests.rs b/codex-rs/core-plugins/src/startup_sync_tests.rs index de3e306de1..8695c7a439 100644 --- a/codex-rs/core-plugins/src/startup_sync_tests.rs +++ b/codex-rs/core-plugins/src/startup_sync_tests.rs @@ -2,6 +2,9 @@ use super::*; use crate::git_policy::REPOSITORY_LOCAL_GIT_ENVIRONMENT_VARIABLES; use crate::test_support::RecordingHttpClientSelector; use crate::test_support::recorded_http_client_urls; +use codex_http_client::DestinationPolicy::Restricted; +use codex_http_client::NetworkPolicyDenied::Destination; +use codex_http_client::RouteAwareRequestError::Policy; use pretty_assertions::assert_eq; use std::ffi::OsStr; use std::io::Write; @@ -185,6 +188,16 @@ async fn backup_archive_routes_metadata_and_backend_supplied_download_urls() { .expect("backup archive download should succeed"); assert_eq!(body, b"archive"); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + let allowed_hosts = Default::default(); + controller.publish(policy.revision(), Restricted { allowed_hosts }); + let factory = crate::test_support::test_http_client_factory().with_network_policy(policy); + let clients = StartupSyncHttpClient::new(&factory); + for url in [&metadata_url, &download_url] { + let result = clients.request(Method::GET, url).send().await; + assert!(matches!(result, Err(Policy(Destination)))); + } assert_eq!( recorded_http_client_urls(&selected_urls), vec![metadata_url, download_url] diff --git a/codex-rs/core-plugins/src/test_support.rs b/codex-rs/core-plugins/src/test_support.rs index 9bb63ae284..7366c19d8a 100644 --- a/codex-rs/core-plugins/src/test_support.rs +++ b/codex-rs/core-plugins/src/test_support.rs @@ -19,8 +19,8 @@ use codex_exec_server::LOCAL_FS; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; +use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; use codex_login::AuthHeaders; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -320,7 +320,7 @@ impl RecordingHttpClientSelector { } impl HttpClientSelector for RecordingHttpClientSelector { - fn request(&self, method: Method, url: &str) -> RouteAwareRequestBuilder { + fn request(&self, method: Method, url: &str) -> RequestBuilder { match self.selected_urls.lock() { Ok(mut selected_urls) => selected_urls.push(url.to_string()), Err(error) => panic!("selected URL recorder lock should not be poisoned: {error}"), diff --git a/codex-rs/core/src/realtime_conversation/sideband.rs b/codex-rs/core/src/realtime_conversation/sideband.rs index 08d2aeb5a3..d0b835550b 100644 --- a/codex-rs/core/src/realtime_conversation/sideband.rs +++ b/codex-rs/core/src/realtime_conversation/sideband.rs @@ -143,7 +143,12 @@ pub(super) fn spawn_webrtc_sideband_input_task( if !realtime_active.load(Ordering::Relaxed) || stop_token.is_cancelled() { break; } - if event_parser != RealtimeEventParser::FramelessBidi { + if event_parser != RealtimeEventParser::FramelessBidi + || matches!( + &err, + ApiError::Transport(codex_client::TransportError::Policy(_)) + ) + { report_realtime_transport_loss(&events_tx, err).await; break; } diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 0d2711b7a8..6fa005ba50 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -295,7 +295,7 @@ impl EnvironmentRegistryClient { let body = response .text() .await - .map_err(|error| ExecServerError::EnvironmentRegistryRequest(error.into()))?; + .map_err(ExecServerError::EnvironmentRegistryRequest)?; return serde_json::from_str(&body).map_err(ExecServerError::Json); } @@ -372,7 +372,7 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { let response = response .json::() .await - .map_err(|error| ExecServerError::EnvironmentRegistryRequest(error.into()))?; + .map_err(ExecServerError::EnvironmentRegistryRequest)?; if !response.valid { return Err(ExecServerError::Protocol( "environment registry rejected Noise relay harness key".to_string(), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs index 5841259415..1411c4c5d9 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs @@ -42,6 +42,7 @@ pub(super) fn sampler_failure_reason(error: &LunaSamplerError) -> &'static str { ApiError::Transport(TransportError::RetryLimit) => "retry_limit", ApiError::Transport(TransportError::Build(_)) => "request_build_error", ApiError::Transport(TransportError::ResponseTooLarge { .. }) => "response_too_large", + ApiError::Transport(TransportError::Policy(_)) => "network_policy_denied", ApiError::Stream(_) => "stream_error", ApiError::ContextWindowExceeded => "context_window_exceeded", ApiError::QuotaExceeded => "quota_exceeded", diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index 0265fb647f..2cd5cd7b94 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -85,7 +85,9 @@ impl SamplingExecution { | LunaSamplerError::InputTooLarge | LunaSamplerError::Api( ApiError::Transport( - TransportError::Build(_) | TransportError::ResponseTooLarge { .. }, + TransportError::Build(_) + | TransportError::ResponseTooLarge { .. } + | TransportError::Policy(_), ) | ApiError::ContextWindowExceeded | ApiError::QuotaExceeded diff --git a/codex-rs/feedback/src/upload.rs b/codex-rs/feedback/src/upload.rs index 0c9722b181..5162e4cd65 100644 --- a/codex-rs/feedback/src/upload.rs +++ b/codex-rs/feedback/src/upload.rs @@ -7,8 +7,8 @@ use std::time::SystemTime; use anyhow::Context; use anyhow::Result; use bytes::Bytes; +use codex_http_client::RequestBuilder; use codex_http_client::RouteAwareClientPool; -use codex_http_client::RouteAwareRequestBuilder; use flate2::Compression; use flate2::write::GzEncoder; use http::HeaderMap; @@ -34,7 +34,7 @@ pub(super) fn envelope_request( dsn: &Dsn, body: Bytes, timeout: Duration, -) -> RouteAwareRequestBuilder { +) -> RequestBuilder { let sentry_options = ClientOptions::default(); let sentry_auth = dsn.to_auth(Some(sentry_options.user_agent.as_ref())); let request = client_pool diff --git a/codex-rs/http-client/Cargo.toml b/codex-rs/http-client/Cargo.toml index a390c46ab2..a18d653d96 100644 --- a/codex-rs/http-client/Cargo.toml +++ b/codex-rs/http-client/Cargo.toml @@ -17,6 +17,7 @@ rustls-native-certs = { workspace = true } rustls-pki-types = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +serde_urlencoded = "0.7" sha2 = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "time", "sync"] } diff --git a/codex-rs/http-client/src/client.rs b/codex-rs/http-client/src/client.rs index b5b1ba9145..c5c2e40bba 100644 --- a/codex-rs/http-client/src/client.rs +++ b/codex-rs/http-client/src/client.rs @@ -1,6 +1,5 @@ -//! Reusable HTTP client wrappers and request construction shared by fixed and route-aware transports. +//! Selects direct or routed HTTP execution and owns transport diagnostics. -use http::Error as HttpRequestBuildError; use http::HeaderMap; use http::HeaderName; use http::HeaderValue; @@ -8,90 +7,138 @@ use opentelemetry::global; use opentelemetry::propagation::Injector; use reqwest::IntoUrl; use reqwest::Method; -use serde::Serialize; -use std::fmt::Display; -use std::time::Duration; +use std::sync::Arc; use tracing::Span; use tracing_opentelemetry::OpenTelemetrySpanExt; -pub type HttpError = reqwest::Error; -pub type HttpResponse = reqwest::Response; +use crate::RequestBuilder; +use crate::RouteAwareClientPool; -/// Reusable HTTP client wrapper with shared tracing and request-diagnostic behavior. +pub type HttpError = crate::RouteAwareRequestError; + +/// Reusable HTTP client with shared tracing and request diagnostics. /// /// Product callers should obtain this through [`crate::HttpClientFactory`] for a fixed /// destination or use [`crate::RouteAwareClientPool`] when request and redirect URLs can vary. #[derive(Clone, Debug)] pub struct HttpClient { - inner: reqwest::Client, - request_logging: RequestLogging, + pub(crate) backend: HttpClientBackend, +} + +#[derive(Clone, Debug)] +pub(crate) enum HttpClientBackend { + Direct(TransportClient), + Routed(Arc), +} + +/// A resolved transport used by direct clients and the route cache. +#[derive(Clone)] +pub(crate) struct TransportClient { + pub(crate) inner: reqwest::Client, + pub(crate) request_logging: RequestLogging, + default_headers: Arc, +} + +impl std::fmt::Debug for TransportClient { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("TransportClient") + .field("request_logging", &self.request_logging) + .finish_non_exhaustive() + } } impl HttpClient { pub fn new(inner: reqwest::Client) -> Self { - Self::from_parts(inner, RequestLogging::Enabled) + Self::from_parts(inner, RequestLogging::Enabled, HeaderMap::new()) } - /// Creates a client that suppresses request URL and response-header diagnostics. - /// - /// Use this for endpoints whose URLs or headers may contain credentials that are redacted by - /// the caller above the HTTP transport boundary. + /// Suppresses diagnostics for endpoints whose URLs or headers may contain credentials. pub fn new_without_request_logging(inner: reqwest::Client) -> Self { - Self::from_parts(inner, RequestLogging::Disabled) + Self::from_parts(inner, RequestLogging::Disabled, HeaderMap::new()) } /// Suppresses URL and response-header diagnostics while preserving this client's routing. pub fn without_request_logging(mut self) -> Self { - self.request_logging = RequestLogging::Disabled; + match &mut self.backend { + HttpClientBackend::Direct(client) => client.request_logging = RequestLogging::Disabled, + HttpClientBackend::Routed(pool) => { + Arc::make_mut(pool).client_builder.request_logging = RequestLogging::Disabled; + } + } self } - pub(crate) fn from_parts(inner: reqwest::Client, request_logging: RequestLogging) -> Self { + pub(crate) fn from_parts( + inner: reqwest::Client, + request_logging: RequestLogging, + default_headers: HeaderMap, + ) -> Self { Self { - inner, - request_logging, + backend: HttpClientBackend::Direct(TransportClient::new( + inner, + request_logging, + default_headers, + )), } } - pub fn get(&self, url: U) -> RequestBuilder - where - U: IntoUrl, - { + pub fn get(&self, url: U) -> RequestBuilder { self.request(Method::GET, url) } - pub fn head(&self, url: U) -> RequestBuilder - where - U: IntoUrl, - { + pub fn head(&self, url: U) -> RequestBuilder { self.request(Method::HEAD, url) } - pub fn post(&self, url: U) -> RequestBuilder - where - U: IntoUrl, - { + pub fn post(&self, url: U) -> RequestBuilder { self.request(Method::POST, url) } - pub fn delete(&self, url: U) -> RequestBuilder - where - U: IntoUrl, - { + pub fn delete(&self, url: U) -> RequestBuilder { self.request(Method::DELETE, url) } - pub fn request(&self, method: Method, url: U) -> RequestBuilder - where - U: IntoUrl, - { - let url_str = url.as_str().to_string(); - RequestBuilder::new( - self.inner.request(method.clone(), url), - method, - url_str, - self.request_logging, - ) + pub fn request(&self, method: Method, url: U) -> RequestBuilder { + match &self.backend { + HttpClientBackend::Direct(client) => RequestBuilder::direct(client, method, url), + HttpClientBackend::Routed(pool) => { + RequestBuilder::routed(Arc::clone(pool), method, url) + } + } + } + + pub(crate) fn request_logging_enabled(&self) -> bool { + match &self.backend { + HttpClientBackend::Direct(client) => client.request_logging == RequestLogging::Enabled, + HttpClientBackend::Routed(pool) => { + pool.client_builder.request_logging == RequestLogging::Enabled + } + } + } +} + +pub(crate) fn apply_default_headers(headers: &mut HeaderMap, defaults: &HeaderMap) { + for name in defaults.keys() { + if !headers.contains_key(name) { + for value in defaults.get_all(name) { + headers.append(name.clone(), value.clone()); + } + } + } +} + +impl TransportClient { + pub(crate) fn new( + inner: reqwest::Client, + request_logging: RequestLogging, + default_headers: HeaderMap, + ) -> Self { + Self { + inner, + request_logging, + default_headers: Arc::new(default_headers), + } } pub(crate) async fn execute( @@ -117,6 +164,7 @@ impl HttpClient { &self, mut request: reqwest::Request, ) -> Result { + apply_default_headers(request.headers_mut(), &self.default_headers); request.headers_mut().extend(trace_headers()); self.inner.execute(request).await } @@ -157,10 +205,6 @@ impl HttpClient { ); } } - - pub(crate) const fn request_logging_enabled(&self) -> bool { - matches!(self.request_logging, RequestLogging::Enabled) - } } #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -170,124 +214,6 @@ pub(crate) enum RequestLogging { Disabled, } -#[must_use = "requests are not sent unless `send` is awaited"] -#[derive(Debug)] -pub struct RequestBuilder { - builder: reqwest::RequestBuilder, - method: Method, - url: String, - request_logging: RequestLogging, -} - -impl RequestBuilder { - fn new( - builder: reqwest::RequestBuilder, - method: Method, - url: String, - request_logging: RequestLogging, - ) -> Self { - Self { - builder, - method, - url, - request_logging, - } - } - - fn map(self, f: impl FnOnce(reqwest::RequestBuilder) -> reqwest::RequestBuilder) -> Self { - Self { - builder: f(self.builder), - method: self.method, - url: self.url, - request_logging: self.request_logging, - } - } - - pub fn headers(self, headers: HeaderMap) -> Self { - self.map(|builder| builder.headers(headers)) - } - - pub fn header(self, key: K, value: V) -> Self - where - HeaderName: TryFrom, - >::Error: Into, - HeaderValue: TryFrom, - >::Error: Into, - { - self.map(|builder| builder.header(key, value)) - } - - pub fn bearer_auth(self, token: T) -> Self - where - T: Display, - { - self.map(|builder| builder.bearer_auth(token)) - } - - pub fn timeout(self, timeout: Duration) -> Self { - self.map(|builder| builder.timeout(timeout)) - } - - pub fn json(self, value: &T) -> Self - where - T: ?Sized + Serialize, - { - self.map(|builder| builder.json(value)) - } - - pub fn query(self, query: &T) -> Self - where - T: ?Sized + Serialize, - { - self.map(|builder| builder.query(query)) - } - - pub fn body(self, body: B) -> Self - where - B: Into, - { - self.map(|builder| builder.body(body)) - } - - pub(crate) fn build(self) -> Result { - self.builder.build() - } - - pub async fn send(self) -> Result { - let headers = trace_headers(); - - match self.builder.headers(headers).send().await { - Ok(response) => { - if self.request_logging == RequestLogging::Enabled { - tracing::debug!( - method = %self.method, - url = %self.url, - status = %response.status(), - headers = ?response.headers(), - version = ?response.version(), - "Request completed" - ); - } - - Ok(response) - } - Err(error) => { - if self.request_logging == RequestLogging::Enabled { - let status = error.status(); - tracing::debug!( - method = %self.method, - url = %self.url, - status = status.map(|s| s.as_u16()), - error = %error, - "Request failed" - ); - } - Err(error) - } - } - } -} - struct HeaderMapInjector<'a>(&'a mut HeaderMap); impl<'a> Injector for HeaderMapInjector<'a> { diff --git a/codex-rs/http-client/src/client_builder.rs b/codex-rs/http-client/src/client_builder.rs index 8b7ae3de5c..6ffd914f99 100644 --- a/codex-rs/http-client/src/client_builder.rs +++ b/codex-rs/http-client/src/client_builder.rs @@ -20,7 +20,9 @@ use crate::HttpClient; use crate::HttpClientFactory; use crate::OutboundProxyRoute; use crate::chatgpt_cloudflare_cookies::ChatGptCookieStore; +use crate::client::HttpClientBackend; use crate::client::RequestLogging; +use crate::client::TransportClient; use crate::custom_ca::build_reqwest_client_with_custom_ca; use crate::with_chatgpt_cloudflare_cookie_store; @@ -31,13 +33,13 @@ use crate::with_chatgpt_cloudflare_cookie_store; /// bypass the factory and are restricted to documented exceptional or legacy compatibility paths. #[derive(Clone)] pub struct HttpClientBuilder { - default_headers: Option, + pub(crate) default_headers: Option, follow_redirects: bool, - redirect_observed: Option>, + pub(crate) redirect_observed: Option>, connect_timeout: Option, chatgpt_cloudflare_cookie_store: bool, chatgpt_cookie_store: Option>, - request_logging: RequestLogging, + pub(crate) request_logging: RequestLogging, tls_backend: TlsBackend, } @@ -84,6 +86,7 @@ impl HttpClientBuilder { Self::default() } + /// Applies every configured value unless the request sets that header explicitly. pub fn default_headers(mut self, headers: HeaderMap) -> Self { self.default_headers = Some(headers); self @@ -149,10 +152,22 @@ impl HttpClientBuilder { request_url: &str, route_class: ClientRouteClass, ) -> Result { + if http_client_factory.network_policy().is_managed() { + return Ok(crate::RouteAwareClientPool::with_builder( + http_client_factory.clone(), + route_class, + self, + ) + .into_client()); + } self.chatgpt_cookie_store = http_client_factory.chatgpt_cookie_store(); - let (builder, request_logging) = self.into_reqwest_parts(); + let (builder, request_logging, default_headers) = self.into_reqwest_parts(); let inner = http_client_factory.build_reqwest_client(builder, request_url, route_class)?; - Ok(HttpClient::from_parts(inner, request_logging)) + Ok(HttpClient::from_parts( + inner, + request_logging, + default_headers, + )) } /// Builds a client for a route that was already resolved by a route-aware caller. @@ -161,15 +176,19 @@ impl HttpClientBuilder { http_client_factory: &HttpClientFactory, route_class: ClientRouteClass, route: &OutboundProxyRoute, - ) -> Result { + ) -> Result { self.chatgpt_cookie_store = http_client_factory.chatgpt_cookie_store(); - let (builder, request_logging) = self.into_reqwest_parts(); + let (builder, request_logging, default_headers) = self.into_reqwest_parts(); let inner = http_client_factory.build_reqwest_client_for_resolved_route( builder, route_class, route, )?; - Ok(HttpClient::from_parts(inner, request_logging)) + Ok(TransportClient::new( + inner, + request_logging, + default_headers, + )) } /// Builds a client using the transport's default proxy behavior. @@ -212,7 +231,11 @@ impl HttpClientBuilder { note = "legacy custom-CA fallback only; use HttpClientFactory::build_client or build_respecting_outbound_proxy_policy" )] pub fn build_with_transport_default_proxy_and_custom_ca_fallback(self) -> HttpClient { - self.build_with_custom_ca_fallback(ProxyRouting::TransportDefault) + HttpClient { + backend: HttpClientBackend::Direct( + self.build_with_custom_ca_fallback(ProxyRouting::TransportDefault), + ), + } } /// Builds a direct client while preserving the legacy custom-CA fallback. @@ -225,33 +248,42 @@ impl HttpClientBuilder { note = "legacy custom-CA fallback only; use build_direct and propagate construction errors" )] pub fn build_direct_with_custom_ca_fallback(self) -> HttpClient { - self.build_with_custom_ca_fallback(ProxyRouting::Direct) + HttpClient { + backend: HttpClientBackend::Direct( + self.build_with_custom_ca_fallback(ProxyRouting::Direct), + ), + } } fn build_with_proxy_routing( - self, + mut self, proxy_routing: ProxyRouting, ) -> Result { let request_logging = self.request_logging; + let default_headers = self.default_headers.take().unwrap_or_default(); build_reqwest_client_with_custom_ca(self.reqwest_builder(proxy_routing)) - .map(|inner| HttpClient::from_parts(inner, request_logging)) + .map(|inner| HttpClient::from_parts(inner, request_logging, default_headers)) } - fn build_with_custom_ca_fallback(self, proxy_routing: ProxyRouting) -> HttpClient { + pub(crate) fn build_with_custom_ca_fallback( + self, + proxy_routing: ProxyRouting, + ) -> TransportClient { self.build_with_custom_ca_fallback_using(proxy_routing, build_reqwest_client_with_custom_ca) } fn build_with_custom_ca_fallback_using( - self, + mut self, proxy_routing: ProxyRouting, build_with_custom_ca: impl FnOnce( reqwest::ClientBuilder, ) -> Result, - ) -> HttpClient { + ) -> TransportClient { let request_logging = self.request_logging; - match build_with_custom_ca(self.clone().reqwest_builder(proxy_routing)) { - Ok(inner) => HttpClient::from_parts(inner, request_logging), + let default_headers = self.default_headers.take().unwrap_or_default(); + let inner = match build_with_custom_ca(self.clone().reqwest_builder(proxy_routing)) { + Ok(inner) => inner, Err(error) => { tracing::event!( target: "codex_otel.log_only", @@ -262,21 +294,26 @@ impl HttpClientBuilder { tracing::warn!(error = %error, "failed to build HTTP client with custom CA"); self.reqwest_builder(proxy_routing) .build() - .map(|inner| HttpClient::from_parts(inner, request_logging)) .unwrap_or_else(|fallback_error| { tracing::warn!( error = %fallback_error, "failed to build fallback HTTP client" ); - HttpClient::from_parts(reqwest::Client::new(), request_logging) + reqwest::Client::new() }) } - } + }; + TransportClient::new(inner, request_logging, default_headers) } - fn into_reqwest_parts(self) -> (reqwest::ClientBuilder, RequestLogging) { + fn into_reqwest_parts(mut self) -> (reqwest::ClientBuilder, RequestLogging, HeaderMap) { let request_logging = self.request_logging; - (self.base_reqwest_builder(), request_logging) + let default_headers = self.default_headers.take().unwrap_or_default(); + ( + self.base_reqwest_builder(), + request_logging, + default_headers, + ) } fn reqwest_builder(self, proxy_routing: ProxyRouting) -> reqwest::ClientBuilder { @@ -293,9 +330,6 @@ impl HttpClientBuilder { ensure_rustls_crypto_provider(); builder = builder.use_rustls_tls(); } - if let Some(default_headers) = self.default_headers { - builder = builder.default_headers(default_headers); - } if !self.follow_redirects { builder = builder.redirect(reqwest::redirect::Policy::none()); } else if let Some(redirect_observed) = self.redirect_observed { @@ -333,7 +367,7 @@ impl Default for HttpClientBuilder { } #[derive(Clone, Copy)] -enum ProxyRouting { +pub(crate) enum ProxyRouting { TransportDefault, Direct, } diff --git a/codex-rs/http-client/src/client_builder_tests.rs b/codex-rs/http-client/src/client_builder_tests.rs index 0b4cbb4b7d..99c0662ea9 100644 --- a/codex-rs/http-client/src/client_builder_tests.rs +++ b/codex-rs/http-client/src/client_builder_tests.rs @@ -37,6 +37,10 @@ async fn custom_ca_fallback_preserves_builder_configuration() { }) }); + let client = HttpClient { + backend: HttpClientBackend::Direct(client), + }; + let response = client .get(format!("http://{address}/fallback")) .send() diff --git a/codex-rs/http-client/src/error.rs b/codex-rs/http-client/src/error.rs index 43ace4f7a7..6105b2eea8 100644 --- a/codex-rs/http-client/src/error.rs +++ b/codex-rs/http-client/src/error.rs @@ -7,6 +7,8 @@ use thiserror::Error; #[derive(Debug, Error)] pub enum TransportError { + #[error(transparent)] + Policy(#[from] crate::NetworkPolicyDenied), #[error("http {status}: {body:?}")] Http { status: StatusCode, diff --git a/codex-rs/http-client/src/lib.rs b/codex-rs/http-client/src/lib.rs index e2f81e823e..9d72958620 100644 --- a/codex-rs/http-client/src/lib.rs +++ b/codex-rs/http-client/src/lib.rs @@ -7,7 +7,9 @@ mod error; mod network_policy; mod outbound_proxy; mod request; +mod request_builder; mod request_draft; +mod response; mod route_aware_client_pool; mod route_aware_redirect; mod tls_backend_fallback; @@ -17,8 +19,6 @@ pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store; pub use crate::chatgpt_hosts::is_allowed_chatgpt_host; pub use crate::client::HttpClient; pub use crate::client::HttpError; -pub use crate::client::HttpResponse; -pub use crate::client::RequestBuilder; pub use crate::client_builder::HttpClientBuilder; pub use crate::custom_ca::BuildCustomCaTransportError; /// Test-only subprocess hook for custom CA coverage. @@ -57,9 +57,10 @@ pub use crate::request::Request; pub use crate::request::RequestBody; pub use crate::request::RequestCompression; pub use crate::request::Response; +pub use crate::request_builder::RequestBuilder; +pub use crate::response::HttpResponse; pub use crate::route_aware_client_pool::RouteAwareClientPool; pub use crate::route_aware_client_pool::RouteAwareClientPoolError; -pub use crate::route_aware_client_pool::RouteAwareRequestBuilder; pub use crate::route_aware_client_pool::RouteAwareRequestError; pub use crate::transport::ByteStream; pub use crate::transport::HttpTransport; diff --git a/codex-rs/http-client/src/request_builder.rs b/codex-rs/http-client/src/request_builder.rs new file mode 100644 index 0000000000..4f9e561ba3 --- /dev/null +++ b/codex-rs/http-client/src/request_builder.rs @@ -0,0 +1,193 @@ +//! Builds direct and routed requests without constructing a second transport client. +//! Routed requests select their transport only when sent; debug output redacts their URLs. + +use bytes::Bytes; +use futures::TryStream; +use http::HeaderMap; +use http::HeaderName; +use http::HeaderValue; +use http::Method; +use http::header::AUTHORIZATION; +use http::header::CONTENT_TYPE; +use reqwest::IntoUrl; +use serde::Serialize; +use std::fmt; +use std::fmt::Display; +use std::sync::Arc; +use std::time::Duration; + +use crate::HttpError; +use crate::HttpResponse; +use crate::RouteAwareClientPool; +use crate::client::HttpClientBackend; +use crate::client::TransportClient; +use crate::request_draft::HeaderUpdate; +use crate::request_draft::RequestDraft; + +#[must_use = "requests are not sent unless `send` is awaited"] +pub struct RequestBuilder { + backend: HttpClientBackend, + request: Result, +} + +impl fmt::Debug for RequestBuilder { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RequestBuilder") + .field( + "method", + &self + .request + .as_ref() + .ok() + .map(|request| request.request.method()), + ) + .field("url", &"") + .finish_non_exhaustive() + } +} + +impl RequestBuilder { + pub(crate) fn direct(client: &TransportClient, method: Method, url: U) -> Self { + Self { + backend: HttpClientBackend::Direct(client.clone()), + request: RequestDraft::new(method, url).map_err(HttpError::Request), + } + } + + pub(crate) fn routed( + pool: Arc, + method: Method, + url: U, + ) -> Self { + Self { + backend: HttpClientBackend::Routed(pool), + request: RequestDraft::new(method, url).map_err(HttpError::Request), + } + } + + fn map(mut self, update: impl FnOnce(&mut RequestDraft) -> Result<(), HttpError>) -> Self { + self.request = self.request.and_then(|mut request| { + update(&mut request)?; + Ok(request) + }); + self + } + + pub fn headers(self, headers: HeaderMap) -> Self { + self.map(|request| { + request.extend_headers(headers); + Ok(()) + }) + } + + pub fn header(self, key: K, value: V) -> Self + where + HeaderName: TryFrom, + >::Error: Into, + HeaderValue: TryFrom, + >::Error: Into, + { + self.map(|request| { + let key = HeaderName::try_from(key) + .map_err(Into::into) + .map_err(|error: http::Error| HttpError::Build(error.to_string()))?; + let value = HeaderValue::try_from(value) + .map_err(Into::into) + .map_err(|error: http::Error| HttpError::Build(error.to_string()))?; + request.headers.push(HeaderUpdate::Append(key, value)); + Ok(()) + }) + } + + pub fn bearer_auth(self, token: T) -> Self { + self.map(|request| { + let mut value = HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(|error| HttpError::Build(error.to_string()))?; + value.set_sensitive(true); + request + .headers + .push(HeaderUpdate::Append(AUTHORIZATION, value)); + Ok(()) + }) + } + + /// Bounds route selection, connection establishment, and the response headers. + /// Use [`crate::HttpClientBuilder::connect_timeout`] to bound only connection establishment. + pub fn timeout(self, timeout: Duration) -> Self { + self.map(|request| { + *request.request.timeout_mut() = Some(timeout); + Ok(()) + }) + } + + pub fn version(self, version: http::Version) -> Self { + self.map(|request| { + *request.request.version_mut() = version; + Ok(()) + }) + } + + pub fn json(self, value: &T) -> Self { + self.map(|request| { + let body = + serde_json::to_vec(value).map_err(|error| HttpError::Build(error.to_string()))?; + let has_content_type = request.headers.iter().any(|update| match update { + HeaderUpdate::Replace(headers) => headers.contains_key(CONTENT_TYPE), + HeaderUpdate::Append(name, _) => *name == CONTENT_TYPE, + }); + if !has_content_type { + request.headers.push(HeaderUpdate::Append( + CONTENT_TYPE, + HeaderValue::from_static("application/json"), + )); + } + *request.request.body_mut() = Some(body.into()); + Ok(()) + }) + } + + pub fn query(self, query: &T) -> Self { + self.map(|request| { + query + .serialize(serde_urlencoded::Serializer::new( + &mut request.request.url_mut().query_pairs_mut(), + )) + .map_err(|error| HttpError::Build(error.to_string()))?; + if request.request.url().query() == Some("") { + request.request.url_mut().set_query(None); + } + Ok(()) + }) + } + + pub fn body>(self, body: B) -> Self { + self.map(|request| { + *request.request.body_mut() = Some(body.into()); + Ok(()) + }) + } + + /// Sets a streaming body without exposing the underlying HTTP implementation. + pub fn body_stream(self, stream: S) -> Self + where + S: TryStream + Send + 'static, + S::Error: Into>, + Bytes: From, + { + self.body(reqwest::Body::wrap_stream(stream)) + } + + pub async fn send(self) -> Result { + match self.backend { + HttpClientBackend::Routed(pool) => pool.send(self.request?).await, + HttpClientBackend::Direct(client) => { + Ok(client.execute(self.request?.build(&client)?).await?.into()) + } + } + } +} + +#[cfg(test)] +#[path = "request_builder_tests.rs"] +mod tests; diff --git a/codex-rs/http-client/src/request_builder_tests.rs b/codex-rs/http-client/src/request_builder_tests.rs new file mode 100644 index 0000000000..9431729c79 --- /dev/null +++ b/codex-rs/http-client/src/request_builder_tests.rs @@ -0,0 +1,117 @@ +use super::*; +use crate::ClientRouteClass; +use crate::HttpClientBuilder; +use crate::HttpClientFactory; +use crate::OutboundProxyPolicy; +use pretty_assertions::assert_eq; + +#[test] +fn request_draft_preserves_reqwest_url_auth_and_header_precedence() { + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy), + ClientRouteClass::Api, + ); + let client = HttpClientBuilder::new() + .build_direct() + .expect("client should build"); + let HttpClientBackend::Direct(transport) = &client.backend else { + panic!("client should use direct transport"); + }; + let reqwest = &transport.inner; + for url in [ + "https://example.com/path", + "https://user:p%40ss@example.com/path", + "https://:%E2%98%83@example.com/path", + "https://%FF:password@example.com/path", + ] { + let headers = HeaderMap::from_iter([ + ( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer replacement"), + ), + ( + CONTENT_TYPE, + HeaderValue::from_static("application/custom+json"), + ), + ]); + for (actual, expected) in [ + (pool.post(url), reqwest.post(url)), + ( + pool.post(url) + .header(http::header::AUTHORIZATION, "Bearer appended"), + reqwest + .post(url) + .header(http::header::AUTHORIZATION, "Bearer appended"), + ), + ( + pool.post(url) + .header(http::header::AUTHORIZATION, "Bearer first") + .headers(headers.clone()) + .header(http::header::AUTHORIZATION, "Bearer last"), + reqwest + .post(url) + .header(http::header::AUTHORIZATION, "Bearer first") + .headers(headers) + .header(http::header::AUTHORIZATION, "Bearer last"), + ), + ( + pool.post(url) + .header(CONTENT_TYPE, "application/custom+json") + .headers(HeaderMap::new()), + reqwest + .post(url) + .header(CONTENT_TYPE, "application/custom+json") + .headers(HeaderMap::new()), + ), + ( + pool.post(url).json(&serde_json::json!({"previous": true})), + reqwest + .post(url) + .json(&serde_json::json!({"previous": true})), + ), + ] { + let body = serde_json::json!({"query": "test"}); + let timeout = Duration::from_secs(7); + let actual = actual + .json(&body) + .timeout(timeout) + .request + .expect("request inputs should be valid") + .build(transport) + .expect("request should build"); + let expected = expected + .json(&body) + .timeout(timeout) + .build() + .expect("reqwest should build"); + assert_eq!( + ( + actual.url(), + actual.method(), + actual.headers(), + actual.body().and_then(reqwest::Body::as_bytes), + actual.timeout() + ), + ( + expected.url(), + expected.method(), + expected.headers(), + expected.body().and_then(reqwest::Body::as_bytes), + expected.timeout() + ), + ); + assert_eq!( + actual + .headers() + .values() + .map(HeaderValue::is_sensitive) + .collect::>(), + expected + .headers() + .values() + .map(HeaderValue::is_sensitive) + .collect::>(), + ); + } + } +} diff --git a/codex-rs/http-client/src/request_draft.rs b/codex-rs/http-client/src/request_draft.rs index 3d62437b1c..7b7b5fcbbe 100644 --- a/codex-rs/http-client/src/request_draft.rs +++ b/codex-rs/http-client/src/request_draft.rs @@ -10,7 +10,7 @@ use http::HeaderValue; use http::Method; use reqwest::IntoUrl; -use crate::HttpClient; +use crate::client::TransportClient; /// Request inputs awaiting construction through reqwest with a selected HTTP client. /// @@ -44,9 +44,14 @@ impl RequestDraft { self.headers.push(HeaderUpdate::Replace(headers)); } - pub(crate) fn build(self, client: &HttpClient) -> Result { + pub(crate) fn build( + self, + client: &TransportClient, + ) -> Result { let mut request = self.request; - let mut builder = client.request(request.method().clone(), request.url().clone()); + let mut builder = client + .inner + .request(request.method().clone(), request.url().clone()); for update in self.headers { builder = match update { HeaderUpdate::Replace(headers) => builder.headers(headers), diff --git a/codex-rs/http-client/src/response.rs b/codex-rs/http-client/src/response.rs new file mode 100644 index 0000000000..ef45a3102b --- /dev/null +++ b/codex-rs/http-client/src/response.rs @@ -0,0 +1,90 @@ +//! HTTP responses retain application authorization while their bodies are consumed. + +use std::ops::Deref; + +use bytes::Bytes; +use futures::Stream; +use futures::StreamExt; +use futures::stream; +use serde::de::DeserializeOwned; + +use crate::HttpError; +use crate::NetworkPermit; +use crate::NetworkPolicy; + +#[derive(Debug)] +pub struct HttpResponse { + inner: reqwest::Response, + permit: NetworkPermit, +} + +impl HttpResponse { + pub(crate) fn new(inner: reqwest::Response, permit: NetworkPermit) -> Self { + Self { inner, permit } + } + + pub async fn bytes(self) -> Result { + Ok(Box::pin(self.permit.run(self.inner.bytes())).await??) + } + + pub async fn text(self) -> Result { + Ok(Box::pin(self.permit.run(self.inner.text())).await??) + } + + pub async fn json(self) -> Result { + Ok(Box::pin(self.permit.run(self.inner.json())).await??) + } + + pub async fn chunk(&mut self) -> Result, HttpError> { + Ok(Box::pin(self.permit.run(self.inner.chunk())).await??) + } + + pub fn error_for_status(self) -> Result { + self.permit.check()?; + Ok(Self { + inner: self.inner.error_for_status()?, + permit: self.permit, + }) + } + + pub fn error_for_status_ref(&self) -> Result<&Self, HttpError> { + self.permit.check()?; + self.inner.error_for_status_ref()?; + Ok(self) + } + + pub fn bytes_stream(self) -> impl Stream> + Send + Unpin { + Box::pin(stream::try_unfold( + (Box::pin(self.inner.bytes_stream()), self.permit), + |(mut body, permit)| async move { + let next = permit.run(body.next()).await?; + next.transpose() + .map(|next| next.map(|bytes| (bytes, (body, permit)))) + .map_err(HttpError::from) + }, + )) + } +} + +impl Deref for HttpResponse { + type Target = reqwest::Response; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl From for HttpResponse { + fn from(inner: reqwest::Response) -> Self { + let permit = NetworkPolicy::unmanaged() + .acquire(inner.url()) + .unwrap_or_else(|_| unreachable!("unmanaged policy permits every URL")); + Self { inner, permit } + } +} + +impl> From> for HttpResponse { + fn from(response: http::Response) -> Self { + reqwest::Response::from(response).into() + } +} diff --git a/codex-rs/http-client/src/route_aware_client_pool.rs b/codex-rs/http-client/src/route_aware_client_pool.rs index 9fdeada936..8a2fec5339 100644 --- a/codex-rs/http-client/src/route_aware_client_pool.rs +++ b/codex-rs/http-client/src/route_aware_client_pool.rs @@ -8,27 +8,24 @@ use std::sync::Arc; use std::sync::Mutex; use std::time::Duration; -use bytes::Bytes; -use futures::TryStream; use http::HeaderMap; -use http::HeaderName; -use http::HeaderValue; use http::Method; use http::StatusCode; -use http::header::CONTENT_TYPE; use reqwest::IntoUrl; -use serde::Serialize; use crate::BuildRouteAwareHttpClientError; use crate::ClientRouteClass; use crate::HttpClient; use crate::HttpClientBuilder; use crate::HttpClientFactory; +use crate::NetworkPolicyDenied; use crate::OutboundProxyPolicy; use crate::OutboundProxyRoute; +use crate::RequestBuilder; use crate::RouteFailureClass; -use crate::request_draft::HeaderUpdate; -use crate::request_draft::RequestDraft; +use crate::client::HttpClientBackend; +use crate::client::TransportClient; +use crate::client_builder::ProxyRouting; use crate::tls_backend_fallback::RustlsClientCache; const MAX_CACHED_ROUTES: usize = 16; @@ -55,9 +52,10 @@ enum SelectedTlsBackend { pub struct RouteAwareClientPool { http_client_factory: HttpClientFactory, route_class: ClientRouteClass, - client_builder: HttpClientBuilder, + pub(crate) client_builder: HttpClientBuilder, + default_headers: HeaderMap, custom_ca_fallback: CustomCaFallback, - clients: Arc>>, + clients: Arc>>, rustls_clients: Option, } @@ -83,6 +81,8 @@ pub enum RouteAwareClientPoolError { /// Error returned while building, routing, or sending a route-aware request. #[derive(Debug, thiserror::Error)] pub enum RouteAwareRequestError { + #[error(transparent)] + Policy(#[from] NetworkPolicyDenied), #[error(transparent)] Request(#[from] reqwest::Error), #[error(transparent)] @@ -137,6 +137,7 @@ impl RouteAwareRequestError { Some(RouteFailureClass::ProxyResolutionUnavailable) } Self::Request(_) + | Self::Policy(_) | Self::Build(_) | Self::UnsupportedRedirectScheme(_) | Self::TooManyRedirects @@ -148,6 +149,7 @@ impl RouteAwareRequestError { match self { Self::Request(error) => error.status(), Self::Route(_) + | Self::Policy(_) | Self::Build(_) | Self::UnsupportedRedirectScheme(_) | Self::TooManyRedirects @@ -163,6 +165,10 @@ impl RouteAwareRequestError { matches!(self, Self::Request(error) if error.is_connect()) } + pub fn is_builder(&self) -> bool { + matches!(self, Self::Request(error) if error.is_builder()) + } + pub fn is_body(&self) -> bool { matches!(self, Self::Request(error) if error.is_body()) } @@ -171,6 +177,13 @@ impl RouteAwareRequestError { matches!(self, Self::Request(error) if error.is_request()) } + pub fn url_mut(&mut self) -> Option<&mut reqwest::Url> { + match self { + Self::Request(error) => error.url_mut(), + _ => None, + } + } + /// Removes a request URL from the underlying transport error before it is logged or returned. /// /// Use this for requests whose URL can contain credentials, such as signed blob uploads. @@ -182,135 +195,6 @@ impl RouteAwareRequestError { } } -#[must_use = "requests are not sent unless `send` is awaited"] -pub struct RouteAwareRequestBuilder { - pool: RouteAwareClientPool, - request: Result, -} - -impl fmt::Debug for RouteAwareRequestBuilder { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - let request = self.request.as_ref().ok(); - formatter - .debug_struct("RouteAwareRequestBuilder") - .field("pool", &self.pool) - .field("method", &request.map(|request| request.request.method())) - .field("url", &request.map(|_| "")) - .finish_non_exhaustive() - } -} - -impl RouteAwareRequestBuilder { - fn new(pool: RouteAwareClientPool, method: Method, url: U) -> Self - where - U: IntoUrl, - { - let request = RequestDraft::new(method, url).map_err(RouteAwareRequestError::Request); - Self { pool, request } - } - - pub fn headers(mut self, headers: HeaderMap) -> Self { - if let Ok(request) = &mut self.request { - request.extend_headers(headers); - } - self - } - - pub fn header(mut self, key: K, value: V) -> Self - where - HeaderName: TryFrom, - >::Error: Into, - HeaderValue: TryFrom, - >::Error: Into, - { - if let Ok(request) = &mut self.request { - let header = HeaderName::try_from(key) - .map_err(Into::into) - .and_then(|key| { - HeaderValue::try_from(value) - .map(|value| (key, value)) - .map_err(Into::into) - }); - match header { - Ok((key, value)) => { - request.headers.push(HeaderUpdate::Append(key, value)); - } - Err(error) => { - self.request = Err(RouteAwareRequestError::Build(error.to_string())); - } - } - } - self - } - - /// Sets a timeout for the request as a whole. - /// - /// The budget starts before outbound-route resolution and covers selecting or constructing a - /// pooled client, establishing a connection, sending the request, and awaiting the response. - /// Use [`HttpClientBuilder::connect_timeout`] when only connection establishment should be - /// bounded. - pub fn timeout(mut self, timeout: Duration) -> Self { - if let Ok(request) = &mut self.request { - *request.request.timeout_mut() = Some(timeout); - } - self - } - - pub fn json(mut self, value: &T) -> Self - where - T: ?Sized + Serialize, - { - if let Ok(request) = &mut self.request { - match serde_json::to_vec(value) { - Ok(body) => { - let has_content_type = request.headers.iter().any(|update| match update { - HeaderUpdate::Replace(headers) => headers.contains_key(CONTENT_TYPE), - HeaderUpdate::Append(name, _) => *name == CONTENT_TYPE, - }); - if !has_content_type { - request.headers.push(HeaderUpdate::Append( - CONTENT_TYPE, - HeaderValue::from_static("application/json"), - )); - } - *request.request.body_mut() = Some(body.into()); - } - Err(error) => { - self.request = Err(RouteAwareRequestError::Build(error.to_string())); - } - } - } - self - } - - pub fn body(mut self, body: B) -> Self - where - B: Into, - { - if let Ok(request) = &mut self.request { - *request.request.body_mut() = Some(body.into()); - } - self - } - - /// Sets a streaming request body without exposing the underlying HTTP implementation. - pub fn body_stream(mut self, stream: S) -> Self - where - S: TryStream + Send + 'static, - S::Error: Into>, - Bytes: From, - { - if let Ok(request) = &mut self.request { - *request.request.body_mut() = Some(reqwest::Body::wrap_stream(stream)); - } - self - } - - pub async fn send(self) -> Result { - self.pool.send(self.request?).await - } -} - impl RouteAwareClientPool { pub(crate) fn request_logging_enabled(&self) -> bool { self.client_builder.request_logging_enabled() @@ -330,6 +214,13 @@ impl RouteAwareClientPool { self } + /// Exposes the ordinary request builder while sending through this policy-aware pool. + pub(crate) fn into_client(self) -> HttpClient { + HttpClient { + backend: HttpClientBackend::Routed(Arc::new(self)), + } + } + /// Creates a pool with the shared default HTTP transport settings. pub fn new(http_client_factory: HttpClientFactory, route_class: ClientRouteClass) -> Self { Self::with_builder(http_client_factory, route_class, HttpClientBuilder::new()) @@ -379,21 +270,37 @@ impl RouteAwareClientPool { ) } - fn with_builder( + /// Configures protocol-specific transport behavior without bypassing destination checks. + pub fn with_builder( http_client_factory: HttpClientFactory, route_class: ClientRouteClass, - client_builder: HttpClientBuilder, + mut client_builder: HttpClientBuilder, ) -> Self { + let default_headers = client_builder.default_headers.take().unwrap_or_default(); Self { http_client_factory, route_class, client_builder, + default_headers, custom_ca_fallback: CustomCaFallback::Disabled, clients: Arc::new(Mutex::new(HashMap::new())), rustls_clients: None, } } + /// Starts a new account-scoped operation while sharing the existing connection cache. + /// Capture this before resolving credentials or reading account-specific content. + pub fn for_current_account(&self) -> Self { + let mut pool = self.clone(); + pool.http_client_factory = pool.http_client_factory.clone().with_network_policy( + pool.http_client_factory + .network_policy() + .clone() + .for_current_account(), + ); + pool + } + /// Retries recognized TLS protocol-negotiation failures once using rustls. /// /// Successful fallback is remembered for each HTTPS origin and resolved outbound route. @@ -496,46 +403,46 @@ impl RouteAwareClientPool { ) } - pub fn get(&self, url: U) -> RouteAwareRequestBuilder + pub fn get(&self, url: U) -> RequestBuilder where U: IntoUrl, { self.request(Method::GET, url) } - pub fn post(&self, url: U) -> RouteAwareRequestBuilder + pub fn post(&self, url: U) -> RequestBuilder where U: IntoUrl, { self.request(Method::POST, url) } - pub fn put(&self, url: U) -> RouteAwareRequestBuilder + pub fn put(&self, url: U) -> RequestBuilder where U: IntoUrl, { self.request(Method::PUT, url) } - pub fn delete(&self, url: U) -> RouteAwareRequestBuilder + pub fn delete(&self, url: U) -> RequestBuilder where U: IntoUrl, { self.request(Method::DELETE, url) } - pub fn request(&self, method: Method, url: U) -> RouteAwareRequestBuilder + pub fn request(&self, method: Method, url: U) -> RequestBuilder where U: IntoUrl, { - RouteAwareRequestBuilder::new(self.clone(), method, url) + RequestBuilder::routed(Arc::new(self.clone()), method, url) } async fn client_for_url_with_resolver( &self, request_url: &str, resolve_route: F, - ) -> Result<(OutboundProxyRoute, HttpClient, SelectedTlsBackend), RouteAwareClientPoolError> + ) -> Result<(OutboundProxyRoute, TransportClient, SelectedTlsBackend), RouteAwareClientPoolError> where F: FnOnce(String) -> Fut, Fut: Future>, @@ -564,16 +471,12 @@ impl RouteAwareClientPool { } else { self.client_builder.clone() }; - #[expect( - deprecated, - reason = "explicitly opted-in pools preserve the legacy custom-CA fallback" - )] let client = match ( self.http_client_factory.outbound_proxy_policy(), self.custom_ca_fallback, ) { (OutboundProxyPolicy::ReqwestDefault, CustomCaFallback::LegacyTransportDefault) => { - client_builder.build_with_transport_default_proxy_and_custom_ca_fallback() + client_builder.build_with_custom_ca_fallback(ProxyRouting::TransportDefault) } (OutboundProxyPolicy::ReqwestDefault, CustomCaFallback::Disabled) | (OutboundProxyPolicy::RespectSystemProxy, CustomCaFallback::Disabled) @@ -609,13 +512,14 @@ impl RouteAwareClientPool { self.client_builder.follows_redirects() && (self.http_client_factory.outbound_proxy_policy() == OutboundProxyPolicy::RespectSystemProxy + || self.http_client_factory.network_policy().is_managed() || self.rustls_clients.is_some()) } fn rustls_client_for_route( &self, route: &OutboundProxyRoute, - ) -> Result { + ) -> Result { let mut client_builder = self.client_builder.clone().with_rustls_tls(); if self.follows_redirects_manually() { client_builder = client_builder.without_redirects(); @@ -633,3 +537,7 @@ mod tests; #[cfg(test)] #[path = "route_aware_tls_fallback_tests.rs"] mod tls_fallback_tests; + +#[cfg(test)] +#[path = "route_aware_policy_tests.rs"] +mod policy_tests; diff --git a/codex-rs/http-client/src/route_aware_client_pool/execution.rs b/codex-rs/http-client/src/route_aware_client_pool/execution.rs index df4d11045e..7ef0d362c4 100644 --- a/codex-rs/http-client/src/route_aware_client_pool/execution.rs +++ b/codex-rs/http-client/src/route_aware_client_pool/execution.rs @@ -2,13 +2,18 @@ use std::future::Future; use std::io; +use std::sync::atomic::Ordering; +use futures::FutureExt; +use futures::future::BoxFuture; use http::header::PROXY_AUTHORIZATION; use super::RouteAwareClientPool; use super::RouteAwareRequestError; use super::SelectedTlsBackend; +use crate::HttpResponse; use crate::OutboundProxyRoute; +use crate::client::apply_default_headers; use crate::request_draft::RequestDraft; use crate::route_aware_redirect::MAX_REDIRECTS; use crate::route_aware_redirect::insert_referer; @@ -19,32 +24,55 @@ use crate::route_aware_redirect::remove_sensitive_headers; use crate::tls_backend_fallback::should_retry_with_rustls; impl RouteAwareClientPool { - pub(crate) async fn send( + /// Executes requests through shared routing, destination, and redirect checks. + /// + /// Box the execution future so nested callers do not inherit policy and routing state. + pub fn execute( + &self, + request: reqwest::Request, + ) -> BoxFuture<'_, Result> { + self.send(request.into()) + } + + pub(crate) fn send( &self, request: RequestDraft, - ) -> Result { - let http_client_factory = self.http_client_factory.clone(); - self.send_with_resolver(request, move |request_url| { - let http_client_factory = http_client_factory.clone(); - async move { - http_client_factory - .resolve_proxy_route_async(request_url) - .await - } + ) -> BoxFuture<'_, Result> { + self.send_with_resolver(request, |request_url| { + self.http_client_factory + .resolve_proxy_route_async(request_url) }) - .await + .boxed() } pub(super) async fn send_with_resolver( &self, request: impl Into, resolve_route: F, - ) -> Result + ) -> Result + where + F: Fn(String) -> Fut, + Fut: Future>, + { + let request = request.into(); + let permit = self + .http_client_factory + .network_policy() + .acquire(request.request.url())?; + permit + .run(Box::pin(self.send_authorized(request, resolve_route))) + .await? + } + + async fn send_authorized( + &self, + draft: RequestDraft, + resolve_route: F, + ) -> Result where F: Fn(String) -> Fut, Fut: Future>, { - let draft = request.into(); let mut request = draft.request; let mut draft_headers = Some(draft.headers); let request_method = request.method().clone(); @@ -57,102 +85,105 @@ impl RouteAwareClientPool { let mut redirects = 0; let mut previous_route = None; loop { - let (current_route, client, selected_tls_backend) = match timeout_deadline { - Some(timeout_deadline) => tokio::time::timeout_at( - timeout_deadline, - self.client_for_url_with_resolver(request.url().as_str(), &resolve_route), - ) - .await - .map_err(|_| RouteAwareRequestError::Timeout)??, - None => { - self.client_for_url_with_resolver(request.url().as_str(), &resolve_route) - .await? - } - }; - if let Some(headers) = draft_headers.take() { - request = RequestDraft { request, headers }.build(&client)?; - request_url = request.url().to_string(); - } let current_url = request.url().clone(); - if previous_route - .as_ref() - .is_some_and(|previous_route| previous_route != ¤t_route) - { - request.headers_mut().remove(PROXY_AUTHORIZATION); - } - previous_route = Some(current_route.clone()); - if let Some(timeout_deadline) = timeout_deadline { - let remaining = timeout_deadline - .checked_duration_since(tokio::time::Instant::now()) - .ok_or(RouteAwareRequestError::Timeout)?; - if remaining.is_zero() { - return Err(RouteAwareRequestError::Timeout); + let permit = self + .http_client_factory + .network_policy() + .acquire(¤t_url)?; + let hop = async { + let (current_route, mut client, selected_tls_backend) = self + .client_for_url_with_resolver(current_url.as_str(), &resolve_route) + .await?; + client.request_logging = self.client_builder.request_logging; + if let Some(headers) = draft_headers.take() { + request = RequestDraft { request, headers }.build(&client)?; + request_url = request.url().to_string(); + apply_default_headers(request.headers_mut(), &self.default_headers); } - *request.timeout_mut() = Some(remaining); - } - let method = request.method().clone(); - let headers = request.headers().clone(); - let version = request.version(); - let timeout = request.timeout().copied(); - let replay = request.try_clone(); - let execute_request = async { - if follows_redirects_manually { + let current_url = request.url().clone(); + if previous_route + .as_ref() + .is_some_and(|previous_route| previous_route != ¤t_route) + { + request.headers_mut().remove(PROXY_AUTHORIZATION); + } + previous_route = Some(current_route.clone()); + if let Some(timeout_deadline) = timeout_deadline { + let remaining = timeout_deadline + .checked_duration_since(tokio::time::Instant::now()) + .ok_or(RouteAwareRequestError::Timeout)?; + if remaining.is_zero() { + return Err(RouteAwareRequestError::Timeout); + } + *request.timeout_mut() = Some(remaining); + } + let method = request.method().clone(); + let headers = request.headers().clone(); + let version = request.version(); + let timeout = request.timeout().copied(); + let replay = request.try_clone(); + let result = if follows_redirects_manually { client.execute_without_request_logging(request).await } else { client.execute(request).await - } - }; - let response = match match timeout_deadline { - Some(timeout_deadline) => { - tokio::time::timeout_at(timeout_deadline, execute_request) - .await - .map_err(|_| RouteAwareRequestError::Timeout)? - } - None => execute_request.await, - } { - Ok(response) => response, - Err(error) => { - let result = self - .retry_with_rustls( - ¤t_url, - ¤t_route, - selected_tls_backend, - replay.as_ref(), - error, - timeout_deadline, - ) - .await; - - if follows_redirects_manually - && let Err(RouteAwareRequestError::Request(error)) = &result - { - client.log_error_summary(&request_method, &request_url, error); + }; + let response = match result { + Ok(response) => response, + Err(error) => { + let result = self + .retry_with_rustls( + ¤t_url, + ¤t_route, + selected_tls_backend, + replay.as_ref(), + error, + timeout_deadline, + ) + .await; + if follows_redirects_manually + && let Err(RouteAwareRequestError::Request(error)) = &result + { + client.log_error_summary(&request_method, &request_url, error); + } + result? } - - result? - } - }; - let status = response.status(); - if !follows_redirects_manually || !is_redirect(status) { - if follows_redirects_manually { + }; + let next_request = if follows_redirects_manually && is_redirect(response.status()) { + redirect_url(&response).and_then(|next_url| { + redirect_request( + response.status(), + method, + headers, + version, + timeout, + replay, + next_url, + ) + }) + } else { + None + }; + if follows_redirects_manually && next_request.is_none() { client.log_response(&request_method, &request_url, &response); } - return Ok(response); + Ok::<_, RouteAwareRequestError>((response, next_request, current_url)) + }; + let (response, next_request, current_url) = permit + .run(async { + match timeout_deadline { + Some(deadline) => tokio::time::timeout_at(deadline, hop) + .await + .map_err(|_| RouteAwareRequestError::Timeout)?, + None => hop.await, + } + }) + .await??; + let Some(mut next_request) = next_request else { + return Ok(HttpResponse::new(response, permit)); + }; + if let Some(redirect_observed) = &self.client_builder.redirect_observed { + redirect_observed.store(/*val*/ true, Ordering::Relaxed); } - let Some(next_url) = redirect_url(&response) else { - if follows_redirects_manually { - client.log_response(&request_method, &request_url, &response); - } - return Ok(response); - }; - let Some(mut next_request) = - redirect_request(status, method, headers, version, timeout, replay, next_url) - else { - if follows_redirects_manually { - client.log_response(&request_method, &request_url, &response); - } - return Ok(response); - }; let next_request_url = next_request.url().clone(); if !matches!(next_request_url.scheme(), "http" | "https") { return Err(RouteAwareRequestError::UnsupportedRedirectScheme( @@ -193,10 +224,11 @@ impl RouteAwareClientPool { return Err(error.into()); }; - let fallback_client = match rustls_clients.client_for_route(current_route) { + let mut fallback_client = match rustls_clients.client_for_route(current_route) { Some(client) => client, None => self.rustls_client_for_route(current_route)?, }; + fallback_client.request_logging = self.client_builder.request_logging; if let Some(timeout_deadline) = timeout_deadline { let remaining = timeout_deadline @@ -208,20 +240,12 @@ impl RouteAwareClientPool { *retry_request.timeout_mut() = Some(remaining); } - let execute_retry = async { - if self.follows_redirects_manually() { - fallback_client - .execute_without_request_logging(retry_request) - .await - } else { - fallback_client.execute(retry_request).await - } - }; - let response = match timeout_deadline { - Some(timeout_deadline) => tokio::time::timeout_at(timeout_deadline, execute_retry) + let response = if self.follows_redirects_manually() { + fallback_client + .execute_without_request_logging(retry_request) .await - .map_err(|_| RouteAwareRequestError::Timeout)?, - None => execute_retry.await, + } else { + fallback_client.execute(retry_request).await }?; rustls_clients.remember(current_url, current_route, fallback_client); diff --git a/codex-rs/http-client/src/route_aware_client_pool_tests.rs b/codex-rs/http-client/src/route_aware_client_pool_tests.rs index 4906f571ea..af1a0cdcd1 100644 --- a/codex-rs/http-client/src/route_aware_client_pool_tests.rs +++ b/codex-rs/http-client/src/route_aware_client_pool_tests.rs @@ -13,116 +13,14 @@ use std::time::Instant; use bytes::Bytes; use futures::stream; +use http::HeaderValue; use pretty_assertions::assert_eq; use tracing_subscriber::Layer; use tracing_subscriber::layer::SubscriberExt; use super::*; use crate::OutboundProxyPolicy; - -#[test] -fn request_draft_preserves_reqwest_url_auth_and_header_precedence() { - let pool = manual_redirect_pool(); - let client = HttpClientBuilder::new() - .build_direct() - .expect("client should build"); - for url in [ - "https://example.com/path", - "https://user:p%40ss@example.com/path", - "https://:%E2%98%83@example.com/path", - "https://%FF:password@example.com/path", - ] { - let headers = HeaderMap::from_iter([ - ( - http::header::AUTHORIZATION, - HeaderValue::from_static("Bearer replacement"), - ), - ( - CONTENT_TYPE, - HeaderValue::from_static("application/custom+json"), - ), - ]); - for (actual, expected) in [ - (pool.post(url), client.post(url)), - ( - pool.post(url) - .header(http::header::AUTHORIZATION, "Bearer appended"), - client - .post(url) - .header(http::header::AUTHORIZATION, "Bearer appended"), - ), - ( - pool.post(url) - .header(http::header::AUTHORIZATION, "Bearer first") - .headers(headers.clone()) - .header(http::header::AUTHORIZATION, "Bearer last"), - client - .post(url) - .header(http::header::AUTHORIZATION, "Bearer first") - .headers(headers) - .header(http::header::AUTHORIZATION, "Bearer last"), - ), - ( - pool.post(url) - .header(CONTENT_TYPE, "application/custom+json") - .headers(HeaderMap::new()), - client - .post(url) - .header(CONTENT_TYPE, "application/custom+json") - .headers(HeaderMap::new()), - ), - ( - pool.post(url).json(&serde_json::json!({"previous": true})), - client - .post(url) - .json(&serde_json::json!({"previous": true})), - ), - ] { - let body = serde_json::json!({"query": "test"}); - let timeout = Duration::from_secs(7); - let actual = actual - .json(&body) - .timeout(timeout) - .request - .expect("request inputs should be valid") - .build(&client) - .expect("request should build"); - let expected = expected - .json(&body) - .timeout(timeout) - .build() - .expect("reqwest should build"); - assert_eq!( - ( - actual.url(), - actual.method(), - actual.headers(), - actual.body().and_then(reqwest::Body::as_bytes), - actual.timeout() - ), - ( - expected.url(), - expected.method(), - expected.headers(), - expected.body().and_then(reqwest::Body::as_bytes), - expected.timeout() - ), - ); - assert_eq!( - actual - .headers() - .values() - .map(HeaderValue::is_sensitive) - .collect::>(), - expected - .headers() - .values() - .map(HeaderValue::is_sensitive) - .collect::>(), - ); - } - } -} +use crate::request_draft::RequestDraft; #[tokio::test] async fn initial_url_auth_is_built_once_after_routing_and_retained_only_on_same_origin_redirects() { @@ -153,11 +51,8 @@ async fn initial_url_auth_is_built_once_after_routing_and_retained_only_on_same_ .collect(), ); let pool = manual_redirect_pool(); - let request = pool - .get(urls[0]) - .timeout(Duration::from_secs(3)) - .request - .expect("valid request"); + let mut request = RequestDraft::new(Method::GET, urls[0]).expect("valid request"); + *request.request.timeout_mut() = Some(Duration::from_secs(/*secs*/ 3)); let response = pool .send_with_resolver(request, |url| resolver.resolve(url)) .await @@ -220,6 +115,64 @@ async fn fixed_transport_preserves_url_auth_and_explicit_authorization() { } } +#[tokio::test] +async fn direct_and_routed_clients_build_equivalent_requests() { + let (address, server) = spawn_response_server(vec![ + "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_string(); 4 + ]); + let mut defaults = HeaderMap::new(); + defaults.append("x-values", "default-one".parse().unwrap()); + defaults.append("x-values", "default-two".parse().unwrap()); + let builder = HttpClientBuilder::new().default_headers(defaults.clone()); + let direct = builder.build_direct().unwrap(); + assert!(!format!("{direct:?}").contains("default-one")); + let controller = crate::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish(policy.revision(), crate::DestinationPolicy::Unrestricted); + let routed = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault).with_network_policy(policy), + ClientRouteClass::Api, + HttpClientBuilder::new().default_headers(defaults), + ) + .into_client(); + for client in [direct, routed] { + client + .post(format!("http://{address}/submit?existing=1")) + .query(&[("value", "hello world"), ("value", "second")]) + .header("x-values", "old") + .headers(HeaderMap::from_iter([( + "x-values".parse().unwrap(), + "first".parse().unwrap(), + )])) + .header("x-values", "second") + .bearer_auth("test-token") + .timeout(Duration::from_secs(/*secs*/ 5)) + .version(http::Version::HTTP_11) + .json(&serde_json::json!({"value": 1})) + .send() + .await + .unwrap(); + let auth = client.get(format!("http://alice:p%40ss@{address}/auth")); + auth.send().await.unwrap(); + } + let requests = server.join().unwrap(); + assert_eq!(requests[0], requests[2]); + assert_eq!(requests[1], requests[3]); + assert!(requests[1].contains("authorization: Basic YWxpY2U6cEBzcw==\r\n")); + assert!(requests[1].contains("x-values: default-one\r\n")); + assert!(requests[1].contains("x-values: default-two\r\n")); + let request = &requests[0]; + assert!( + request.starts_with("POST /submit?existing=1&value=hello+world&value=second HTTP/1.1\r\n") + ); + assert!(request.contains("authorization: Bearer test-token\r\n")); + assert!(request.contains("content-type: application/json\r\n")); + assert!(request.contains("x-values: first\r\n")); + assert!(request.contains("x-values: second\r\n")); + assert!(!request.contains("x-values: old") && !request.contains("x-values: default")); + assert!(request.ends_with("\r\n\r\n{\"value\":1}")); +} + #[tokio::test] async fn request_failures_classify_real_untrusted_certificate_handshakes() { codex_utils_rustls_provider::ensure_rustls_crypto_provider(); @@ -252,11 +205,11 @@ async fn request_failures_classify_real_untrusted_certificate_handshakes() { ClientRouteClass::Api, ); - let request = pool - .get(format!("https://localhost:{}/", address.port())) - .timeout(Duration::from_secs(3)) - .request - .expect("TLS request should build"); + let mut request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&format!("https://localhost:{}/", address.port())).unwrap(), + ); + *request.timeout_mut() = Some(Duration::from_secs(3)); let error = pool .send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) }) .await @@ -322,14 +275,7 @@ fn request_builder_debug_redacts_url_secrets() { assert_eq!( format!("{request:?}"), - concat!( - "RouteAwareRequestBuilder { pool: RouteAwareClientPool { ", - "http_client_factory: HttpClientFactory { outbound_proxy_policy: ReqwestDefault, ", - "system_proxy_fallback: false, ", - "network_policy: NetworkPolicy { managed: false, .. } }, ", - "route_class: Api, .. }, method: Some(GET), ", - "url: Some(\"\"), .. }" - ) + "RequestBuilder { method: Some(GET), url: \"\", .. }" ); } @@ -515,7 +461,11 @@ async fn cached_tls_backend_only_changes_its_destination_and_route() { ])); let fallback_client = HttpClientBuilder::new() .with_rustls_tls() - .build_direct() + .build_for_resolved_route( + &HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + &OutboundProxyRoute::Direct, + ) .expect("rustls client should build without proxy autodiscovery"); pool.rustls_clients .as_ref() @@ -976,7 +926,7 @@ async fn resolve_with( pool: &RouteAwareClientPool, resolver: &FakeRouteResolver, request_url: &str, -) -> Result { +) -> Result { let resolver = resolver.clone(); let (_, client, _) = pool .client_for_url_with_resolver(request_url, move |request_url| async move { @@ -994,7 +944,7 @@ fn manual_redirect_pool() -> RouteAwareClientPool { ) } -fn spawn_response_server( +pub(super) fn spawn_response_server( responses: Vec, ) -> (std::net::SocketAddr, std::thread::JoinHandle>) { let listener = TcpListener::bind(("127.0.0.1", 0)).expect("response listener should bind"); diff --git a/codex-rs/http-client/src/route_aware_policy_tests.rs b/codex-rs/http-client/src/route_aware_policy_tests.rs new file mode 100644 index 0000000000..ec88589bef --- /dev/null +++ b/codex-rs/http-client/src/route_aware_policy_tests.rs @@ -0,0 +1,95 @@ +//! Unavailable policy rejects requests before a socket is opened. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn unavailable_policy_does_not_connect() { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let url = format!("https://{}/", listener.local_addr().unwrap()); + let controller = crate::NetworkPolicyController::default(); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(controller.policy()); + let client = factory.build_client(&url, ClientRouteClass::Api).unwrap(); + assert!(matches!( + client.get(&url).send().await, + Err(RouteAwareRequestError::Policy( + NetworkPolicyDenied::Unavailable + )) + )); + assert_eq!( + listener.accept().unwrap_err().kind(), + std::io::ErrorKind::WouldBlock + ); +} + +#[tokio::test] +async fn account_change_cancels_redirect_route_selection() { + let (address, server) = super::tests::spawn_response_server(vec![ + "HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + .into(), + ]); + let controller = crate::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish(policy.revision(), crate::DestinationPolicy::Unrestricted); + let pool = RouteAwareClientPool::new( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy.clone()), + ClientRouteClass::Api, + ); + let request = reqwest::Request::new( + Method::GET, + reqwest::Url::parse(&format!("http://{address}/start")).unwrap(), + ); + let result = tokio::time::timeout( + std::time::Duration::from_secs(/*secs*/ 5), + pool.send_with_resolver(request, |url| { + let policy = policy.clone(); + async move { + if url.ends_with("/final") { + policy.invalidate(); + return std::future::pending().await; + } + Ok(OutboundProxyRoute::Direct) + } + }), + ) + .await + .unwrap(); + assert!(matches!( + result, + Err(RouteAwareRequestError::Policy(NetworkPolicyDenied::Revoked)) + )); + assert_eq!(server.join().unwrap().len(), 1); +} + +#[tokio::test] +async fn managed_redirect_failure_is_tracked_before_proxy_retry() { + let unused = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let dead_address = unused.local_addr().unwrap(); + drop(unused); + let (address, server) = super::tests::spawn_response_server(vec![format!( + "HTTP/1.1 307 Temporary Redirect\r\nLocation: http://{dead_address}/token\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + )]); + let controller = crate::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish(policy.revision(), crate::DestinationPolicy::Unrestricted); + let redirected = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(/*v*/ false)); + let pool = RouteAwareClientPool::with_builder( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy) + .with_system_proxy_fallback(), + ClientRouteClass::Auth, + HttpClientBuilder::new().with_redirect_tracking(redirected.clone()), + ); + let result = pool + .post(format!("http://{address}/token")) + .body("code=one-time-code") + .timeout(std::time::Duration::from_secs(/*secs*/ 5)) + .send() + .await; + assert!(result.is_err()); + assert!(redirected.load(std::sync::atomic::Ordering::Relaxed)); + assert_eq!(server.join().unwrap().len(), 1); +} diff --git a/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs b/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs index e596803090..dc716e3a77 100644 --- a/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs +++ b/codex-rs/http-client/src/route_aware_tls_fallback_tests.rs @@ -16,7 +16,6 @@ use rcgen::generate_simple_self_signed; use rustls_pki_types::PrivateKeyDer; use super::ClientRouteClass; -use super::HttpClient; use super::HttpClientFactory; use super::Method; use super::OutboundProxyPolicy; @@ -24,11 +23,17 @@ use super::OutboundProxyRoute; use super::RouteAwareClientPool; use super::RouteAwareRequestError; use super::SelectedTlsBackend; +use crate::client::RequestLogging; +use crate::client::TransportClient; use crate::tls_backend_fallback::should_retry_with_rustls; const PROTOCOL_VERSION_TLS_ALERT: &[u8] = &[21, 3, 3, 0, 2, 2, 70]; -type SuccessfulTlsFallbackServer = (String, HttpClient, mpsc::Receiver>>); +type SuccessfulTlsFallbackServer = ( + String, + TransportClient, + mpsc::Receiver>>, +); #[tokio::test] async fn default_pool_does_not_retry_a_native_tls_protocol_failure() { @@ -346,7 +351,7 @@ fn spawn_successful_tls_fallback_server() -> io::Result io::Result, - clients: HashMap, + clients: HashMap, } #[derive(Clone, Hash, PartialEq, Eq)] @@ -46,7 +46,7 @@ impl RustlsClientCache { .contains(&destination) } - pub(crate) fn client_for_route(&self, route: &OutboundProxyRoute) -> Option { + pub(crate) fn client_for_route(&self, route: &OutboundProxyRoute) -> Option { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -59,7 +59,7 @@ impl RustlsClientCache { &self, url: &reqwest::Url, route: &OutboundProxyRoute, - client: HttpClient, + client: TransportClient, ) { let Some(destination) = DestinationRoute::new(url, route) else { return; diff --git a/codex-rs/http-client/src/tls_backend_fallback_tests.rs b/codex-rs/http-client/src/tls_backend_fallback_tests.rs index c7ca0adceb..f1b2778d52 100644 --- a/codex-rs/http-client/src/tls_backend_fallback_tests.rs +++ b/codex-rs/http-client/src/tls_backend_fallback_tests.rs @@ -107,10 +107,7 @@ fn rustls_fallback_decisions_are_scoped_to_origin_and_outbound_route() { no_proxy: None, }; - let client = HttpClientBuilder::new() - .with_rustls_tls() - .build_direct() - .expect("rustls client should build without proxy autodiscovery"); + let client = direct_rustls_client(); cache.remember(&destination, &direct, client); assert_eq!( @@ -134,10 +131,7 @@ fn cached_rustls_clients_are_reused_for_the_same_outbound_route() { let second_destination = reqwest::Url::parse("https://second.example.com").expect("valid second URL"); let direct = OutboundProxyRoute::Direct; - let client = HttpClientBuilder::new() - .with_rustls_tls() - .build_direct() - .expect("rustls client should build without proxy autodiscovery"); + let client = direct_rustls_client(); cache.remember(&first_destination, &direct, client.clone()); cache.remember(&second_destination, &direct, client); @@ -161,10 +155,7 @@ fn cached_rustls_clients_are_reused_for_the_same_outbound_route() { #[test] fn cached_rustls_destinations_remain_bounded_while_sharing_a_route_client() { let cache = RustlsClientCache::default(); - let client = HttpClientBuilder::new() - .with_rustls_tls() - .build_direct() - .expect("rustls client should build without proxy autodiscovery"); + let client = direct_rustls_client(); for index in 0..=MAX_CACHED_RUSTLS_DESTINATIONS { let destination = @@ -182,10 +173,7 @@ fn cached_rustls_destinations_remain_bounded_while_sharing_a_route_client() { #[test] fn evicting_a_destination_removes_its_unshared_route_client() { let cache = RustlsClientCache::default(); - let client = HttpClientBuilder::new() - .with_rustls_tls() - .build_direct() - .expect("rustls client should build without proxy autodiscovery"); + let client = direct_rustls_client(); for index in 0..=MAX_CACHED_RUSTLS_DESTINATIONS { let destination = @@ -214,3 +202,14 @@ fn evicting_a_destination_removes_its_unshared_route_client() { ) ); } + +fn direct_rustls_client() -> crate::client::TransportClient { + HttpClientBuilder::new() + .with_rustls_tls() + .build_for_resolved_route( + &crate::HttpClientFactory::new(crate::OutboundProxyPolicy::ReqwestDefault), + crate::ClientRouteClass::Api, + &OutboundProxyRoute::Direct, + ) + .expect("rustls client should build without proxy autodiscovery") +} diff --git a/codex-rs/http-client/src/transport.rs b/codex-rs/http-client/src/transport.rs index 08792b9e74..9329f5e804 100644 --- a/codex-rs/http-client/src/transport.rs +++ b/codex-rs/http-client/src/transport.rs @@ -3,9 +3,11 @@ //! Each request's limit applies to observed body bytes, including unsuccessful responses; //! declared content lengths are only an additional early-rejection check. +use crate::HttpResponse; use crate::RouteAwareClientPool; use crate::RouteAwareRequestError; use crate::client::HttpClient; +use crate::client::HttpClientBackend; use crate::error::TransportError; use crate::request::Request; use crate::request::RequestBody; @@ -70,7 +72,7 @@ impl ReqwestTransport { } } - async fn send(&self, req: Request) -> Result { + async fn send(&self, req: Request) -> Result { let prepared = req.prepare_body_for_send().map_err(TransportError::Build)?; let Request { @@ -84,16 +86,29 @@ impl ReqwestTransport { } = req; let method = Method::from_bytes(method.as_str().as_bytes()).unwrap_or(Method::GET); - let mut request = RequestDraft::new(method, url).map_err(Self::map_error)?; + let mut request = + RequestDraft::new(method, url).map_err(|error| Self::map_error(error.into()))?; request.extend_headers(prepared.headers); *request.request.timeout_mut() = timeout; *request.request.body_mut() = prepared.body.map(Into::into); match &self.client { - TransportClient::Fixed(client) => { - let request = request.build(client).map_err(Self::map_error)?; - client.execute(request).await.map_err(Self::map_error) - } + TransportClient::Fixed(client) => match &client.backend { + HttpClientBackend::Routed(pool) => pool + .send(request) + .await + .map_err(Self::map_route_aware_error), + HttpClientBackend::Direct(client) => { + let request = request + .build(client) + .map_err(|error| Self::map_error(error.into()))?; + client + .execute(request) + .await + .map(HttpResponse::from) + .map_err(|error| Self::map_error(error.into())) + } + }, TransportClient::RouteAware(client) => client .send(request) .await @@ -101,9 +116,11 @@ impl ReqwestTransport { } } - fn map_error(err: reqwest::Error) -> TransportError { + fn map_error(err: crate::HttpError) -> TransportError { let err = err.without_url(); - if err.is_connect() { + if let crate::HttpError::Policy(denied) = err { + TransportError::Policy(denied) + } else if err.is_connect() { TransportError::Connection(err.without_url()) } else if err.is_timeout() { TransportError::Timeout @@ -114,7 +131,8 @@ impl ReqwestTransport { fn map_route_aware_error(error: RouteAwareRequestError) -> TransportError { match error.without_url() { - RouteAwareRequestError::Request(error) => Self::map_error(error), + RouteAwareRequestError::Policy(denied) => TransportError::Policy(denied), + RouteAwareRequestError::Request(error) => Self::map_error(error.into()), RouteAwareRequestError::Timeout => TransportError::Timeout, RouteAwareRequestError::Route(error) => TransportError::Build(error.to_string()), RouteAwareRequestError::Build(error) => TransportError::Build(error), @@ -168,7 +186,9 @@ impl HttpTransport for ReqwestTransport { if !status.is_success() { let body = match bytes { Ok(bytes) => String::from_utf8(bytes.to_vec()).ok(), - Err(error @ TransportError::ResponseTooLarge { .. }) => return Err(error), + Err( + error @ (TransportError::ResponseTooLarge { .. } | TransportError::Policy(_)), + ) => return Err(error), // Keep bounded diagnostic-body failures from hiding HTTP auth/retry status. Err(_) if response_body_limit_bytes.is_some() => None, Err(error) => return Err(error), @@ -205,11 +225,17 @@ impl HttpTransport for ReqwestTransport { *buffered.headers_mut() = headers.clone(); reqwest::Response::from(buffered).text().await.ok() } - Err(error @ TransportError::ResponseTooLarge { .. }) => return Err(error), + Err( + error @ (TransportError::ResponseTooLarge { .. } + | TransportError::Policy(_)), + ) => return Err(error), // A failed diagnostic body must not hide HTTP auth or retry semantics. Err(_) => None, }, - None => resp.text().await.ok(), + None => match resp.text().await { + Err(crate::HttpError::Policy(denied)) => return Err(denied.into()), + body => body.ok(), + }, }; return Err(TransportError::Http { status, @@ -234,7 +260,7 @@ impl HttpTransport for ReqwestTransport { } fn bounded_response_stream( - response: reqwest::Response, + response: HttpResponse, max_bytes: usize, ) -> Result { if response @@ -265,7 +291,7 @@ fn bounded_response_stream( } async fn bounded_response_bytes( - response: reqwest::Response, + response: HttpResponse, max_bytes: usize, ) -> Result { let mut stream = bounded_response_stream(response, max_bytes)?; diff --git a/codex-rs/http-client/src/transport_limit_tests.rs b/codex-rs/http-client/src/transport_limit_tests.rs index 1c977854a1..a1a70b384c 100644 --- a/codex-rs/http-client/src/transport_limit_tests.rs +++ b/codex-rs/http-client/src/transport_limit_tests.rs @@ -124,6 +124,46 @@ async fn rejects_oversized_declared_lengths_before_reading_the_body() { } } +#[tokio::test] +async fn bounded_stream_retains_account_revocation() { + let (finish, finished) = std::sync::mpsc::channel(); + let (url, server) = start_server(move |connection| { + connection + .write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\npart\r\n") + .unwrap(); + finished.recv().unwrap(); + }); + let controller = crate::NetworkPolicyController::default(); + let policy = controller.policy(); + assert!(controller.publish(policy.revision(), crate::DestinationPolicy::Unrestricted)); + let factory = crate::HttpClientFactory::new(crate::OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy.clone().for_current_account()); + let client = factory + .build_client(&url, crate::ClientRouteClass::Api) + .unwrap(); + let transport = ReqwestTransport::from_http_client(client); + let mut request = Request::new(Method::GET, url); + request.response_body_limit_bytes = Some(5); + let mut response = transport.stream(request).await.unwrap(); + assert_eq!( + response.bytes.next().await.unwrap().unwrap(), + Bytes::from_static(b"part") + ); + policy.invalidate(); + let error = tokio::time::timeout(Duration::from_secs(/*secs*/ 1), response.bytes.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(matches!( + error, + TransportError::Policy(crate::NetworkPolicyDenied::Revoked) + )); + assert!(response.bytes.next().await.is_none()); + finish.send(()).unwrap(); + server.join().unwrap(); +} + #[tokio::test] async fn streaming_limit_counts_across_chunks_and_terminates_on_error() { let (continue_tx, continue_rx) = std::sync::mpsc::channel(); diff --git a/codex-rs/http-client/src/transport_tests.rs b/codex-rs/http-client/src/transport_tests.rs index 07e49f97f8..f8240ba245 100644 --- a/codex-rs/http-client/src/transport_tests.rs +++ b/codex-rs/http-client/src/transport_tests.rs @@ -25,6 +25,20 @@ async fn disabled_request_logging_suppresses_transport_url_and_body() { assert!(logs.contains("log capture sentinel")); assert!(!logs.contains("url-secret")); assert!(!logs.contains("body-secret")); + + let controller = crate::NetworkPolicyController::default(); + let policy = controller.policy(); + controller.publish(policy.revision(), crate::DestinationPolicy::Unrestricted); + let factory = crate::HttpClientFactory::new(crate::OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy); + let client = + crate::RouteAwareClientPool::new(factory, crate::ClientRouteClass::Api).into_client(); + let logs = capture_transport_logs(client.clone()).await; + assert!(logs.contains("url-secret")); + let logs = capture_transport_logs(client.clone().without_request_logging()).await; + assert!(!logs.contains("url-secret"), "{logs}"); + assert!(!logs.contains("body-secret")); + assert!(capture_transport_logs(client).await.contains("url-secret")); } #[tokio::test] @@ -72,7 +86,8 @@ async fn capture_transport_logs(client: HttpClient) -> String { .with_writer(move || TestLogWriter(Arc::clone(&writer_buffer))) .with_filter( tracing_subscriber::filter::Targets::new() - .with_target("codex_http_client::transport", tracing::Level::TRACE), + .with_target("codex_http_client::transport", tracing::Level::TRACE) + .with_target("codex_http_client::client", tracing::Level::DEBUG), ), ); let _guard = tracing::subscriber::set_default(subscriber); diff --git a/codex-rs/protocol/src/error_tests.rs b/codex-rs/protocol/src/error_tests.rs index 921fb79ce7..8819c7fd68 100644 --- a/codex-rs/protocol/src/error_tests.rs +++ b/codex-rs/protocol/src/error_tests.rs @@ -267,10 +267,10 @@ fn to_error_event_handles_response_stream_failed() { .status(StatusCode::TOO_MANY_REQUESTS) .body("") .unwrap(); - let source = HttpResponse::from(response) + let mut source = HttpResponse::from(response) .error_for_status_ref() - .unwrap_err() - .with_url("http://example.com".parse().unwrap()); + .unwrap_err(); + *source.url_mut().unwrap() = "http://example.com".parse().unwrap(); let err = CodexErr::ResponseStreamFailed(ResponseStreamFailed { source, request_id: Some("req-123".to_string()), diff --git a/codex-rs/response-debug-context/src/lib.rs b/codex-rs/response-debug-context/src/lib.rs index 9014af3e80..ad65c0b1fc 100644 --- a/codex-rs/response-debug-context/src/lib.rs +++ b/codex-rs/response-debug-context/src/lib.rs @@ -66,6 +66,7 @@ pub fn telemetry_transport_error_message(error: &TransportError) -> String { TransportError::Connection(err) => err.to_string(), TransportError::Network(err) => err.to_string(), TransportError::Build(err) => err.to_string(), + TransportError::Policy(denied) => denied.to_string(), } } diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 272b25d56a..5be7b53631 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -5981,14 +5981,14 @@ async fn clear_ui_header_shows_fast_status_for_fast_capable_models() { assert_app_snapshot!("clear_ui_header_fast_status_fast_capable_models", rendered); } -async fn make_test_app() -> App { +async fn make_test_app() -> Box { let (chat_widget, app_event_tx, _rx, _op_rx) = make_chatwidget_manual_with_sender().await; let config = chat_widget.config_ref().clone(); let file_search = FileSearchManager::new(config.cwd.to_path_buf(), app_event_tx.clone()); let model = get_model_offline_for_tests(config.model.as_deref()); let session_telemetry = test_session_telemetry(&config, model.as_str()); - App { + Box::new(App { feature_write_lock: Arc::default(), model_catalog: chat_widget.model_catalog(), session_telemetry, @@ -6079,11 +6079,11 @@ async fn make_test_app() -> App { pending_plugin_enabled_writes: HashMap::new(), pending_hook_enabled_writes: HashMap::new(), recap: recap::RecapState::default(), - } + }) } pub(super) async fn make_test_app_with_channels() -> ( - App, + Box, tokio::sync::mpsc::UnboundedReceiver, tokio::sync::mpsc::UnboundedReceiver, ) { @@ -6094,7 +6094,7 @@ pub(super) async fn make_test_app_with_channels() -> ( let session_telemetry = test_session_telemetry(&config, model.as_str()); ( - App { + Box::new(App { feature_write_lock: Arc::default(), model_catalog: chat_widget.model_catalog(), session_telemetry, @@ -6185,7 +6185,7 @@ pub(super) async fn make_test_app_with_channels() -> ( pending_plugin_enabled_writes: HashMap::new(), pending_hook_enabled_writes: HashMap::new(), recap: recap::RecapState::default(), - }, + }), rx, op_rx, ) diff --git a/codex-rs/tui/src/app/tests/history_hydration_tests.rs b/codex-rs/tui/src/app/tests/history_hydration_tests.rs index 9f28aecc6f..97c983cb66 100644 --- a/codex-rs/tui/src/app/tests/history_hydration_tests.rs +++ b/codex-rs/tui/src/app/tests/history_hydration_tests.rs @@ -27,7 +27,9 @@ use codex_state::SqliteConfig; use pretty_assertions::assert_eq; use pretty_assertions::assert_ne; -async fn history_fixture(item_counts: &[usize]) -> Result<(App, tempfile::TempDir, SessionTarget)> { +async fn history_fixture( + item_counts: &[usize], +) -> Result<(Box, tempfile::TempDir, SessionTarget)> { let mut app = make_test_app().await; let codex_home = tempdir()?; app.config.codex_home = codex_home.path().to_path_buf().abs(); diff --git a/codex-rs/tui/src/app/tests/key_chords.rs b/codex-rs/tui/src/app/tests/key_chords.rs index b8f165dede..2143be9401 100644 --- a/codex-rs/tui/src/app/tests/key_chords.rs +++ b/codex-rs/tui/src/app/tests/key_chords.rs @@ -25,7 +25,7 @@ use crossterm::event::KeyEvent; use crossterm::event::KeyModifiers; use pretty_assertions::assert_eq; -async fn chord_app() -> Result<(App, Tui, AppServerSession)> { +async fn chord_app() -> Result<(Box, Tui, AppServerSession)> { let mut app = make_test_app().await; let mut config = TuiKeymap::default(); config.global.open_transcript = Some(KeybindingsSpec::One(KeybindingSpec( diff --git a/codex-rs/tui/src/app/tests/pagination_completion_tests.rs b/codex-rs/tui/src/app/tests/pagination_completion_tests.rs index 424a1a01d8..d3d98496c6 100644 --- a/codex-rs/tui/src/app/tests/pagination_completion_tests.rs +++ b/codex-rs/tui/src/app/tests/pagination_completion_tests.rs @@ -19,7 +19,7 @@ use pretty_assertions::assert_eq; pub(super) async fn completed_history_app( names: &[&str], -) -> Result<(App, tempfile::TempDir, ThreadId)> { +) -> Result<(Box, tempfile::TempDir, ThreadId)> { let mut app = make_test_app().await; let codex_home = tempdir()?; app.config.codex_home = codex_home.path().to_path_buf().abs(); diff --git a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs index 28344b49be..46357d659a 100644 --- a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs +++ b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs @@ -664,7 +664,7 @@ pub(super) fn recorded_params(requests: &RecordedRequests, method: &str) -> Vec< .collect() } -async fn make_history_test_app() -> Result<(App, tempfile::TempDir)> { +async fn make_history_test_app() -> Result<(Box, tempfile::TempDir)> { let mut app = make_test_app().await; let codex_home = tempdir()?; app.config.codex_home = codex_home.path().to_path_buf().abs(); diff --git a/codex-rs/tui/src/app/tests/thread_usage.rs b/codex-rs/tui/src/app/tests/thread_usage.rs index b782d6a564..8bd1900150 100644 --- a/codex-rs/tui/src/app/tests/thread_usage.rs +++ b/codex-rs/tui/src/app/tests/thread_usage.rs @@ -9,8 +9,8 @@ use pretty_assertions::assert_eq; use ratatui::layout::Rect; use ratatui::layout::Size; -async fn app_with_pending_thread_usage() -> Result<(App, AppServerSession, tui::Tui, ThreadId, u64)> -{ +async fn app_with_pending_thread_usage() +-> Result<(Box, AppServerSession, tui::Tui, ThreadId, u64)> { let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; let thread_id = ThreadId::new(); app.chat_widget.handle_thread_session(test_thread_session( diff --git a/codex-rs/tui/src/app/tests/transcript_composer.rs b/codex-rs/tui/src/app/tests/transcript_composer.rs index ecf63a6c21..7fe8788585 100644 --- a/codex-rs/tui/src/app/tests/transcript_composer.rs +++ b/codex-rs/tui/src/app/tests/transcript_composer.rs @@ -307,7 +307,7 @@ async fn transcript_flag_off_preserves_viewer_and_backtracking() -> Result<()> { Ok(()) } -async fn assert_transcript_close_repaints_inline_draft(mut app: App) -> Result<()> { +async fn assert_transcript_close_repaints_inline_draft(mut app: Box) -> Result<()> { let mut app_server = start_config_write_test_app_server(&app).await?; let mut tui = crate::tui::test_support::make_test_tui()?; app.chat_widget.insert_str("EDGE-DRAFT-MUST-SURVIVE"); diff --git a/codex-rs/websocket-client/src/cookie_tests.rs b/codex-rs/websocket-client/src/cookie_tests.rs index 543232f1da..1d48eab20c 100644 --- a/codex-rs/websocket-client/src/cookie_tests.rs +++ b/codex-rs/websocket-client/src/cookie_tests.rs @@ -182,10 +182,10 @@ async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { .connect_with_route( request, WebSocketConfig::default(), - OutboundProxyRoute::Proxy { + std::future::ready(Ok(OutboundProxyRoute::Proxy { url: proxy_url.clone(), no_proxy: None, - }, + })), /*loopback_direct*/ false, ) .await; diff --git a/codex-rs/websocket-client/src/dialer.rs b/codex-rs/websocket-client/src/dialer.rs index 680b144aeb..e705bf891d 100644 --- a/codex-rs/websocket-client/src/dialer.rs +++ b/codex-rs/websocket-client/src/dialer.rs @@ -53,7 +53,7 @@ pub(crate) async fn connect( tls_config.map(Connector::Rustls), ) .await?; - return Ok((ConnectionInner::TransportDefault(stream), response)); + return Ok((ConnectionInner::Left(stream), response)); } OutboundProxyRoute::Direct => None, OutboundProxyRoute::Proxy { @@ -76,7 +76,7 @@ pub(crate) async fn connect( .await { Ok((stream, response)) => { - return Ok((ConnectionInner::TransportDefault(stream), response)); + return Ok((ConnectionInner::Left(stream), response)); } Err(WebSocketError::Url(UrlError::UnsupportedProxyScheme)) => Some(url), Err(error) => return Err(error), @@ -131,7 +131,7 @@ pub(crate) async fn connect( tls_config.map(Connector::Rustls), ) .await?; - Ok((ConnectionInner::Routed(stream), response)) + Ok((ConnectionInner::Right(stream), response)) } #[derive(Debug, PartialEq, Eq)] diff --git a/codex-rs/websocket-client/src/dialer_tests.rs b/codex-rs/websocket-client/src/dialer_tests.rs index 8a8112bce7..e965915ee5 100644 --- a/codex-rs/websocket-client/src/dialer_tests.rs +++ b/codex-rs/websocket-client/src/dialer_tests.rs @@ -198,7 +198,7 @@ async fn direct_route_connects_secure_websocket() { ) .await .expect("direct websocket handshake should succeed"); - drop(WebSocketConnection { inner }); + drop(test_connection(inner)); target_task.await.expect("target task should finish"); } @@ -286,7 +286,7 @@ async fn no_proxy_subprocess_probe() { ) .await .expect("websocket handshake should succeed"); - let mut websocket = WebSocketConnection { inner }; + let mut websocket = test_connection(inner); websocket .send(Message::Text("probe".into())) .await @@ -411,7 +411,7 @@ async fn routed_tcp_connections_only_enable_nodelay_when_requested() { } fn websocket_tcp_nodelay(websocket: &WebSocketConnection) -> bool { - let ConnectionInner::TransportDefault(stream) = &websocket.inner else { + let Some(ConnectionInner::Left(stream)) = &websocket.inner else { panic!("default connector should use Tungstenite's transport"); }; let MaybeTlsStream::Plain(stream) = stream.get_ref() else { @@ -645,7 +645,7 @@ async fn assert_proxy_tunnels_secure_websocket(proxy_tls: bool) { ) .await .expect("proxied websocket handshake should succeed"); - drop(WebSocketConnection { inner }); + drop(test_connection(inner)); target_task.await.expect("target task should finish"); proxy_task.await.expect("proxy task should finish"); @@ -705,3 +705,11 @@ fn test_tls_configs() -> (Arc, TlsAcceptor, CertificateDer<'static certificate, ) } + +fn test_connection(inner: ConnectionInner) -> WebSocketConnection { + let policy = codex_http_client::NetworkPolicy::unmanaged(); + let lease = policy + .acquire(&url::Url::parse("wss://localhost/").unwrap()) + .unwrap(); + WebSocketConnection::new(inner, lease) +} diff --git a/codex-rs/websocket-client/src/lib.rs b/codex-rs/websocket-client/src/lib.rs index 19cd8ccc96..28dbe21484 100644 --- a/codex-rs/websocket-client/src/lib.rs +++ b/codex-rs/websocket-client/src/lib.rs @@ -12,11 +12,15 @@ use std::task::Poll; use codex_http_client::BuildCustomCaTransportError; use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPermit; +use codex_http_client::NetworkPolicyDenied; use codex_http_client::OutboundProxyRoute; use codex_http_client::build_rustls_client_config_with_custom_ca; use futures::FutureExt; use futures::Sink; use futures::Stream; +use futures::future::BoxFuture; +use futures::future::Either; use rustls::ClientConfig; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; @@ -31,6 +35,18 @@ use tokio_tungstenite::tungstenite::http::Uri; use tokio_tungstenite::tungstenite::http::header::COOKIE; use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; +/// Recovers a deterministic policy denial from Tungstenite's I/O error boundary. +/// Callers must preserve this distinction when deciding whether to retry. +pub fn network_policy_denial(error: &WebSocketError) -> Option { + match error { + WebSocketError::Io(error) => error + .get_ref()? + .downcast_ref::() + .copied(), + _ => None, + } +} + /// Connects WebSockets using the outbound proxy policy resolved by application configuration. /// /// Construct this from the effective [`HttpClientFactory`] rather than selecting proxy behavior at @@ -103,19 +119,15 @@ impl WebSocketConnector { request: Request, config: WebSocketConfig, ) -> Result<(WebSocketConnection, Response), WebSocketError> { - let proxy_route = self + let route = self .http_client_factory - .resolve_proxy_route_async(request.uri().to_string()) - .await - .map_err(WebSocketError::Io)?; - self.connect_with_route(request, config, proxy_route, /*loopback_direct*/ false) + .resolve_proxy_route_async(request.uri().to_string()); + self.connect_with_route(request, config, route, /*loopback_direct*/ false) .await } /// Connects to a validated loopback destination without consulting proxy settings. - /// - /// This is limited to loopback destinations because bypassing configured proxy policy is - /// only safe for local connections. + /// Application destination policy still applies before dialing. pub async fn connect_loopback_direct( &self, request: Request, @@ -130,7 +142,7 @@ impl WebSocketConnector { self.connect_with_route( request, config, - OutboundProxyRoute::Direct, + std::future::ready(Ok(OutboundProxyRoute::Direct)), /*loopback_direct*/ true, ) .await @@ -140,37 +152,52 @@ impl WebSocketConnector { &self, mut request: Request, config: WebSocketConfig, - proxy_route: OutboundProxyRoute, + route: impl std::future::Future>, loopback_direct: bool, ) -> Result<(WebSocketConnection, Response), WebSocketError> { let uri = request.uri().clone(); - if !request.headers().contains_key(COOKIE) - && let Some(cookies) = self.http_client_factory.chatgpt_cookie_header(&uri) - { - request.headers_mut().insert(COOKIE, cookies); - } - let result = dialer::connect( - request, - config, - self.tls_config.clone(), - proxy_route, - self.tcp_nodelay, - loopback_direct, - ) - .boxed() - .await; - // Like HTTP responses, rejected upgrades can also refresh infrastructure cookies. - match &result { - Ok((_, response)) => self - .http_client_factory - .store_chatgpt_response_cookies(&uri, response.headers()), - Err(WebSocketError::Http(response)) => self - .http_client_factory - .store_chatgpt_response_cookies(&uri, response.headers()), - Err(_) => {} - } - let (inner, response) = result?; - Ok((WebSocketConnection { inner }, response)) + let url = url::Url::parse(&uri.to_string()).map_err(|_| { + WebSocketError::Io(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid WebSocket URL", + )) + })?; + let permit = self + .http_client_factory + .network_policy() + .acquire(&url) + .map_err(policy_error)?; + let connect = async { + if !request.headers().contains_key(COOKIE) + && let Some(cookies) = self.http_client_factory.chatgpt_cookie_header(&uri) + { + request.headers_mut().insert(COOKIE, cookies); + } + let proxy_route = route.await.map_err(WebSocketError::Io)?; + let result = dialer::connect( + request, + config, + self.tls_config.clone(), + proxy_route, + self.tcp_nodelay, + loopback_direct, + ) + .boxed() + .await; + // Like HTTP responses, rejected upgrades can also refresh infrastructure cookies. + match &result { + Ok((_, response)) => self + .http_client_factory + .store_chatgpt_response_cookies(&uri, response.headers()), + Err(WebSocketError::Http(response)) => self + .http_client_factory + .store_chatgpt_response_cookies(&uri, response.headers()), + Err(_) => {} + } + result + }; + let (inner, response) = permit.run(connect).await.map_err(policy_error)??; + Ok((WebSocketConnection::new(inner, permit), response)) } } @@ -191,19 +218,85 @@ fn is_loopback_destination(uri: &Uri) -> bool { /// An established WebSocket independent of its direct, proxy, and TLS transport layers. /// /// This implements [`Stream`] and [`Sink`] so protocol clients can process Tungstenite messages -/// without knowing which concrete network stream route selection produced. +/// without knowing which concrete network stream route selection produced. Independent revocation +/// waits preserve both task wakers when the stream and sink are split. pub struct WebSocketConnection { - inner: ConnectionInner, + inner: Option, + permit: NetworkPermit, + read_revoked: BoxFuture<'static, ()>, + write_revoked: BoxFuture<'static, ()>, + read_terminated: bool, +} + +enum ConnectionDirection { + Read, + Write, +} + +impl std::fmt::Debug for WebSocketConnection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("WebSocketConnection") + .finish_non_exhaustive() + } +} + +impl WebSocketConnection { + fn new(inner: ConnectionInner, permit: NetworkPermit) -> Self { + let read_permit = permit.clone(); + let write_permit = permit.clone(); + Self { + inner: Some(inner), + permit, + read_revoked: async move { read_permit.revoked().await }.boxed(), + write_revoked: async move { write_permit.revoked().await }.boxed(), + read_terminated: false, + } + } + + fn poll_policy( + &mut self, + context: &mut Context<'_>, + direction: ConnectionDirection, + ) -> Result<(), WebSocketError> { + let revoked = match direction { + ConnectionDirection::Read => &mut self.read_revoked, + ConnectionDirection::Write => &mut self.write_revoked, + }; + if self.permit.check().is_err() || revoked.poll_unpin(context).is_ready() { + self.inner = None; + return Err(policy_error(NetworkPolicyDenied::Revoked)); + } + Ok(()) + } +} + +fn policy_error(denied: NetworkPolicyDenied) -> WebSocketError { + WebSocketError::Io(std::io::Error::new( + std::io::ErrorKind::PermissionDenied, + denied, + )) } impl Stream for WebSocketConnection { type Item = Result; fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { - match &mut self.get_mut().inner { - ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_next(context), - ConnectionInner::Routed(stream) => Pin::new(stream).poll_next(context), + let connection = self.get_mut(); + if connection.read_terminated { + return Poll::Ready(None); } + if let Err(error) = connection.poll_policy(context, ConnectionDirection::Read) { + connection.read_terminated = true; + return Poll::Ready(Some(Err(error))); + } + let next = match &mut connection.inner { + Some(stream) => Pin::new(stream).poll_next(context), + None => Poll::Ready(None), + }; + if matches!(&next, Poll::Ready(None)) { + connection.read_terminated = true; + } + next } } @@ -214,16 +307,23 @@ impl Sink for WebSocketConnection { self: Pin<&mut Self>, context: &mut Context<'_>, ) -> Poll> { - match &mut self.get_mut().inner { - ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_ready(context), - ConnectionInner::Routed(stream) => Pin::new(stream).poll_ready(context), + let connection = self.get_mut(); + connection.poll_policy(context, ConnectionDirection::Write)?; + match &mut connection.inner { + Some(stream) => Pin::new(stream).poll_ready(context), + None => Poll::Ready(Err(WebSocketError::ConnectionClosed)), } } fn start_send(self: Pin<&mut Self>, message: Message) -> Result<(), Self::Error> { - match &mut self.get_mut().inner { - ConnectionInner::TransportDefault(stream) => Pin::new(stream).start_send(message), - ConnectionInner::Routed(stream) => Pin::new(stream).start_send(message), + let connection = self.get_mut(); + if let Err(error) = connection.permit.check() { + connection.inner = None; + return Err(policy_error(error)); + } + match &mut connection.inner { + Some(stream) => Pin::new(stream).start_send(message), + None => Err(WebSocketError::ConnectionClosed), } } @@ -231,9 +331,11 @@ impl Sink for WebSocketConnection { self: Pin<&mut Self>, context: &mut Context<'_>, ) -> Poll> { - match &mut self.get_mut().inner { - ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_flush(context), - ConnectionInner::Routed(stream) => Pin::new(stream).poll_flush(context), + let connection = self.get_mut(); + connection.poll_policy(context, ConnectionDirection::Write)?; + match &mut connection.inner { + Some(stream) => Pin::new(stream).poll_flush(context), + None => Poll::Ready(Err(WebSocketError::ConnectionClosed)), } } @@ -241,17 +343,19 @@ impl Sink for WebSocketConnection { self: Pin<&mut Self>, context: &mut Context<'_>, ) -> Poll> { - match &mut self.get_mut().inner { - ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_close(context), - ConnectionInner::Routed(stream) => Pin::new(stream).poll_close(context), + let connection = self.get_mut(); + connection.poll_policy(context, ConnectionDirection::Write)?; + match &mut connection.inner { + Some(stream) => Pin::new(stream).poll_close(context), + None => Poll::Ready(Ok(())), } } } -pub(crate) enum ConnectionInner { - TransportDefault(TungsteniteStream>), - Routed(TungsteniteStream>>), -} +pub(crate) type ConnectionInner = Either< + TungsteniteStream>, + TungsteniteStream>>, +>; /// Async network I/O carried through optional proxy and target TLS handshakes. pub(crate) trait AsyncIo: AsyncRead + AsyncWrite + Send + Unpin {} @@ -265,3 +369,7 @@ mod tests; #[cfg(test)] #[path = "cookie_tests.rs"] mod cookie_tests; + +#[cfg(test)] +#[path = "network_policy_tests.rs"] +mod network_policy_tests; diff --git a/codex-rs/websocket-client/src/network_policy_tests.rs b/codex-rs/websocket-client/src/network_policy_tests.rs new file mode 100644 index 0000000000..8be65d2772 --- /dev/null +++ b/codex-rs/websocket-client/src/network_policy_tests.rs @@ -0,0 +1,54 @@ +//! Split readers keep their revocation wakeup and report denial even when writers fail first. + +use super::*; +use codex_http_client::DestinationPolicy; +use codex_http_client::NetworkPolicyController; +use codex_http_client::OutboundProxyPolicy::ReqwestDefault; +use futures::SinkExt; +use futures::StreamExt; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; + +struct WakeFlag(AtomicBool); + +impl futures::task::ArcWake for WakeFlag { + fn wake_by_ref(waker: &Arc) { + waker.0.store(true, Ordering::SeqCst); + } +} + +#[tokio::test] +async fn revocation_wakes_split_reader_and_survives_writer_observing_it_first() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let _ = websocket.next().await; + }); + let controller = NetworkPolicyController::default(); + let policy = controller.policy(); + assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted)); + let factory = HttpClientFactory::new(ReqwestDefault).with_network_policy(policy.clone()); + let connector = WebSocketConnector::new(&factory).unwrap(); + let request = url.into_client_request().unwrap(); + let (connection, _) = connector + .connect(request, WebSocketConfig::default()) + .await + .unwrap(); + let (mut writer, mut reader) = connection.split(); + let wake = Arc::new(WakeFlag(AtomicBool::new(false))); + let waker = futures::task::waker(Arc::clone(&wake)); + let mut context = Context::from_waker(&waker); + assert!(Pin::new(&mut reader).poll_next(&mut context).is_pending()); + writer.flush().await.unwrap(); + policy.invalidate(); + assert!(wake.0.load(Ordering::SeqCst)); + let error = writer.flush().await.unwrap_err(); + assert!(network_policy_denial(&error).is_some()); + let error = reader.next().await.unwrap().unwrap_err(); + assert!(network_policy_denial(&error).is_some()); + assert!(reader.next().await.is_none()); + server.abort(); +}