diff --git a/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs b/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs index 69124e0315..847bad77c5 100644 --- a/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs +++ b/codex-rs/http-client/src/chatgpt_cloudflare_cookies.rs @@ -1,10 +1,17 @@ +//! ChatGPT cookies shared by HTTP and WebSocket transports. Only infrastructure cookies may be +//! stored globally; configured cookies remain scoped to their factory. + use std::sync::Arc; use std::sync::LazyLock; +use http::HeaderMap; +use http::HeaderValue; +use http::Uri; +use http::header::SET_COOKIE; use reqwest::cookie::CookieStore; use reqwest::cookie::Jar; -use reqwest::header::HeaderValue; +use crate::HttpClientFactory; use crate::chatgpt_hosts::is_allowed_chatgpt_host; // WARNING: this HTTP cookie store is process-global and may be shared across auth contexts. @@ -96,6 +103,38 @@ impl CookieStore for ChatGptCookieStore { } } +impl HttpClientFactory { + /// Returns cookies for a ChatGPT HTTPS or WSS request, using the same store as HTTP clients. + /// Explicit request Cookie headers should take precedence over this value. + pub fn chatgpt_cookie_header(&self, uri: &Uri) -> Option { + let url = chatgpt_cookie_url(uri)?; + let mut cookies = match self.chatgpt_cookie_store() { + Some(store) => store.cookies(&url), + None => SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE.cookies(&url), + }?; + cookies.set_sensitive(true); + Some(cookies) + } + + /// Retains only allowlisted infrastructure cookies from a ChatGPT HTTPS or WSS response. + /// Account and session cookies are never added to the shared store. + pub fn store_chatgpt_response_cookies(&self, uri: &Uri, headers: &HeaderMap) { + if let Some(url) = chatgpt_cookie_url(uri) { + SHARED_CHATGPT_CLOUDFLARE_COOKIE_STORE + .set_cookies(&mut headers.get_all(SET_COOKIE).iter(), &url); + } + } +} + +fn chatgpt_cookie_url(uri: &Uri) -> Option { + let mut url = reqwest::Url::parse(&uri.to_string()).ok()?; + // A secure WebSocket handshake has the same cookie scope as HTTPS. + if url.scheme() == "wss" { + url.set_scheme("https").ok()?; + } + is_chatgpt_cookie_url(&url).then_some(url) +} + /// Adds the process-local ChatGPT infrastructure cookie jar used by Codex HTTP clients. /// /// WARNING: this jar is global within the process. It is only acceptable because it hardcodes a @@ -177,9 +216,52 @@ fn is_allowed_cloudflare_cookie_name(name: &str) -> bool { #[cfg(test)] mod tests { use super::*; + use crate::OutboundProxyPolicy; use pretty_assertions::assert_eq; use reqwest::cookie::CookieStore; + #[test] + fn http_and_websocket_cookies_share_the_factory_store() { + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_chatgpt_cookies([HeaderValue::from_static("configured=owner")]); + let store = factory.chatgpt_cookie_store().unwrap(); + let https = reqwest::Url::parse("https://http-websocket-cookies.chatgpt.com/api").unwrap(); + let wss: Uri = "wss://http-websocket-cookies.chatgpt.com/api" + .parse() + .unwrap(); + let cookie = HeaderValue::from_static("__oailb=from-http; Path=/; Secure"); + store.set_cookies(&mut std::iter::once(&cookie), &https); + + let header = factory.chatgpt_cookie_header(&wss).unwrap(); + assert_eq!(header, "__oailb=from-http; configured=owner"); + assert!(header.is_sensitive()); + + let mut response_headers = HeaderMap::new(); + response_headers.insert( + SET_COOKIE, + HeaderValue::from_static("__oailb=from-wss; Path=/; Secure"), + ); + factory.store_chatgpt_response_cookies(&wss, &response_headers); + assert_eq!( + store.cookies(&https), + Some(HeaderValue::from_static( + "__oailb=from-wss; configured=owner" + )) + ); + assert_eq!( + factory.chatgpt_cookie_header( + &"ws://http-websocket-cookies.chatgpt.com/api" + .parse() + .unwrap() + ), + None + ); + assert_eq!( + factory.chatgpt_cookie_header(&"wss://api.openai.com/api".parse().unwrap()), + None + ); + } + #[test] fn additional_cookies_use_current_path_scoped_cloudflare_cookies() { let cloudflare = Arc::new(ChatGptCloudflareCookieStore::default()); diff --git a/codex-rs/websocket-client/src/cookie_tests.rs b/codex-rs/websocket-client/src/cookie_tests.rs new file mode 100644 index 0000000000..543232f1da --- /dev/null +++ b/codex-rs/websocket-client/src/cookie_tests.rs @@ -0,0 +1,221 @@ +//! Real WebSocket handshakes exercising shared routing cookies and their transport boundaries. + +use std::sync::Arc; + +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_http_client::OutboundProxyRoute; +use codex_utils_rustls_provider::ensure_rustls_crypto_provider; +use pretty_assertions::assert_eq; +use rcgen::CertifiedKey; +use rcgen::generate_simple_self_signed; +use rustls::ClientConfig; +use rustls::RootCertStore; +use rustls::ServerConfig; +use rustls::pki_types::PrivateKeyDer; +use rustls::pki_types::PrivatePkcs8KeyDer; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio_rustls::TlsAcceptor; +use tokio_tungstenite::accept_hdr_async; +use tokio_tungstenite::tungstenite::Error as WebSocketError; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::handshake::server::Request; +use tokio_tungstenite::tungstenite::handshake::server::Response; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::StatusCode; +use tokio_tungstenite::tungstenite::http::header::COOKIE; +use tokio_tungstenite::tungstenite::http::header::SET_COOKIE; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; + +use crate::AsyncIo; +use crate::TcpNodelay; +use crate::WebSocketConnector; + +#[tokio::test] +async fn websocket_handshakes_share_routing_cookies_and_respect_cookie_scope() { + ensure_rustls_crypto_provider(); + let host = "websocket-cookie-test.chatgpt.com"; + let other_host = "other-websocket-cookie-test.chatgpt.com"; + let CertifiedKey { cert, signing_key } = generate_simple_self_signed(vec![ + host.to_string(), + other_host.to_string(), + "api.openai.com".to_string(), + ]) + .unwrap(); + let server_config = ServerConfig::builder() + .with_no_client_auth() + .with_single_cert( + vec![cert.der().clone()], + PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(signing_key.serialize_der())), + ) + .unwrap(); + let acceptor = TlsAcceptor::from(Arc::new(server_config)); + let mut roots = RootCertStore::empty(); + roots.add(cert.der().clone()).unwrap(); + let tls_config = Arc::new( + ClientConfig::builder() + .with_root_certificates(roots) + .with_no_client_auth(), + ); + + let parent_url = format!("wss://{host}/backend-api/codex/responses"); + let review_url = format!("wss://{host}/backend-api/codex/guardian"); + let requests = [ + (parent_url, None), + (review_url.clone(), None), + (review_url.clone(), Some("explicit=keep")), + (format!("wss://{host}/outside"), None), + ( + format!("wss://{other_host}/backend-api/codex/responses"), + None, + ), + ("wss://api.openai.com/v1/responses".to_string(), None), + ("wss://api.openai.com/v1/responses".to_string(), None), + (format!("ws://{host}/backend-api/codex/responses"), None), + (review_url.clone(), None), + (review_url.clone(), None), + (review_url, None), + ]; + let responses: Vec<(StatusCode, Vec<&str>)> = vec![ + ( + StatusCode::SWITCHING_PROTOCOLS, + vec![ + "__oailb=west; Path=/backend-api; Secure; HttpOnly", + "chatgpt_session=never-store; Path=/; Secure", + ], + ), + // Even a failed upgrade can refresh a cookie used by the next connection. + ( + StatusCode::FORBIDDEN, + vec!["__oailb=east; Path=/backend-api; Secure"], + ), + (StatusCode::SWITCHING_PROTOCOLS, vec![]), + (StatusCode::SWITCHING_PROTOCOLS, vec![]), + (StatusCode::SWITCHING_PROTOCOLS, vec![]), + ( + StatusCode::SWITCHING_PROTOCOLS, + vec!["__oailb=untrusted; Path=/; Secure"], + ), + (StatusCode::SWITCHING_PROTOCOLS, vec![]), + ( + StatusCode::SWITCHING_PROTOCOLS, + vec!["__oailb=insecure; Path=/backend-api"], + ), + (StatusCode::SWITCHING_PROTOCOLS, vec![]), + ( + StatusCode::SWITCHING_PROTOCOLS, + vec!["__oailb=; Path=/backend-api; Max-Age=0; Secure"], + ), + (StatusCode::SWITCHING_PROTOCOLS, vec![]), + ]; + let expected_statuses = responses + .iter() + .map(|(status, _)| *status) + .collect::>(); + let secure = requests + .iter() + .map(|(url, _)| url.starts_with("wss:")) + .collect::>(); + + // Terminate CONNECT locally so the real TLS/upgrade path can use ChatGPT hostnames without DNS + // overrides, environment changes, or external traffic. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_url = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let mut received_cookies = Vec::new(); + for ((status, cookies), secure) in responses.into_iter().zip(secure) { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut connect = Vec::new(); + while !connect.ends_with(b"\r\n\r\n") { + connect.push(stream.read_u8().await.unwrap()); + } + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + let stream: Box = if secure { + Box::new(acceptor.accept(stream).await.unwrap()) + } else { + Box::new(stream) + }; + let result = accept_hdr_async(stream, |request: &Request, mut response: Response| { + received_cookies.push(request.headers().get(COOKIE).cloned()); + for cookie in cookies { + response + .headers_mut() + .append(SET_COOKIE, HeaderValue::from_static(cookie)); + } + if status == StatusCode::SWITCHING_PROTOCOLS { + Ok(response) + } else { + *response.status_mut() = status; + Err(response.map(|()| None)) + } + }) + .await; + if status == StatusCode::SWITCHING_PROTOCOLS { + drop(result.unwrap()); + } else { + assert!(matches!(result, Err(WebSocketError::Http(_)))); + } + } + received_cookies + }); + + let mut actual_statuses = Vec::new(); + for (url, explicit_cookie) in requests { + // Separate connector/factory instances must still see the shared routing cookie. + let connector = WebSocketConnector { + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + tls_config: Some(Arc::clone(&tls_config)), + tcp_nodelay: TcpNodelay::Default, + }; + let mut request = url.into_client_request().unwrap(); + if let Some(cookie) = explicit_cookie { + request + .headers_mut() + .insert(COOKIE, HeaderValue::from_static(cookie)); + } + let result = connector + .connect_with_route( + request, + WebSocketConfig::default(), + OutboundProxyRoute::Proxy { + url: proxy_url.clone(), + no_proxy: None, + }, + /*loopback_direct*/ false, + ) + .await; + let status = match result { + Ok((connection, response)) => { + drop(connection); + response.status() + } + Err(WebSocketError::Http(response)) => response.status(), + Err(error) => panic!("unexpected handshake error: {error}"), + }; + actual_statuses.push(status); + } + assert_eq!(actual_statuses, expected_statuses); + assert_eq!( + server.await.unwrap(), + [ + None, + Some("__oailb=west"), + Some("explicit=keep"), + None, + None, + None, + None, + None, + Some("__oailb=east"), + Some("__oailb=east"), + None + ] + .map(|cookie| cookie.map(HeaderValue::from_static)) + .to_vec(), + ); +} diff --git a/codex-rs/websocket-client/src/lib.rs b/codex-rs/websocket-client/src/lib.rs index c709320ec3..e8e5b0a7c4 100644 --- a/codex-rs/websocket-client/src/lib.rs +++ b/codex-rs/websocket-client/src/lib.rs @@ -1,4 +1,5 @@ -//! Proxy-aware WebSocket connection setup shared by Codex API clients. +//! Proxy-aware WebSocket connection setup shared by Codex API clients, reusing the HTTP factory's +//! ChatGPT cookie store for secure handshakes. mod dialer; @@ -27,6 +28,7 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::handshake::client::Request; use tokio_tungstenite::tungstenite::handshake::client::Response; use tokio_tungstenite::tungstenite::http::Uri; +use tokio_tungstenite::tungstenite::http::header::COOKIE; use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; /// Connects WebSockets using the outbound proxy policy resolved by application configuration. @@ -133,12 +135,18 @@ impl WebSocketConnector { async fn connect_with_route( &self, - request: Request, + mut request: Request, config: WebSocketConfig, proxy_route: OutboundProxyRoute, loopback_direct: bool, ) -> Result<(WebSocketConnection, Response), WebSocketError> { - let (inner, response) = dialer::connect( + 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(), @@ -147,7 +155,18 @@ impl WebSocketConnector { loopback_direct, ) .boxed() - .await?; + .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)) } } @@ -239,3 +258,7 @@ impl AsyncIo for T where T: AsyncRead + AsyncWrite + Send + Unpin {} #[cfg(test)] #[path = "lib_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "cookie_tests.rs"] +mod cookie_tests;