diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 9e2bf12fcc..750b874be0 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2094,6 +2094,7 @@ dependencies = [ "codex-app-server-protocol", "codex-config", "codex-core", + "codex-http-state", "codex-login", "codex-model-provider", "codex-state", diff --git a/codex-rs/app-server-transport/Cargo.toml b/codex-rs/app-server-transport/Cargo.toml index 175890962e..84c93890bc 100644 --- a/codex-rs/app-server-transport/Cargo.toml +++ b/codex-rs/app-server-transport/Cargo.toml @@ -25,6 +25,7 @@ clap = { workspace = true, features = ["derive"] } codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-core = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-state = { workspace = true } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs index ac67726371..f12941ec68 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs @@ -6,6 +6,8 @@ use super::protocol::RemoteControlTarget; use super::protocol::StartRemoteControlPairingRequest; use super::protocol::StartRemoteControlPairingResponse; use axum::http::HeaderMap; +use axum::http::HeaderValue; +use axum::http::header::AUTHORIZATION; use codex_api::SharedAuthProvider; use codex_app_server_protocol::RemoteControlPairingStartResponse; use codex_login::default_client::build_reqwest_client; @@ -45,6 +47,7 @@ pub(super) struct RemoteControlEnrollment { impl RemoteControlEnrollment { pub(super) async fn start_pairing( &self, + auth: &RemoteControlConnectionAuth, request: StartRemoteControlPairingRequest, ) -> io::Result { if self.should_refresh_server_token() { @@ -55,10 +58,22 @@ impl RemoteControlEnrollment { .as_deref() .ok_or_else(pairing_unavailable_error)?; + let mut auth_headers = HeaderMap::new(); + auth.server_token_auth_provider + .add_auth_headers_for_url(&self.remote_control_target.pair_url, &mut auth_headers); + auth_headers.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Bearer {remote_control_token}")).map_err(|err| { + io::Error::new( + ErrorKind::InvalidInput, + format!("invalid remote control pairing token: {err}"), + ) + })?, + ); let response = build_reqwest_client() .post(&self.remote_control_target.pair_url) .timeout(REMOTE_CONTROL_PAIRING_TIMEOUT) - .bearer_auth(remote_control_token) + .headers(auth_headers.clone()) .json(&request) .send() .await @@ -69,6 +84,11 @@ impl RemoteControlEnrollment { )) })?; let headers = response.headers().clone(); + auth.server_token_auth_provider.observe_response_headers( + &self.remote_control_target.pair_url, + &auth_headers, + &headers, + ); let status = response.status(); let body = response.bytes().await.map_err(|err| { io::Error::other(format!( @@ -153,6 +173,7 @@ impl RemoteControlEnrollment { pub(super) struct RemoteControlConnectionAuth { pub(super) auth_provider: SharedAuthProvider, + pub(super) server_token_auth_provider: SharedAuthProvider, pub(super) account_id: String, } @@ -430,11 +451,12 @@ where { let client = build_reqwest_client(); let mut auth_headers = HeaderMap::new(); - auth.auth_provider.add_auth_headers(&mut auth_headers); + auth.auth_provider + .add_auth_headers_for_url(url, &mut auth_headers); let response = client .post(url) .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT) - .headers(auth_headers) + .headers(auth_headers.clone()) .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id) .header(REMOTE_CONTROL_INSTALLATION_ID_HEADER, installation_id) .json(request) @@ -446,6 +468,8 @@ where )) })?; let headers = response.headers().clone(); + auth.auth_provider + .observe_response_headers(url, &auth_headers, &headers); let status = response.status(); let body = response.bytes().await.map_err(|err| { io::Error::other(format!( @@ -745,6 +769,7 @@ mod tests { &remote_control_target, &RemoteControlConnectionAuth { auth_provider: codex_model_provider::unauthenticated_auth_provider(), + server_token_auth_provider: codex_model_provider::unauthenticated_auth_provider(), account_id: "account_id".to_string(), }, "11111111-1111-4111-8111-111111111111", diff --git a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs index a55667886f..8c85ea3c39 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs @@ -174,7 +174,7 @@ impl RemoteControlHandle { let pairing_request = || protocol::StartRemoteControlPairingRequest { manual_code: params.manual_code, }; - let pairing_response = match enrollment.start_pairing(pairing_request()).await { + let pairing_response = match enrollment.start_pairing(&auth, pairing_request()).await { Err(err) if err.kind() == io::ErrorKind::PermissionDenied => { clear_pairing_server_token(&self.current_enrollment, &mut enrollment)?; refresh_pairing_enrollment( @@ -185,7 +185,7 @@ impl RemoteControlHandle { &mut enrollment, ) .await?; - enrollment.start_pairing(pairing_request()).await + enrollment.start_pairing(&auth, pairing_request()).await } pairing_response => pairing_response, }; diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs index 176aeb4ed5..5a1a8e3026 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs @@ -1,3 +1,4 @@ +use super::super::enroll::RemoteControlConnectionAuth; use super::super::protocol::StartRemoteControlPairingRequest; use super::*; use pretty_assertions::assert_eq; @@ -21,6 +22,14 @@ fn remote_control_enrollment( } } +fn remote_control_connection_auth() -> RemoteControlConnectionAuth { + RemoteControlConnectionAuth { + auth_provider: codex_model_provider::unauthenticated_auth_provider(), + server_token_auth_provider: codex_model_provider::unauthenticated_auth_provider(), + account_id: "account-id".to_string(), + } +} + async fn pairing_error(status: &'static str, body: &'static str) -> (String, String) { let listener = TcpListener::bind("127.0.0.1:0") .await @@ -41,7 +50,10 @@ async fn pairing_error(status: &'static str, body: &'static str) -> (String, Str }); let err = remote_control_enrollment(&remote_control_url, "remote-control-token") - .start_pairing(StartRemoteControlPairingRequest { manual_code: false }) + .start_pairing( + &remote_control_connection_auth(), + StartRemoteControlPairingRequest { manual_code: false }, + ) .await .expect_err("pairing should fail"); server_task.await.expect("server task should finish"); @@ -59,7 +71,10 @@ async fn pairing_response_error(body: serde_json::Value) -> String { }); let err = remote_control_enrollment(&remote_control_url, "remote-control-token") - .start_pairing(StartRemoteControlPairingRequest { manual_code: false }) + .start_pairing( + &remote_control_connection_auth(), + StartRemoteControlPairingRequest { manual_code: false }, + ) .await .expect_err("pairing should fail"); server_task.await.expect("server task should finish"); diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 5ec4f3d04e..1330ccf86e 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -28,6 +28,8 @@ use base64::Engine; use codex_app_server_protocol::RemoteControlConnectionStatus; use codex_app_server_protocol::RemoteControlStatusChangedNotification; use codex_core::util::backoff; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::UnauthorizedRecovery; use codex_state::StateRuntime; @@ -1124,6 +1126,7 @@ fn set_remote_control_header( fn build_remote_control_websocket_request( websocket_url: &str, + auth_provider: &codex_api::SharedAuthProvider, enrollment: &RemoteControlEnrollment, installation_id: &str, subscribe_cursor: Option<&str>, @@ -1135,6 +1138,7 @@ fn build_remote_control_websocket_request( ) })?; let headers = request.headers_mut(); + auth_provider.add_auth_headers_for_url(websocket_url, headers); set_remote_control_header(headers, "x-codex-server-id", &enrollment.server_id)?; set_remote_control_header( headers, @@ -1206,8 +1210,21 @@ pub(crate) async fn load_remote_control_auth( )); } + let http_state = HttpStateContext::new( + auth_manager.codex_home().to_path_buf(), + HttpStateSurface::CodexRemoteControl, + ); Ok(RemoteControlConnectionAuth { - auth_provider: codex_model_provider::auth_provider_from_auth(&auth), + auth_provider: codex_model_provider::with_native_integrity_state( + codex_model_provider::auth_provider_from_auth(&auth), + Some(&auth), + Some(http_state.clone()), + ), + server_token_auth_provider: codex_model_provider::with_native_integrity_state( + codex_model_provider::unauthenticated_auth_provider(), + Some(&auth), + Some(http_state), + ), account_id: auth.get_account_id().ok_or_else(|| { io::Error::new( ErrorKind::WouldBlock, @@ -1386,10 +1403,12 @@ pub(super) async fn connect_remote_control_websocket( publish_current_enrollment(current_enrollment, enrollment_ref); let request = build_remote_control_websocket_request( &remote_control_target.websocket_url, + &auth.server_token_auth_provider, enrollment_ref, connect_options.installation_id, connect_options.subscribe_cursor, )?; + let request_headers = request.headers().clone(); let websocket_connect_result = tokio::time::timeout( REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT, @@ -1407,8 +1426,22 @@ pub(super) async fn connect_remote_control_websocket( })?; match websocket_connect_result { - Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))), + Ok((websocket_stream, response)) => { + auth.server_token_auth_provider.observe_response_headers( + &remote_control_target.websocket_url, + &request_headers, + response.headers(), + ); + Ok((websocket_stream, response.map(|_| ()))) + } Err(err) => { + if let tungstenite::Error::Http(response) = &err { + auth.server_token_auth_provider.observe_response_headers( + &remote_control_target.websocket_url, + &request_headers, + response.headers(), + ); + } match &err { tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { info!( @@ -1630,6 +1663,7 @@ mod tests { use codex_app_server_protocol::ServerNotification; use codex_config::types::AuthCredentialsStoreMode; use codex_core::test_support::auth_manager_from_auth; + use codex_http_state::HttpStateStore; use codex_login::AuthDotJson; use codex_login::CodexAuth; use codex_login::save_auth; @@ -1673,6 +1707,51 @@ mod tests { } } + #[test] + fn websocket_request_attaches_integrity_state_without_chatgpt_auth_headers() { + let codex_home = TempDir::new().expect("temp dir should create"); + HttpStateStore::new(codex_home.path().to_path_buf()) + .set( + HttpStateSurface::CodexRemoteControl, + "integrity-state".to_string(), + ) + .expect("state should persist"); + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let auth_provider = codex_model_provider::with_native_integrity_state( + codex_model_provider::unauthenticated_auth_provider(), + Some(&auth), + Some(HttpStateContext::new( + codex_home.path().to_path_buf(), + HttpStateSurface::CodexRemoteControl, + )), + ); + + let request = build_remote_control_websocket_request( + "wss://chatgpt.com/backend-api/wham/remote/control/server", + &auth_provider, + &remote_control_enrollment(Some(TEST_REMOTE_CONTROL_SERVER_TOKEN)), + TEST_INSTALLATION_ID, + None, + ) + .expect("request should build"); + + assert_eq!( + request + .headers() + .get("x-oai-is") + .and_then(|value| value.to_str().ok()), + Some("integrity-state") + ); + assert_eq!( + request + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer Remote Control Token") + ); + assert!(!request.headers().contains_key("chatgpt-account-id")); + } + fn test_current_enrollment() -> CurrentRemoteControlEnrollment { Arc::new(std::sync::Mutex::new(None)) }