From 5bcbe797f55d5776b12eb22f99684610b13b16a7 Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Tue, 2 Jun 2026 20:57:54 +0000 Subject: [PATCH] [codex-api] cover native state transport bypasses [ci changed_files] --- .../src/endpoint/responses_websocket.rs | 263 +++++++++++++++--- codex-rs/codex-api/src/files.rs | 121 ++++++-- codex-rs/codex-client/src/default_client.rs | 47 +++- 3 files changed, 375 insertions(+), 56 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 022f2fd0b8..abbe9e44dc 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -160,6 +160,41 @@ const OPENAI_MODEL_HEADER: &str = "openai-model"; const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE: &str = "websocket_connection_limit_reached"; const WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE: &str = "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."; +#[derive(Clone)] +struct WebsocketAuthContext { + auth: SharedAuthProvider, + request_url: String, + request_headers: HeaderMap, +} + +impl WebsocketAuthContext { + fn new(auth: SharedAuthProvider, request_url: String, request_headers: HeaderMap) -> Self { + Self { + auth, + request_url, + request_headers, + } + } + + fn observe_response_headers(&self, response_headers: &HeaderMap) { + self.auth.observe_response_headers( + &self.request_url, + &self.request_headers, + response_headers, + ); + } + + fn observe_error_headers(&self, error: &ApiError) { + if let ApiError::Transport(TransportError::Http { + headers: Some(response_headers), + .. + }) = error + { + self.observe_response_headers(response_headers); + } + } +} + pub struct ResponsesWebsocketConnection { stream: Arc>>, // TODO (pakrym): is this the right place for timeout? @@ -167,6 +202,7 @@ pub struct ResponsesWebsocketConnection { server_reasoning_included: bool, models_etag: Option, server_model: Option, + auth_context: WebsocketAuthContext, telemetry: Option>, } @@ -178,6 +214,7 @@ impl std::fmt::Debug for ResponsesWebsocketConnection { .field("server_reasoning_included", &self.server_reasoning_included) .field("models_etag", &self.models_etag) .field("server_model", &self.server_model) + .field("auth_context", &"") .field("telemetry", &self.telemetry.as_ref().map(|_| "")) .finish() } @@ -190,6 +227,7 @@ impl ResponsesWebsocketConnection { server_reasoning_included: bool, models_etag: Option, server_model: Option, + auth_context: WebsocketAuthContext, telemetry: Option>, ) -> Self { Self { @@ -198,6 +236,7 @@ impl ResponsesWebsocketConnection { server_reasoning_included, models_etag, server_model, + auth_context, telemetry, } } @@ -258,6 +297,7 @@ impl ResponsesWebsocketConnection { let server_reasoning_included = self.server_reasoning_included; let models_etag = self.models_etag.clone(); let server_model = self.server_model.clone(); + let auth_context = self.auth_context.clone(); let telemetry = self.telemetry.clone(); let request_body = serde_json::to_value(&request).map_err(|err| { ApiError::Stream(format!("failed to encode websocket request: {err}")) @@ -297,6 +337,7 @@ impl ResponsesWebsocketConnection { tx_event.clone(), request_body, idle_timeout, + &auth_context, telemetry, connection_reused, ) @@ -383,7 +424,8 @@ impl ResponsesWebsocketClient { merge_request_headers(&self.provider.headers, extra_headers, default_headers); self.auth .add_auth_headers_for_url(&request_url, &mut headers); - let request_headers = headers.clone(); + let auth_context = + WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone()); let ( stream, @@ -395,21 +437,16 @@ impl ResponsesWebsocketClient { ) = connect_websocket(ws_url, headers, turn_state.clone()) .await .inspect_err(|error| { - observe_websocket_auth_error_headers( - self.auth.as_ref(), - &request_url, - &request_headers, - error, - ); + auth_context.observe_error_headers(error); })?; - self.auth - .observe_response_headers(&request_url, &request_headers, &response_headers); + auth_context.observe_response_headers(&response_headers); Ok(ResponsesWebsocketConnection::new( stream, self.provider.stream_idle_timeout, server_reasoning_included, models_etag, server_model, + auth_context, telemetry, )) } @@ -437,21 +474,16 @@ impl ResponsesWebsocketClient { merge_request_headers(&self.provider.headers, extra_headers, default_headers); self.auth .add_auth_headers_for_url(&request_url, &mut headers); - let request_headers = headers.clone(); + let auth_context = + WebsocketAuthContext::new(self.auth.clone(), request_url, headers.clone()); let (mut stream, status, reasoning_included, models_etag, server_model, response_headers) = connect_websocket(ws_url.clone(), headers, /*turn_state*/ None) .await .inspect_err(|error| { - observe_websocket_auth_error_headers( - self.auth.as_ref(), - &request_url, - &request_headers, - error, - ); + auth_context.observe_error_headers(error); })?; - self.auth - .observe_response_headers(&request_url, &request_headers, &response_headers); + auth_context.observe_response_headers(&response_headers); let immediate_close = tokio::time::timeout(immediate_close_timeout, stream.next()) .await .ok() @@ -473,21 +505,6 @@ impl ResponsesWebsocketClient { } } -fn observe_websocket_auth_error_headers( - auth: &dyn crate::auth::AuthProvider, - request_url: &str, - request_headers: &HeaderMap, - error: &ApiError, -) { - if let ApiError::Transport(TransportError::Http { - headers: Some(response_headers), - .. - }) = error - { - auth.observe_response_headers(request_url, request_headers, response_headers); - } -} - fn immediate_close_from_message(message: Message) -> Option { let Message::Close(frame) = message else { return None; @@ -721,6 +738,7 @@ async fn run_websocket_response_stream( tx_event: mpsc::Sender>, request_body: Value, idle_timeout: Duration, + auth_context: &WebsocketAuthContext, telemetry: Option>, connection_reused: bool, ) -> Result<(), ApiError> { @@ -764,6 +782,7 @@ async fn run_websocket_response_stream( && let Some(error) = map_wrapped_websocket_error_event(wrapped_error, text.to_string()) { + auth_context.observe_error_headers(&error); return Err(error); } @@ -873,8 +892,186 @@ async fn send_websocket_request( #[cfg(test)] mod tests { use super::*; + use crate::auth::AuthProvider; + use crate::provider::RetryConfig; use pretty_assertions::assert_eq; use serde_json::json; + use std::sync::Mutex as StdMutex; + use tokio::net::TcpListener; + use tokio_tungstenite::accept_hdr_async_with_config; + + const SENT_STATE: &str = "ois1.sent.nonce.ciphertext"; + const HANDSHAKE_STATE: &str = "ois1.handshake.nonce.ciphertext"; + const WRAPPED_ERROR_STATE: &str = "ois1.error.nonce.ciphertext"; + + #[derive(Default)] + struct RecordingAuthProvider { + observed_updates: StdMutex>, + } + + impl AuthProvider for RecordingAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) { + headers.insert("x-oai-is", HeaderValue::from_static(SENT_STATE)); + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, + ) { + let Some(sent_state) = request_headers + .get("x-oai-is") + .and_then(|value| value.to_str().ok()) + else { + return; + }; + let Some(update_state) = response_headers + .get("x-oai-is-update") + .and_then(|value| value.to_str().ok()) + else { + return; + }; + self.observed_updates + .lock() + .expect("observed updates lock should not be poisoned") + .push(( + request_url.to_string(), + sent_state.to_string(), + update_state.to_string(), + )); + } + } + + #[tokio::test] + async fn connect_attaches_auth_and_observes_handshake_and_wrapped_error_headers() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind websocket server"); + let address = listener.local_addr().expect("websocket server address"); + let observed_request_state = Arc::new(StdMutex::new(None)); + let server_request_state = Arc::clone(&observed_request_state); + let server_task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept websocket client"); + let mut stream = accept_hdr_async_with_config( + stream, + move |request: &tokio_tungstenite::tungstenite::handshake::server::Request, + mut response: tokio_tungstenite::tungstenite::handshake::server::Response| { + *server_request_state + .lock() + .expect("request state lock should not be poisoned") = request + .headers() + .get("x-oai-is") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); + response.headers_mut().insert( + "x-oai-is-update", + HeaderValue::from_static(HANDSHAKE_STATE), + ); + Ok(response) + }, + Some(websocket_config()), + ) + .await + .expect("complete websocket handshake"); + let message = stream + .next() + .await + .expect("receive websocket request") + .expect("read websocket request"); + assert!(matches!(message, Message::Text(_))); + stream + .send(Message::Text( + json!({ + "type": "error", + "status": 401, + "headers": { + "x-oai-is-update": WRAPPED_ERROR_STATE, + }, + }) + .to_string() + .into(), + )) + .await + .expect("send wrapped websocket error"); + }); + + let auth = Arc::new(RecordingAuthProvider::default()); + let request_url = format!("ws://{address}/backend-api/codex/responses"); + let client = ResponsesWebsocketClient::new(websocket_provider(address), auth.clone()); + let connection = client + .connect( + HeaderMap::new(), + HeaderMap::new(), + /*turn_state*/ None, + /*telemetry*/ None, + ) + .await + .expect("connect websocket client"); + let mut response_stream = connection + .stream_request( + ResponsesWsRequest::ResponseProcessed(ResponseProcessedWsRequest { + response_id: "response-id".to_string(), + }), + /*connection_reused*/ false, + ) + .await + .expect("start websocket response stream"); + let error = response_stream + .next() + .await + .expect("receive websocket response") + .expect_err("wrapped error should fail the response stream"); + + server_task.await.expect("websocket server task"); + let ApiError::Transport(TransportError::Http { status, .. }) = error else { + panic!("expected transport HTTP error"); + }; + assert_eq!(status, StatusCode::UNAUTHORIZED); + assert_eq!( + *observed_request_state + .lock() + .expect("request state lock should not be poisoned"), + Some(SENT_STATE.to_string()) + ); + assert_eq!( + *auth + .observed_updates + .lock() + .expect("observed updates lock should not be poisoned"), + vec![ + ( + request_url.clone(), + SENT_STATE.to_string(), + HANDSHAKE_STATE.to_string(), + ), + ( + request_url, + SENT_STATE.to_string(), + WRAPPED_ERROR_STATE.to_string(), + ), + ] + ); + } + + fn websocket_provider(address: std::net::SocketAddr) -> Provider { + Provider { + name: "test".to_string(), + base_url: format!("http://{address}/backend-api/codex"), + query_params: None, + headers: HeaderMap::new(), + retry: RetryConfig { + max_attempts: 1, + base_delay: Duration::from_millis(1), + retry_429: false, + retry_5xx: false, + retry_transport: false, + }, + stream_idle_timeout: Duration::from_secs(1), + } + } #[test] fn websocket_config_enables_permessage_deflate() { diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index d1e2840066..76e8d15c85 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -129,7 +129,9 @@ pub async fn upload_local_file( .unwrap_or("file") .to_string(); let create_url = format!("{}/files", base_url.trim_end_matches('/')); - let create_response = authorized_request(auth, reqwest::Method::POST, &create_url) + let (create_request, create_request_headers) = + authorized_request(auth, reqwest::Method::POST, &create_url); + let create_response = create_request .json(&serde_json::json!({ "file_name": file_name, "file_size": metadata.len(), @@ -141,6 +143,11 @@ pub async fn upload_local_file( url: create_url.clone(), source, })?; + auth.observe_response_headers( + &create_url, + &create_request_headers, + create_response.headers(), + ); let create_status = create_response.status(); let create_body = create_response.text().await.unwrap_or_default(); if !create_status.is_success() { @@ -191,7 +198,9 @@ pub async fn upload_local_file( ); let finalize_started_at = Instant::now(); loop { - let finalize_response = authorized_request(auth, reqwest::Method::POST, &finalize_url) + let (finalize_request, finalize_request_headers) = + authorized_request(auth, reqwest::Method::POST, &finalize_url); + let finalize_response = finalize_request .json(&serde_json::json!({})) .send() .await @@ -199,6 +208,11 @@ pub async fn upload_local_file( url: finalize_url.clone(), source, })?; + auth.observe_response_headers( + &finalize_url, + &finalize_request_headers, + finalize_response.headers(), + ); let finalize_status = finalize_response.status(); let finalize_body = finalize_response.text().await.unwrap_or_default(); if !finalize_status.is_success() { @@ -255,15 +269,18 @@ fn authorized_request( auth: &dyn AuthProvider, method: reqwest::Method, url: &str, -) -> reqwest::RequestBuilder { +) -> (reqwest::RequestBuilder, http::HeaderMap) { let mut headers = http::HeaderMap::new(); - auth.add_auth_headers(&mut headers); + auth.add_auth_headers_for_url(url, &mut headers); let client = build_reqwest_client(); - client - .request(method, url) - .timeout(OPENAI_FILE_REQUEST_TIMEOUT) - .headers(headers) + ( + client + .request(method, url) + .timeout(OPENAI_FILE_REQUEST_TIMEOUT) + .headers(headers.clone()), + headers, + ) } fn build_reqwest_client() -> reqwest::Client { @@ -279,6 +296,7 @@ mod tests { use pretty_assertions::assert_eq; use reqwest::header::HeaderValue; use std::sync::Arc; + use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use tempfile::TempDir; @@ -291,8 +309,16 @@ mod tests { use wiremock::matchers::method; use wiremock::matchers::path; - #[derive(Clone, Copy)] - struct ChatGptTestAuth; + const SENT_STATE: &str = "ois1.sent.nonce.ciphertext"; + const CREATED_STATE: &str = "ois1.created.nonce.ciphertext"; + const RETRY_STATE: &str = "ois1.retry.nonce.ciphertext"; + const FINALIZED_STATE: &str = "ois1.finalized.nonce.ciphertext"; + type ObservedUpdate = (String, Option); + + #[derive(Clone, Default)] + struct ChatGptTestAuth { + observed_updates: Arc>>, + } impl AuthProvider for ChatGptTestAuth { fn add_auth_headers(&self, headers: &mut reqwest::header::HeaderMap) { @@ -302,10 +328,33 @@ mod tests { ); headers.insert("ChatGPT-Account-ID", HeaderValue::from_static("account_id")); } + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut http::HeaderMap) { + self.add_auth_headers(headers); + headers.insert("x-oai-is", HeaderValue::from_static(SENT_STATE)); + } + + fn observe_response_headers( + &self, + request_url: &str, + _request_headers: &http::HeaderMap, + response_headers: &http::HeaderMap, + ) { + self.observed_updates + .lock() + .expect("observed updates lock should not be poisoned") + .push(( + request_url.to_string(), + response_headers + .get("x-oai-is-update") + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string), + )); + } } fn chatgpt_auth() -> ChatGptTestAuth { - ChatGptTestAuth + ChatGptTestAuth::default() } fn base_url_for(server: &MockServer) -> String { @@ -318,6 +367,7 @@ mod tests { Mock::given(method("POST")) .and(path("/backend-api/files")) .and(header("chatgpt-account-id", "account_id")) + .and(header("x-oai-is", SENT_STATE)) .and(body_json(serde_json::json!({ "file_name": "hello.txt", "file_size": 5, @@ -325,6 +375,7 @@ mod tests { }))) .respond_with( ResponseTemplate::new(200) + .insert_header("x-oai-is-update", CREATED_STATE) .set_body_json(serde_json::json!({"file_id": "file_123", "upload_url": format!("{}/upload/file_123", server.uri())})), ) .mount(&server) @@ -340,20 +391,25 @@ mod tests { let download_url = format!("{}/download/file_123", server.uri()); Mock::given(method("POST")) .and(path("/backend-api/files/file_123/uploaded")) + .and(header("x-oai-is", SENT_STATE)) .respond_with(move |_request: &Request| { if finalize_attempts_responder.fetch_add(1, Ordering::SeqCst) == 0 { - return ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "status": "retry" - })); + return ResponseTemplate::new(200) + .insert_header("x-oai-is-update", RETRY_STATE) + .set_body_json(serde_json::json!({ + "status": "retry" + })); } - ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "status": "success", - "download_url": download_url, - "file_name": "hello.txt", - "mime_type": "text/plain", - "file_size_bytes": 5 - })) + ResponseTemplate::new(200) + .insert_header("x-oai-is-update", FINALIZED_STATE) + .set_body_json(serde_json::json!({ + "status": "success", + "download_url": download_url, + "file_name": "hello.txt", + "mime_type": "text/plain", + "file_size_bytes": 5 + })) }) .mount(&server) .await; @@ -363,7 +419,8 @@ mod tests { let path = dir.path().join("hello.txt"); tokio::fs::write(&path, b"hello").await.expect("write file"); - let uploaded = upload_local_file(&base_url, &chatgpt_auth(), &path) + let auth = chatgpt_auth(); + let uploaded = upload_local_file(&base_url, &auth, &path) .await .expect("upload succeeds"); @@ -376,5 +433,25 @@ mod tests { assert_eq!(uploaded.file_name, "hello.txt"); assert_eq!(uploaded.mime_type, Some("text/plain".to_string())); assert_eq!(finalize_attempts.load(Ordering::SeqCst), 2); + assert_eq!( + *auth + .observed_updates + .lock() + .expect("observed updates lock should not be poisoned"), + vec![ + ( + format!("{}/files", base_url.trim_end_matches('/')), + Some(CREATED_STATE.to_string()), + ), + ( + format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')), + Some(RETRY_STATE.to_string()), + ), + ( + format!("{}/files/file_123/uploaded", base_url.trim_end_matches('/')), + Some(FINALIZED_STATE.to_string()), + ), + ] + ); } } diff --git a/codex-rs/codex-client/src/default_client.rs b/codex-rs/codex-client/src/default_client.rs index 56b3ce4b16..a57b1dc338 100644 --- a/codex-rs/codex-client/src/default_client.rs +++ b/codex-rs/codex-client/src/default_client.rs @@ -1,3 +1,5 @@ +use crate::INTEGRITY_STATE_HEADER_NAME; +use crate::INTEGRITY_STATE_UPDATE_HEADER_NAME; use http::Error as HttpError; use http::HeaderMap; use http::HeaderName; @@ -115,11 +117,12 @@ impl CodexRequestBuilder { match self.builder.headers(headers).send().await { Ok(response) => { + let headers = redacted_headers_for_logging(response.headers()); tracing::debug!( method = %self.method, url = %self.url, status = %response.status(), - headers = ?response.headers(), + headers = ?headers, version = ?response.version(), "Request completed" ); @@ -141,6 +144,19 @@ impl CodexRequestBuilder { } } +fn redacted_headers_for_logging(headers: &HeaderMap) -> HeaderMap { + let mut headers = headers.clone(); + for name in [ + INTEGRITY_STATE_HEADER_NAME, + INTEGRITY_STATE_UPDATE_HEADER_NAME, + ] { + if headers.contains_key(name) { + headers.insert(name, HeaderValue::from_static("")); + } + } + headers +} + struct HeaderMapInjector<'a>(&'a mut HeaderMap); impl<'a> Injector for HeaderMapInjector<'a> { @@ -174,10 +190,39 @@ mod tests { use opentelemetry::trace::TracerProvider; use opentelemetry_sdk::propagation::TraceContextPropagator; use opentelemetry_sdk::trace::SdkTracerProvider; + use pretty_assertions::assert_eq; use tracing::trace_span; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::util::SubscriberInitExt; + #[test] + fn redacted_headers_for_logging_hides_integrity_state() { + let mut headers = HeaderMap::new(); + headers.insert( + INTEGRITY_STATE_HEADER_NAME, + HeaderValue::from_static("ois1.request.nonce.ciphertext"), + ); + headers.insert( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + HeaderValue::from_static("ois1.response.nonce.ciphertext"), + ); + headers.insert("x-request-id", HeaderValue::from_static("request-id")); + + let redacted = redacted_headers_for_logging(&headers); + + let mut expected = HeaderMap::new(); + expected.insert( + INTEGRITY_STATE_HEADER_NAME, + HeaderValue::from_static(""), + ); + expected.insert( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + HeaderValue::from_static(""), + ); + expected.insert("x-request-id", HeaderValue::from_static("request-id")); + assert_eq!(redacted, expected); + } + #[test] fn inject_trace_headers_uses_current_span_context() { global::set_text_map_propagator(TraceContextPropagator::new());