diff --git a/codex-rs/exec-server/src/client/connection_failure.rs b/codex-rs/exec-server/src/client/connection_failure.rs index 984378b6cb..af0d1df9aa 100644 --- a/codex-rs/exec-server/src/client/connection_failure.rs +++ b/codex-rs/exec-server/src/client/connection_failure.rs @@ -1,4 +1,4 @@ -//! Preserve policy denials through terminal connection state and allow later policy recovery. +//! Preserve policy denials and registry auth errors through terminal connection state. use codex_http_client::NetworkPolicyDenied; use codex_http_client::RouteAwareRequestError; @@ -11,6 +11,8 @@ pub(super) enum ConnectionFailure { Disconnected(String), #[error(transparent)] Policy(NetworkPolicyDenied), + #[error("environment registry authentication error: {0}")] + RegistryAuth(String), } impl From for ExecServerError { @@ -18,6 +20,7 @@ impl From for ExecServerError { match failure { ConnectionFailure::Disconnected(message) => Self::Disconnected(message), ConnectionFailure::Policy(denial) => Self::ApplicationNetworkPolicy(denial), + ConnectionFailure::RegistryAuth(message) => Self::EnvironmentRegistryAuth(message), } } } diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 15a4f25588..78c945aec9 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -456,6 +456,11 @@ impl Inner { return; } if !is_retryable_recovery_error(&error) { + if let ExecServerError::EnvironmentRegistryAuth(message) = error { + self.fail(super::ConnectionFailure::RegistryAuth(message)) + .await; + return; + } break error.to_string(); } } diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index c0ed2c8748..51c2eb1482 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -207,7 +207,7 @@ struct OpenNoiseRendezvousConnection { connection: JsonRpcConnection, pending_tasks: Vec>, options: ExecServerClientConnectOptions, - handshake_ready: tokio::sync::oneshot::Receiver<()>, + handshake_ready: tokio::sync::oneshot::Receiver>, } struct ReadyNoiseRendezvousConnection { @@ -749,7 +749,7 @@ impl ExecServerClient { initialize_span: &tracing::Span, ) -> Result<(), ExecServerError> { match timeout_at(deadline, &mut connection.handshake_ready).await { - Ok(Ok(())) => Ok(()), + Ok(Ok(result)) => result, Ok(Err(_)) => Err(ExecServerError::Disconnected( "Noise harness handshake failed before connection became ready".to_string(), )), diff --git a/codex-rs/exec-server/src/noise_channel.rs b/codex-rs/exec-server/src/noise_channel.rs index cde2aae4ee..42223b8f16 100644 --- a/codex-rs/exec-server/src/noise_channel.rs +++ b/codex-rs/exec-server/src/noise_channel.rs @@ -3,7 +3,8 @@ //! The harness initiates hybrid IK and pins the exec-server static key returned //! by the registry. The first handshake message lets the exec-server authenticate //! the harness static key; the exec-server then asks the registry whether that -//! key is authorized before completing the handshake. +//! key is authorized before opening a transport. A denied request can receive an +//! authenticated rejection without creating a transport on either peer. //! //! "Hybrid" means the session keys include both X25519 and ML-KEM-768 key //! agreement. Once the two-message handshake finishes, AES-GCM protects the @@ -33,6 +34,7 @@ use serde::Serialize; pub(crate) const NOISE_CHANNEL_SUITE: &str = "Noise_hybridIK_X25519+MLKEM768_AESGCM_SHA256"; const PROLOGUE_DOMAIN: &[u8] = b"codex-exec-server-relay-noise/v1"; +const REGISTRY_AUTH_REJECTED: &[u8] = b"codex-registry-authentication-rejected/v1"; type Handshake = HybridHandshake; type Transport = TransportState; @@ -152,12 +154,16 @@ impl InitiatorHandshake { } /// Consume the executor response and enter transport mode. - /// The v1 response does not carry an application payload. + /// Successful responses have no payload. A fixed authenticated rejection + /// fails this request without opening an application transport. pub(crate) fn finish(mut self, response: &[u8]) -> Result { ensure_noise_frame_len(response.len(), "handshake response is too large")?; let overhead = self.handshake.get_next_message_overhead()?; let mut payload = vec![0u8; response.len().saturating_sub(overhead)]; let payload_len = self.handshake.read_message(response, &mut payload)?; + if &payload[..payload_len] == REGISTRY_AUTH_REJECTED { + return Err(NoiseChannelError::RegistryAuthenticationRejected); + } if payload_len != 0 { return Err(NoiseChannelError::InvalidMessage( "handshake response payload must be empty", @@ -227,6 +233,18 @@ impl PendingResponderHandshake { response, )) } + + /// Authenticate a terminal registry 401 to the harness without authorizing + /// its key or constructing a transport. Older peers reject nonempty payloads. + pub(crate) fn reject_registry_authentication(mut self) -> Result, NoiseChannelError> { + let overhead = self.handshake.get_next_message_overhead()?; + let mut response = vec![0u8; overhead + REGISTRY_AUTH_REJECTED.len()]; + let response_len = self + .handshake + .write_message(REGISTRY_AUTH_REJECTED, &mut response)?; + response.truncate(response_len); + Ok(response) + } } /// Established channel with independent implicit send and receive nonces. @@ -294,6 +312,8 @@ fn ensure_noise_frame_len( #[derive(Debug, thiserror::Error)] pub enum NoiseChannelError { + #[error("environment registry authentication rejected")] + RegistryAuthenticationRejected, #[error("Noise channel key generation failed: {0}")] KeyGeneration(String), #[error("invalid Noise channel public key: {0}")] diff --git a/codex-rs/exec-server/src/noise_channel_tests.rs b/codex-rs/exec-server/src/noise_channel_tests.rs index b298e8e580..b5cefb9291 100644 --- a/codex-rs/exec-server/src/noise_channel_tests.rs +++ b/codex-rs/exec-server/src/noise_channel_tests.rs @@ -83,6 +83,40 @@ fn initiator_rejects_wrong_responder_key() { ); } +#[test] +fn registry_rejection_must_be_authenticated() { + for tamper in [false, true] { + let initiator = NoiseChannelIdentity::generate().expect("initiator"); + let responder = NoiseChannelIdentity::generate().expect("responder"); + let prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + let (handshake, request) = InitiatorHandshake::start( + &initiator, + &responder.public_key(), + &prologue, + b"authorization", + ) + .expect("request"); + let pending = PendingResponderHandshake::read_request(&responder, &prologue, &request) + .expect("read request"); + let mut response = pending.reject_registry_authentication().expect("rejection"); + if tamper { + *response.last_mut().expect("encrypted payload") ^= 1; + } + let error = handshake + .finish(&response) + .err() + .expect("must not open transport"); + if tamper { + assert!(matches!(error, NoiseChannelError::Handshake(_))); + } else { + assert!(matches!( + error, + NoiseChannelError::RegistryAuthenticationRejected + )); + } + } +} + #[test] fn responder_rejects_mismatched_prologue() { let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index d1d6670d02..dc985a8d84 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -31,6 +31,7 @@ use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::noise_channel::InitiatorHandshake; +use crate::noise_channel::NoiseChannelError; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; use crate::noise_channel::NoiseTransport; @@ -73,7 +74,7 @@ pub(crate) struct NoiseHarnessConnectionArgs { /// transport is ready for application messages. pub(crate) struct NoiseHarnessConnection { pub(crate) connection: JsonRpcConnection, - pub(crate) handshake_ready: oneshot::Receiver<()>, + pub(crate) handshake_ready: oneshot::Receiver>, pub(crate) stream_id: String, pub(crate) handshake_started: Instant, } @@ -277,10 +278,21 @@ where "Noise harness handshake completed" ); if let Some(handshake_ready_tx) = handshake_ready_tx.take() { - let _ = handshake_ready_tx.send(()); + let _ = handshake_ready_tx.send(Ok(())); } break transport; } + Err(NoiseChannelError::RegistryAuthenticationRejected) => { + record_handshake_failure("registry_authentication_rejected"); + if let Some(handshake_ready_tx) = handshake_ready_tx.take() { + let _ = handshake_ready_tx.send(Err(ExecServerError::EnvironmentRegistryAuth( + "executor credentials rejected".to_string(), + ))); + } + send_disconnected(&incoming_tx, &disconnected_tx, + "environment registry rejected executor credentials".to_string()); + return; + } Err(error) => { record_handshake_failure("handshake_finish"); send_disconnected( diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 4ba54a43dc..3b8c6f2dfe 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -495,7 +495,7 @@ where /// /// Implementations must consult an authority independent of rendezvous. The /// exec-server invokes this after parsing the first IK message and before -/// completing the responder handshake. +/// opening an application transport. pub(crate) trait HarnessKeyValidator: Send + Sync { fn validate_harness_key( &self, @@ -645,7 +645,20 @@ where reason, registry_status, ); - send_reset(&physical_outgoing_tx, validation_result.stream_id); + // Authenticate the rejection so the harness can stop retrying. + if registry_status == Some(401) + && let Ok(response) = pending.handshake.reject_registry_authentication() + { + let frame = RelayMessageFrame::handshake( + validation_result.stream_id, + response, + ); + let _ = physical_outgoing_tx.try_send( + encode_relay_message_frame(&frame), + ); + } else { + send_reset(&physical_outgoing_tx, validation_result.stream_id); + } record_failed_handshake( &mut failed_handshakes, &mut handshake_cooldown_until, @@ -864,8 +877,8 @@ where ); let validator = validator.clone(); - // Failed validation leaves no transport state and sends only a - // generic reset. + // Failed validation leaves no transport state. Preserve the + // trusted registry status for request-scoped rejection handling. validation_tasks.spawn(async move { let result = match timeout( HARNESS_KEY_VALIDATION_TIMEOUT, diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index e365547487..ed1fa073fc 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -432,11 +432,8 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { if !status.is_success() { // The request contains the short-lived authorization. Do not include // a response body that might echo it in logs or error chains. - if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { - return Err(ExecServerError::EnvironmentRegistryAuth(format!( - "environment registry harness key validation authentication failed ({status})" - ))); - } + // Preserve the trusted HTTP status: a 401 rejects this executor's + // credentials; a 403 can reject a stale harness registration/grant. return Err(ExecServerError::EnvironmentRegistryHttp { status, code: None, diff --git a/codex-rs/exec-server/tests/relay/auth_rejection_tests.rs b/codex-rs/exec-server/tests/relay/auth_rejection_tests.rs new file mode 100644 index 0000000000..f3e9fe9e71 --- /dev/null +++ b/codex-rs/exec-server/tests/relay/auth_rejection_tests.rs @@ -0,0 +1,186 @@ +#![cfg(test)] +//! Exercise authorization failures in the real session-recovery loop. + +use super::*; +use codex_exec_server::EnvironmentObservedStatus; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use prost::Message as _; +use std::collections::BTreeSet; +use std::sync::atomic::AtomicBool; +use std::time::Instant; +use tokio_tungstenite::tungstenite::Message; + +#[path = "../../src/proto/codex.exec_server.relay.v1.rs"] +mod wire; + +// Keep the executor connection alive while closing/replacing harness connections, +// as rendezvous does. This is the same topology as the original retry reproduction. +async fn relay_reconnecting_harness( + mut executor: tokio_tungstenite::WebSocketStream, + mut harness: tokio_tungstenite::WebSocketStream, + listener: TcpListener, + mut reset: mpsc::Receiver<()>, +) -> Result<()> { + loop { + let mut stream_id = String::new(); + loop { + tokio::select! { + Some(()) = reset.recv() => { + harness.close(None).await?; + break; + } + message = harness.next() => { + let Some(Ok(message)) = message else { break; }; + if matches!(message, Message::Close(_)) { break; } + if let Message::Binary(bytes) = &message { + stream_id = wire::RelayMessageFrame::decode(bytes.as_ref())?.stream_id; + } + executor.send(message).await?; + } + message = executor.next() => { + let message = message.context("executor must stay connected")??; + if let Message::Ping(bytes) = message { + executor.send(Message::Pong(bytes)).await?; + } else { + harness.send(message).await?; + } + } + } + } + let reset = wire::RelayMessageFrame { + version: 1, + stream_id, + body: Some(wire::relay_message_frame::Body::Reset(wire::RelayReset { + reason: "harness disconnected".to_owned(), + })), + ..Default::default() + }; + executor + .send(Message::Binary(reset.encode_to_vec().into())) + .await?; + harness = accept_websocket(&listener, "next harness").await?; + } +} + +#[test_case::test_case(401; "rejected_executor_credentials")] +#[test_case::test_case(403; "stale_harness_grant")] +#[test_case::test_case(502; "identity_service_unavailable")] +#[test_case::test_case(503; "service_unavailable")] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn recovery_preserves_auth_failures(status: u16) -> Result<()> { + let fixture = RegistryFixture::new(vec![200], Duration::ZERO).await?; + let reject = Arc::new(AtomicBool::new(false)); + let validations = Arc::new(AtomicUsize::new(0)); + Mock::given(method("POST")) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/validate" + ))) + .respond_with({ + let reject = reject.clone(); + let validations = validations.clone(); + move |_: &wiremock::Request| { + let attempt = validations.fetch_add(1, Ordering::SeqCst); + if reject.load(Ordering::SeqCst) && (status == 401 || attempt < 4) { + ResponseTemplate::new(status).set_body_string("secret-must-not-leak") + } else { + ResponseTemplate::new(200).set_body_json(serde_json::json!({"valid": true})) + } + } + }) + .with_priority(1) + .mount(&fixture.registry) + .await; + let (shutdown, remote) = fixture.start()?; + let executor = accept_websocket(&fixture.listener, "executor").await?; + let provider = Arc::new(FreshBundleNoiseConnectProvider { + websocket_url: format!("ws://{}/relay?role=harness", fixture.listener.local_addr()?), + executor_public_key: registered_executor_public_key(&fixture.registry).await?, + calls: AtomicUsize::new(0), + }); + let manager = EnvironmentManager::without_environments(HttpClientFactory::new( + OutboundProxyPolicy::ReqwestDefault, + )); + let environment = manager + .materialize_pending_noise_environment(ENVIRONMENT_ID.to_owned(), provider.clone())?; + let mut state = environment.subscribe_connection_state().context("state")?; + manager.report_environment_provisioning_status( + ENVIRONMENT_ID.to_owned(), + Ok(EnvironmentReadyInfo { + selected_capability_roots: vec![], + }), + provider.clone(), + )?; + let first = AbortOnDropHandle::new(tokio::spawn({ + let environment = environment.clone(); + async move { environment.force_info().await } + })); + let harness = accept_websocket(&fixture.listener, "harness").await?; + let (reset_tx, reset_rx) = mpsc::channel(1); + let router = AbortOnDropHandle::new(tokio::spawn(relay_reconnecting_harness( + executor, + harness, + fixture.listener, + reset_rx, + ))); + let original_info = timeout(TEST_TIMEOUT, first).await???; + state.borrow_and_update(); + reject.store(true, Ordering::SeqCst); + let started = Instant::now(); + reset_tx.send(()).await?; + let terminal = timeout(TEST_TIMEOUT, async { + loop { + state.changed().await?; + let current = *state.borrow_and_update(); + if matches!( + current, + EnvironmentConnectionState::Connected | EnvironmentConnectionState::Disconnected + ) { + break anyhow::Ok(current); + } + } + }) + .await??; + if status == 401 { + assert_eq!(terminal, EnvironmentConnectionState::Disconnected); + assert!(matches!(environment.status().await, + EnvironmentObservedStatus::Disconnected { error } + if error == "environment registry authentication error: executor credentials rejected")); + assert_eq!(provider.calls(), 2); // Initial connection plus one rejected recovery. + } else { + assert_eq!(terminal, EnvironmentConnectionState::Connected); + assert_eq!(environment.force_info().await?, original_info); + assert_eq!(provider.calls(), 5); // Initial, three rejected recoveries, then success. + } + assert_eq!(validations.load(Ordering::SeqCst), provider.calls()); + let requests = fixture + .registry + .received_requests() + .await + .context("requests")?; + let validations: Vec<_> = requests + .iter() + .filter(|r| r.url.path().ends_with("/validate")) + .collect(); + let bearers: BTreeSet<_> = validations + .iter() + .map(|r| r.headers.get("authorization").unwrap()) + .collect(); + let grants: BTreeSet<_> = validations.iter().map(|r| { + serde_json::from_slice::(&r.body).unwrap()["harness_key_authorization"].to_string() + }).collect(); + assert_eq!(bearers.len(), 1); + assert_eq!(grants.len(), validations.len()); + println!( + "status={status}, recovery_attempts={}, elapsed_ms={}, terminal={terminal:?}", + provider.calls() - 1, + started.elapsed().as_millis() + ); + assert!(!router.is_finished() && !remote.is_finished()); + router.abort(); + let _ = router.await; + let _ = shutdown.send(()); + timeout(TEST_TIMEOUT, remote).await???; + Ok(()) +} diff --git a/codex-rs/exec-server/tests/relay/registration_retry_tests.rs b/codex-rs/exec-server/tests/relay/registration_retry_tests.rs index 2df73491ed..987b4b34a8 100644 --- a/codex-rs/exec-server/tests/relay/registration_retry_tests.rs +++ b/codex-rs/exec-server/tests/relay/registration_retry_tests.rs @@ -13,6 +13,9 @@ use tokio_tungstenite::tungstenite::handshake::server::Response; use super::*; +#[path = "auth_rejection_tests.rs"] +mod auth_rejection; + type RemoteTask = AbortOnDropHandle>; type RelayTask = AbortOnDropHandle>;