diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 5d96f89d89..b347a3e441 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -1,7 +1,10 @@ use super::enroll::REMOTE_CONTROL_ACCOUNT_ID_HEADER; use super::enroll::REMOTE_CONTROL_INSTALLATION_ID_HEADER; +use super::enroll::RemoteControlConnectionAuth; use super::enroll::RemoteControlEnrollment; +use super::enroll::enroll_remote_control_server; use super::enroll::load_persisted_remote_control_enrollment; +use super::enroll::refresh_remote_control_server; use super::enroll::update_persisted_remote_control_enrollment; use super::protocol::ClientEnvelope; use super::protocol::ClientEvent; @@ -15,7 +18,10 @@ use crate::outgoing_message::QueuedOutgoingMessage; use crate::transport::CHANNEL_CAPACITY; use crate::transport::ConnectionOrigin; use crate::transport::TransportEvent; +use axum::http::HeaderMap; +use axum::http::HeaderValue; use base64::Engine; +use codex_api::AuthProvider; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::ConfigWarningNotification; use codex_app_server_protocol::JSONRPCMessage; @@ -26,6 +32,8 @@ use codex_app_server_protocol::ServerNotification; use codex_config::types::AuthCredentialsStoreMode; use codex_core::test_support::auth_manager_from_auth; use codex_core::test_support::auth_manager_from_auth_with_home; +use codex_http_state::HttpStateStore; +use codex_http_state::HttpStateSurface; use codex_login::AuthDotJson; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -67,6 +75,60 @@ const TEST_REMOTE_CONTROL_URL: &str = "http://127.0.0.1:1/backend-api/wham/remot const TEST_REMOTE_CONTROL_SERVER_TOKEN: &str = "Remote Control Token"; const TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN: &str = "Refreshed Remote Control Token"; const TEST_REMOTE_CONTROL_SERVER_TOKEN_EXPIRES_AT: &str = "2999-01-01T00:00:00Z"; +const SENT_HTTP_STATE: &str = "ois1.sent.nonce.ciphertext"; +const ROTATED_HTTP_STATE: &str = "ois1.rotated.nonce.ciphertext"; + +struct RemoteControlHttpStateAuthProvider { + store: HttpStateStore, +} + +impl AuthProvider for RemoteControlHttpStateAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) { + if let Ok(Some(state)) = self.store.get(HttpStateSurface::CodexRemoteControl) + && let Ok(state) = HeaderValue::from_str(&state) + { + headers.insert("x-oai-is", 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(rotated_state) = response_headers + .get("x-oai-is-update") + .and_then(|value| value.to_str().ok()) + else { + return; + }; + let _ = self.store.compare_and_set( + HttpStateSurface::CodexRemoteControl, + sent_state, + rotated_state.to_string(), + ); + } +} + +fn remote_control_connection_auth_with_http_state( + store: HttpStateStore, +) -> RemoteControlConnectionAuth { + let auth_provider = Arc::new(RemoteControlHttpStateAuthProvider { store }); + RemoteControlConnectionAuth { + auth_provider: auth_provider.clone(), + server_token_auth_provider: auth_provider, + account_id: "account-id".to_string(), + } +} fn remote_control_auth_manager() -> Arc { auth_manager_from_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) @@ -1116,6 +1178,158 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() { let _ = remote_task.await; } +#[tokio::test] +async fn enroll_remote_control_server_attaches_and_rotates_http_state() { + let codex_home = TempDir::new().expect("temp dir should create"); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set( + HttpStateSurface::CodexRemoteControl, + SENT_HTTP_STATE.to_string(), + ) + .expect("HTTP state should store"); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_target = + normalize_remote_control_url(&remote_control_url_for_listener(&listener)) + .expect("target should normalize"); + let server_task = tokio::spawn(async move { + let enroll_request = accept_http_request(&listener).await; + assert_eq!( + enroll_request.request_line, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + assert_eq!( + enroll_request.headers.get("x-oai-is"), + Some(&SENT_HTTP_STATE.to_string()) + ); + assert_eq!( + enroll_request.headers.get(REMOTE_CONTROL_ACCOUNT_ID_HEADER), + Some(&"account-id".to_string()) + ); + assert_eq!( + enroll_request + .headers + .get(REMOTE_CONTROL_INSTALLATION_ID_HEADER), + Some(&TEST_INSTALLATION_ID.to_string()) + ); + respond_with_status_and_headers( + enroll_request.stream, + "200 OK", + &[("x-oai-is-update", ROTATED_HTTP_STATE)], + &remote_control_server_token_response( + "server-id", + "environment-id", + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ) + .to_string(), + ) + .await; + }); + + let enrollment = enroll_remote_control_server( + &remote_control_target, + &remote_control_connection_auth_with_http_state(store.clone()), + TEST_INSTALLATION_ID, + "server-name", + ) + .await + .expect("enrollment should succeed"); + server_task.await.expect("server task should finish"); + assert_eq!( + enrollment.remote_control_token.as_deref(), + Some(TEST_REMOTE_CONTROL_SERVER_TOKEN) + ); + assert_eq!( + store + .get(HttpStateSurface::CodexRemoteControl) + .expect("HTTP state should load"), + Some(ROTATED_HTTP_STATE.to_string()) + ); +} + +#[tokio::test] +async fn refresh_remote_control_server_attaches_and_rotates_http_state() { + let codex_home = TempDir::new().expect("temp dir should create"); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set( + HttpStateSurface::CodexRemoteControl, + SENT_HTTP_STATE.to_string(), + ) + .expect("HTTP state should store"); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_target = + normalize_remote_control_url(&remote_control_url_for_listener(&listener)) + .expect("target should normalize"); + let mut enrollment = RemoteControlEnrollment { + remote_control_target, + account_id: "account-id".to_string(), + environment_id: "environment-id".to_string(), + server_id: "server-id".to_string(), + server_name: "server-name".to_string(), + remote_control_token: Some("stale-token".to_string()), + expires_at: None, + }; + let server_task = tokio::spawn(async move { + let refresh_request = accept_http_request(&listener).await; + assert_eq!( + refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + assert_eq!( + refresh_request.headers.get("x-oai-is"), + Some(&SENT_HTTP_STATE.to_string()) + ); + assert_eq!( + refresh_request + .headers + .get(REMOTE_CONTROL_ACCOUNT_ID_HEADER), + Some(&"account-id".to_string()) + ); + assert_eq!( + refresh_request + .headers + .get(REMOTE_CONTROL_INSTALLATION_ID_HEADER), + Some(&TEST_INSTALLATION_ID.to_string()) + ); + respond_with_status_and_headers( + refresh_request.stream, + "200 OK", + &[("x-oai-is-update", ROTATED_HTTP_STATE)], + &remote_control_server_token_response( + "server-id", + "environment-id", + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ) + .to_string(), + ) + .await; + }); + + refresh_remote_control_server( + &remote_control_connection_auth_with_http_state(store.clone()), + TEST_INSTALLATION_ID, + &mut enrollment, + ) + .await + .expect("refresh should succeed"); + server_task.await.expect("server task should finish"); + assert_eq!( + enrollment.remote_control_token.as_deref(), + Some(TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN) + ); + assert_eq!( + store + .get(HttpStateSurface::CodexRemoteControl) + .expect("HTTP state should load"), + Some(ROTATED_HTTP_STATE.to_string()) + ); +} + #[tokio::test] async fn remote_control_http_mode_enrolls_before_connecting() { let listener = TcpListener::bind("127.0.0.1:0") 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 5a1a8e3026..bd6c717cbf 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,4 +1,3 @@ -use super::super::enroll::RemoteControlConnectionAuth; use super::super::protocol::StartRemoteControlPairingRequest; use super::*; use pretty_assertions::assert_eq; @@ -81,6 +80,66 @@ async fn pairing_response_error(body: serde_json::Value) -> String { err.to_string() } +#[tokio::test] +async fn start_pairing_attaches_and_rotates_http_state_without_replacing_authorization() { + let codex_home = TempDir::new().expect("temp dir should create"); + let store = HttpStateStore::new(codex_home.path().to_path_buf()); + store + .set( + HttpStateSurface::CodexRemoteControl, + SENT_HTTP_STATE.to_string(), + ) + .expect("HTTP state should store"); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let server_task = tokio::spawn(async move { + let pairing_request = accept_http_request(&listener).await; + assert_eq!( + pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + pairing_request.headers.get("authorization"), + Some(&"Bearer remote-control-token".to_string()) + ); + assert_eq!( + pairing_request.headers.get("x-oai-is"), + Some(&SENT_HTTP_STATE.to_string()) + ); + respond_with_status_and_headers( + pairing_request.stream, + "200 OK", + &[("x-oai-is-update", ROTATED_HTTP_STATE)], + &json!({ + "pairing_code": "pairing-code", + "manual_pairing_code": null, + "server_id": "server-id", + "environment_id": "environment-id", + "expires_at": "3026-05-22T12:34:56Z", + }) + .to_string(), + ) + .await; + }); + + remote_control_enrollment(&remote_control_url, "remote-control-token") + .start_pairing( + &remote_control_connection_auth_with_http_state(store.clone()), + StartRemoteControlPairingRequest { manual_code: false }, + ) + .await + .expect("pairing should succeed"); + server_task.await.expect("server task should finish"); + assert_eq!( + store + .get(HttpStateSurface::CodexRemoteControl) + .expect("HTTP state should load"), + Some(ROTATED_HTTP_STATE.to_string()) + ); +} + #[tokio::test] async fn remote_control_handle_starts_pairing_before_websocket_connects() { let listener = TcpListener::bind("127.0.0.1:0")