diff --git a/codex-rs/exec-server/src/client_refresh_tests.rs b/codex-rs/exec-server/src/client_refresh_tests.rs index 9a5cb0a453..b5d5c7c607 100644 --- a/codex-rs/exec-server/src/client_refresh_tests.rs +++ b/codex-rs/exec-server/src/client_refresh_tests.rs @@ -111,6 +111,7 @@ impl HarnessKeyValidator for Validator { &self, _: &NoiseChannelPublicKey, _: &str, + _: Option, ) -> Result<(), ExecServerError> { if let Some(handshake) = &self.handshake { self.started.notify_one(); diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 9de23ca083..1854a83aee 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -179,7 +179,10 @@ pub(crate) struct ReconnectAttempt { } struct OpenNoiseRendezvousConnection { + environment_id: String, executor_registration_id: String, + stream_id: String, + handshake_started: Instant, connection: JsonRpcConnection, options: ExecServerClientConnectOptions, handshake_ready: tokio::sync::oneshot::Receiver<()>, @@ -686,7 +689,7 @@ impl ExecServerClient { stream, NoiseHarnessConnectionArgs { connection_label, - environment_id, + environment_id: environment_id.clone(), executor_registration_id: executor_registration_id.clone(), identity: harness_identity, responder_public_key: executor_public_key, @@ -694,7 +697,10 @@ impl ExecServerClient { }, ); Ok(OpenNoiseRendezvousConnection { + environment_id, executor_registration_id, + stream_id: connection.stream_id, + handshake_started: connection.handshake_started, connection: connection.connection, options: ExecServerClientConnectOptions { client_name, @@ -715,12 +721,12 @@ impl ExecServerClient { ) )] async fn wait_for_noise_handshake( - handshake_ready: &mut tokio::sync::oneshot::Receiver<()>, + connection: &mut OpenNoiseRendezvousConnection, deadline: Instant, initialize_timeout: Duration, initialize_span: &tracing::Span, ) -> Result<(), ExecServerError> { - match timeout_at(deadline, handshake_ready).await { + match timeout_at(deadline, &mut connection.handshake_ready).await { Ok(Ok(())) => Ok(()), Ok(Err(_)) => Err(ExecServerError::Disconnected( "Noise harness handshake failed before connection became ready".to_string(), @@ -736,6 +742,10 @@ impl ExecServerClient { noise_side = "harness", noise_reason = "timeout", timeout_ms = initialize_timeout.as_millis() as u64, + elapsed_ms = connection.handshake_started.elapsed().as_millis() as u64, + environment_id = %connection.environment_id, + executor_registration_id = %connection.executor_registration_id, + stream_id = %connection.stream_id, ); Err(ExecServerError::InitializeTimedOut { timeout: initialize_timeout, @@ -751,7 +761,7 @@ impl ExecServerClient { // startup parent while making its two child operations visible. let initialize_timeout = connection.options.initialize_timeout; let noise_context = NoiseInitializeContext { - executor_registration_id: connection.executor_registration_id, + executor_registration_id: connection.executor_registration_id.clone(), span: tracing::info_span!( "codex.exec_server.request", otel.kind = "client", @@ -762,7 +772,7 @@ impl ExecServerClient { }; let deadline = Instant::now() + initialize_timeout; let readiness = Self::wait_for_noise_handshake( - &mut connection.handshake_ready, + &mut connection, deadline, initialize_timeout, &noise_context.span, diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index 61cd78156b..408841663f 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -18,6 +18,8 @@ use tokio::io::duplex; use tokio::net::TcpListener; use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; +use tracing::instrument::WithSubscriber; +use tracing_subscriber::layer::SubscriberExt; use super::ExecServerClient; use super::ExecServerReconnectStrategy; @@ -244,7 +246,10 @@ async fn noise_handshake_uses_initialize_timeout() -> Result<()> { let mut websocket = accept_async(socket).await?; // Drain the frames sent before the harness waits for the responder, // then verify that a timed-out readiness wait closes the socket. - assert!(websocket.next().await.is_some()); + let Some(Ok(Message::Binary(resume))) = websocket.next().await else { + anyhow::bail!("expected resume frame"); + }; + let stream_id = decode_relay_message_frame(&resume)?.stream_id; assert!(websocket.next().await.is_some()); let closed = tokio::time::timeout(std::time::Duration::from_secs(1), websocket.next()).await?; @@ -252,7 +257,7 @@ async fn noise_handshake_uses_initialize_timeout() -> Result<()> { matches!(closed, None | Some(Ok(Message::Close(_))) | Some(Err(_))), "timed-out Noise handshake must close its websocket" ); - anyhow::Ok(()) + anyhow::Ok(stream_id) }); let initialize_timeout = std::time::Duration::from_millis(1); let opened = ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { @@ -274,7 +279,16 @@ async fn noise_handshake_uses_initialize_timeout() -> Result<()> { }) .await?; + let stream_id = opened.stream_id.clone(); + let mut log_file = tempfile::tempfile()?; + let writer = log_file.try_clone()?; + let subscriber = tracing_subscriber::registry().with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(move || writer.try_clone().expect("clone log file")), + ); let error = ExecServerClient::finish_noise_rendezvous_connection(opened) + .with_subscriber(subscriber) .await .err() .expect("stalled Noise handshake must time out"); @@ -283,7 +297,19 @@ async fn noise_handshake_uses_initialize_timeout() -> Result<()> { ExecServerError::InitializeTimedOut { timeout } if timeout == initialize_timeout )); - server.await??; + assert_eq!(server.await??, stream_id); + std::io::Seek::rewind(&mut log_file)?; + let mut logs = String::new(); + std::io::Read::read_to_string(&mut log_file, &mut logs)?; + for field in [ + "noise_reason=\"timeout\"", + "environment_id=environment", + "executor_registration_id=registration", + "elapsed_ms=", + &stream_id, + ] { + assert!(logs.contains(field), "missing {field} in {logs}"); + } Ok(()) } diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 35b574d4f1..66db172867 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -14,7 +14,10 @@ use futures::StreamExt; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::watch; +use tokio::time::Instant; +use tokio_tungstenite::tungstenite::Error as WebSocketError; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::error::ProtocolError; use tracing::Instrument; use tracing::debug; use tracing::info; @@ -66,6 +69,8 @@ pub(crate) struct NoiseHarnessConnectionArgs { pub(crate) struct NoiseHarnessConnection { pub(crate) connection: JsonRpcConnection, pub(crate) handshake_ready: oneshot::Receiver<()>, + pub(crate) stream_id: String, + pub(crate) handshake_started: Instant, } // Reset frames are cleartext relay control and are not authenticated by Noise. @@ -80,13 +85,16 @@ const MAX_FRAMES_DRAINED_AFTER_PONG_DEADLINE: usize = 32; /// Callers that send application messages immediately after opening the /// websocket must await handshake_ready first. Dropping that receiver keeps /// the legacy fire-and-forget behavior for existing connection owners. -pub(crate) fn noise_harness_connection_from_websocket_with_readiness( +pub(crate) fn noise_harness_connection_from_websocket_with_readiness( stream: T, args: NoiseHarnessConnectionArgs, ) -> NoiseHarnessConnection where - T: Sink + Stream> + Unpin + Send + 'static, - E: std::fmt::Display + Send + 'static, + T: Sink + + Stream> + + Unpin + + Send + + 'static, { let NoiseHarnessConnectionArgs { connection_label, @@ -97,12 +105,13 @@ where harness_key_authorization, } = args; let stream_id = Uuid::new_v4().to_string(); + let connection_stream_id = stream_id.clone(); + let handshake_started = Instant::now(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); let (disconnected_tx, disconnected_rx) = watch::channel(false); let (handshake_ready_tx, handshake_ready_rx) = oneshot::channel(); - let stream_span = - tracing::debug_span!("noise_relay.stream", noise_side = "harness",).or_current(); + let stream_span = tracing::debug_span!("noise_relay.stream", noise_side = "harness",); debug!( environment_id, executor_registration_id, stream_id, "Noise harness relay details" @@ -122,6 +131,7 @@ where noise_outcome = "failed", noise_side = "harness", noise_reason = reason, + elapsed_ms = handshake_started.elapsed().as_millis() as u64, environment_id, executor_registration_id, stream_id, @@ -207,7 +217,7 @@ where return; } Err(error) => { - record_handshake_failure("websocket_read"); + record_handshake_failure(websocket_read_error_reason(&error)); send_disconnected( &incoming_tx, &disconnected_tx, @@ -253,6 +263,11 @@ where info!( noise_event = "handshake", noise_outcome = "ok", + noise_side = "harness", + elapsed_ms = handshake_started.elapsed().as_millis() as u64, + environment_id, + executor_registration_id, + stream_id, "Noise harness handshake completed" ); if let Some(handshake_ready_tx) = handshake_ready_tx.take() { @@ -574,6 +589,28 @@ where transport: JsonRpcTransport::Plain, }, handshake_ready: handshake_ready_rx, + stream_id: connection_stream_id, + handshake_started, + } +} + +// Classify locally typed errors only; their text can contain untrusted wire data. +fn websocket_read_error_reason(error: &WebSocketError) -> &'static str { + match error { + WebSocketError::Io(error) => match error.kind() { + std::io::ErrorKind::ConnectionReset => "websocket_read_connection_reset", + std::io::ErrorKind::ConnectionAborted => "websocket_read_connection_aborted", + std::io::ErrorKind::UnexpectedEof => "websocket_read_unexpected_eof", + std::io::ErrorKind::TimedOut => "websocket_read_timed_out", + _ => "websocket_read_io_other", + }, + WebSocketError::Protocol(ProtocolError::ResetWithoutClosingHandshake) => { + "websocket_read_reset_without_close" + } + WebSocketError::Protocol(_) => "websocket_read_protocol", + WebSocketError::ConnectionClosed | WebSocketError::AlreadyClosed => "websocket_read_closed", + WebSocketError::Tls(_) => "websocket_read_tls", + _ => "websocket_read_other", } } diff --git a/codex-rs/exec-server/src/noise_relay/harness_tests.rs b/codex-rs/exec-server/src/noise_relay/harness_tests.rs index 5530941ca0..e8b9e98831 100644 --- a/codex-rs/exec-server/src/noise_relay/harness_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/harness_tests.rs @@ -29,13 +29,16 @@ use super::*; use crate::connection::JsonRpcConnectionEvent; use crate::noise_channel::PendingResponderHandshake; -fn noise_harness_connection_from_websocket( +fn noise_harness_connection_from_websocket( stream: T, args: NoiseHarnessConnectionArgs, ) -> JsonRpcConnection where - T: Sink + Stream> + Unpin + Send + 'static, - E: std::fmt::Display + Send + 'static, + T: Sink + + Stream> + + Unpin + + Send + + 'static, { noise_harness_connection_from_websocket_with_readiness(stream, args).connection } @@ -43,6 +46,24 @@ where const ENVIRONMENT_ID: &str = "environment-1"; const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; +#[test] +fn websocket_read_diagnostics_use_typed_errors_only() { + let reset = std::io::Error::new(std::io::ErrorKind::ConnectionReset, "private data"); + for (error, expected) in [ + (WebSocketError::Io(reset), "websocket_read_connection_reset"), + ( + WebSocketError::Protocol(ProtocolError::ResetWithoutClosingHandshake), + "websocket_read_reset_without_close", + ), + ( + WebSocketError::Protocol(ProtocolError::InvalidOpcode(255)), + "websocket_read_protocol", + ), + ] { + assert_eq!(websocket_read_error_reason(&error), expected); + } +} + #[tokio::test(start_paused = true)] async fn first_encrypted_request_frame_exposes_only_its_trace_context() -> Result<()> { let (connection, mut control, mut outbound_rx) = connected_controlled_harness().await?; @@ -368,7 +389,7 @@ async fn connected_controlled_harness() -> Result<( } struct ControlledWebSocket { - inbound_rx: futures_mpsc::UnboundedReceiver>, + inbound_rx: futures_mpsc::UnboundedReceiver>, outbound_tx: futures_mpsc::UnboundedSender, write_permit_rx: futures_mpsc::UnboundedReceiver<()>, blocked_write_tx: futures_mpsc::UnboundedSender, @@ -378,7 +399,7 @@ struct ControlledWebSocket { } struct ControlledWebSocketHandle { - inbound_tx: futures_mpsc::UnboundedSender>, + inbound_tx: futures_mpsc::UnboundedSender>, write_permit_tx: futures_mpsc::UnboundedSender<()>, blocked_write_rx: futures_mpsc::UnboundedReceiver, inbound_reads: Arc, @@ -452,7 +473,7 @@ impl ControlledWebSocketHandle { } impl Sink for ControlledWebSocket { - type Error = std::convert::Infallible; + type Error = WebSocketError; fn poll_ready(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { let this = self.get_mut(); @@ -497,7 +518,7 @@ impl Sink for ControlledWebSocket { } impl futures::Stream for ControlledWebSocket { - type Item = Result; + type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll> { let result = Pin::new(&mut self.inbound_rx).poll_next(cx); diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index a3fcd95a4a..4c878d6723 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -59,6 +59,8 @@ const NOISE_HANDSHAKE_FAILURE_COOLDOWN: Duration = Duration::from_secs(10); const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096; const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32; const HARNESS_KEY_VALIDATION_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE: usize = 100; +const HANDSHAKE_DIAGNOSTIC_WINDOW: Duration = Duration::from_secs(60); #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub(crate) enum RendezvousDisconnectReason { @@ -499,6 +501,7 @@ pub(crate) trait HarnessKeyValidator: Send + Sync { &self, harness_public_key: &NoiseChannelPublicKey, authorization: &str, + stream_id: Option, ) -> impl std::future::Future> + Send; } @@ -546,6 +549,42 @@ where let mut handshake_cooldown_until = None; let mut next_validation_id = 0u64; let mut disconnect_reason = RendezvousDisconnectReason::LocalShutdown; + let mut diagnostic_window = (Instant::now(), 0); + let mut record_diagnostic = + |stream_id: &str, reason, registry_status| match handshake_diagnostic_slot( + &mut diagnostic_window, + Instant::now(), + ) { + slot if slot < MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE => record_handshake_diagnostic( + &environment_id, + &executor_registration_id, + stream_id, + reason, + registry_status, + ), + MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE => { + let span = tracing::info_span!( + parent: None, + "codex.exec_server.noise_handshake_diagnostics_rate_limited", + noise_event = "handshake_diagnostics", + noise_outcome = "rate_limited", + noise_side = "executor", + environment_id, + executor_registration_id, + ); + let _entered = span.enter(); + warn!( + target: "codex_otel.log_only", + noise_event = "handshake_diagnostics", + noise_outcome = "rate_limited", + noise_side = "executor", + environment_id, + executor_registration_id, + "Noise executor handshake diagnostics rate limited" + ); + } + _ => {} + }; loop { // Registry calls run separately so a slow check does not block the relay. @@ -589,7 +628,7 @@ where else { continue; }; - if validation_result.result.is_err() { + if let Err((reason, registry_status)) = validation_result.result { // Validator errors may contain authorization details. warn!( noise_event = "authorization", @@ -601,6 +640,11 @@ where stream_id = validation_result.stream_id, "Noise harness authorization failure details" ); + record_diagnostic( + &validation_result.stream_id, + reason, + registry_status, + ); send_reset(&physical_outgoing_tx, validation_result.stream_id); record_failed_handshake( &mut failed_handshakes, @@ -610,6 +654,7 @@ where } if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { warn!("Noise relay has too many active streams"); + record_diagnostic(&validation_result.stream_id, "active_stream_limit", None); send_reset(&physical_outgoing_tx, validation_result.stream_id); continue; } @@ -621,6 +666,7 @@ where Ok(completed) => completed, Err(error) => { warn!("failed to complete Noise relay handshake: {error}"); + record_diagnostic(&validation_result.stream_id, "handshake_completion", None); send_reset(&physical_outgoing_tx, validation_result.stream_id); record_failed_handshake( &mut failed_handshakes, @@ -635,10 +681,12 @@ where ); // Do not leave a half-open stream if the handshake reply // cannot be queued immediately. - if physical_outgoing_tx - .try_send(encode_relay_message_frame(&response)) - .is_err() - { + if let Err(error) = physical_outgoing_tx.try_send(encode_relay_message_frame(&response)) { + let reason = match error { + mpsc::error::TrySendError::Full(_) => "response_queue_full", + mpsc::error::TrySendError::Closed(_) => "response_queue_closed", + }; + record_diagnostic(&validation_result.stream_id, reason, None); break; } info!( @@ -669,6 +717,7 @@ where let stream_ids = pending_handshakes.keys().cloned().collect::>(); pending_handshakes.clear(); for stream_id in stream_ids { + record_diagnostic(&stream_id, "validator_task_failed", None); send_reset(&physical_outgoing_tx, stream_id); } } @@ -719,6 +768,7 @@ where // Existing streams and already-admitted validations keep running. if let Some(deadline) = handshake_cooldown_until { if Instant::now() < deadline { + record_diagnostic(&stream_id, "cooldown", None); send_reset(&physical_outgoing_tx, stream_id); continue; } @@ -729,22 +779,26 @@ where // handshake. Malformed attempts that reach cryptography are // covered by the connection-wide failure budget below. if streams.contains_key(&stream_id) { + record_diagnostic(&stream_id, "duplicate_active_stream", None); send_reset(&physical_outgoing_tx, stream_id); continue; } // Removing pending state makes the in-flight validation result stale. if pending_handshakes.remove(&stream_id).is_some() { + record_diagnostic(&stream_id, "duplicate_pending_stream", None); send_reset(&physical_outgoing_tx, stream_id); record_failed_handshake(&mut failed_handshakes, &mut handshake_cooldown_until); continue; } if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { warn!("Noise relay has too many active streams"); + record_diagnostic(&stream_id, "active_stream_limit", None); send_reset(&physical_outgoing_tx, stream_id); continue; } if validation_tasks.len() >= MAX_PENDING_HANDSHAKE_VALIDATIONS { warn!("Noise relay has too many pending harness key validations"); + record_diagnostic(&stream_id, "pending_validation_limit", None); send_reset(&physical_outgoing_tx, stream_id); continue; } @@ -754,6 +808,7 @@ where Ok(request) => request, Err(error) => { warn!("failed to read Noise relay handshake frame: {error}"); + record_diagnostic(&stream_id, "handshake_frame", None); send_reset(&physical_outgoing_tx, stream_id); continue; } @@ -763,6 +818,7 @@ where Ok(pending) => pending, Err(error) => { warn!("failed to read Noise relay handshake request: {error}"); + record_diagnostic(&stream_id, "handshake_request", None); send_reset(&physical_outgoing_tx, stream_id); record_failed_handshake( &mut failed_handshakes, @@ -782,10 +838,12 @@ where } Ok(_) => { warn!("Noise relay handshake authorization is too long"); + record_diagnostic(&stream_id, "authorization_too_long", None); None } Err(_) => { warn!("Noise relay handshake authorization is not UTF-8"); + record_diagnostic(&stream_id, "authorization_not_utf8", None); None } }; @@ -811,14 +869,18 @@ where validation_tasks.spawn(async move { let result = match timeout( HARNESS_KEY_VALIDATION_TIMEOUT, - validator.validate_harness_key(&harness_public_key, &authorization), + validator.validate_harness_key( + &harness_public_key, + &authorization, + Uuid::parse_str(&stream_id).ok(), + ), ) .await { - Ok(result) => result, - Err(_) => Err(ExecServerError::Protocol( - "timed out validating Noise relay harness key".to_string(), - )), + Ok(result) => { + result.map_err(|error| classify_harness_key_validation_error(&error)) + } + Err(_) => Err(("validation_timeout", None)), }; HarnessKeyValidationResult { stream_id, @@ -832,6 +894,9 @@ where let Some(stream) = streams.get_mut(&stream_id) else { let canceled_pending_handshake = pending_handshakes.remove(&stream_id).is_some(); + if canceled_pending_handshake { + record_diagnostic(&stream_id, "data_before_handshake", None); + } send_reset(&physical_outgoing_tx, stream_id); if canceled_pending_handshake { record_failed_handshake( @@ -857,7 +922,11 @@ where } } RelayFrameBodyKind::Reset => { - pending_handshakes.remove(&stream_id); + if pending_handshakes.remove(&stream_id).is_some() { + // A reset can come from the harness or relay. Its reason is + // untrusted and does not tell us why the stream was closed. + record_diagnostic(&stream_id, "reset_pending_validation", None); + } if let Some(stream) = streams.remove(&stream_id) { // The reset reason is unauthenticated, so do not log it. stream.disconnect(); @@ -880,6 +949,62 @@ where disconnect_reason } +// Return the next slot in this window. The limit slot records a suppression +// marker; all later slots suppress diagnostics until the window resets. +fn handshake_diagnostic_slot(window: &mut (Instant, usize), now: Instant) -> usize { + if now.duration_since(window.0) >= HANDSHAKE_DIAGNOSTIC_WINDOW { + *window = (now, 0); + } + let slot = window.1; + window.1 = (slot + 1).min(MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE + 1); + slot +} + +/// Record bounded, non-secret context for an unfinished handshake. +fn record_handshake_diagnostic( + environment_id: &str, + executor_registration_id: &str, + stream_id: &str, + reason: &'static str, + registry_status: Option, +) { + // The peer controls stream_id. Never export unvalidated wire data. + let stream_id = Uuid::parse_str(stream_id) + .map(|id| id.to_string()) + .unwrap_or_else(|_| "invalid".to_string()); + let noise_outcome = if reason == "reset_pending_validation" { + "reset_received" + } else { + "rejected" + }; + // A separate span ends promptly even if the physical relay stays connected. + let span = tracing::info_span!( + parent: None, + "codex.exec_server.noise_handshake_diagnostic", + noise_event = "handshake", + noise_outcome, + noise_side = "executor", + environment_id, + executor_registration_id, + stream_id, + noise_reason = reason, + registry_status = ?registry_status, + ); + let _entered = span.enter(); + warn!( + target: "codex_otel.log_only", + noise_event = "handshake", + noise_outcome, + noise_side = "executor", + noise_reason = reason, + registry_status = ?registry_status, + environment_id, + executor_registration_id, + stream_id, + "Noise executor handshake diagnostic" + ); +} + /// Charge one failed authenticated-channel attempt to this physical relay. /// /// Pause new handshakes after a small fixed budget without disconnecting @@ -909,7 +1034,27 @@ struct PendingHandshake { struct HarnessKeyValidationResult { stream_id: String, validation_id: u64, - result: Result<(), ExecServerError>, + result: Result<(), (&'static str, Option)>, +} + +fn classify_harness_key_validation_error(error: &ExecServerError) -> (&'static str, Option) { + match error { + ExecServerError::ApplicationNetworkPolicy(_) + | ExecServerError::EnvironmentRegistryRequest( + codex_http_client::RouteAwareRequestError::Policy(_), + ) => ("network_policy", None), + ExecServerError::EnvironmentRegistryAuth(_) => ("registry_auth", None), + ExecServerError::EnvironmentRegistryHttp { status, .. } => { + ("registry_http", Some(status.as_u16())) + } + ExecServerError::EnvironmentRegistryRequest(error) if error.is_timeout() => { + ("registry_request_timeout", None) + } + ExecServerError::EnvironmentRegistryRequest(_) => ("registry_request", None), + ExecServerError::EnvironmentRegistryConfig(_) => ("registry_config", None), + ExecServerError::Protocol(_) => ("validation_protocol", None), + _ => ("validation_error", None), + } } /// Queue a best-effort reset without blocking the shared websocket loop. @@ -950,6 +1095,35 @@ mod tests { use super::*; + #[test] + fn handshake_diagnostics_are_capped_per_window() { + let start = Instant::now(); + let mut window = (start, 0); + for slot in 0..MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE { + assert_eq!(handshake_diagnostic_slot(&mut window, start), slot); + } + assert_eq!( + handshake_diagnostic_slot(&mut window, start), + MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE + ); + assert_eq!( + handshake_diagnostic_slot(&mut window, start), + MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE + 1 + ); + let next = start + HANDSHAKE_DIAGNOSTIC_WINDOW; + for slot in 0..MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE { + assert_eq!(handshake_diagnostic_slot(&mut window, next), slot); + } + assert_eq!( + handshake_diagnostic_slot(&mut window, next), + MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE + ); + assert_eq!( + handshake_diagnostic_slot(&mut window, next), + MAX_HANDSHAKE_DIAGNOSTICS_PER_MINUTE + 1 + ); + } + #[tokio::test] async fn harness_connection_sends_keepalive_and_receives_relay_data() -> anyhow::Result<()> { let (client_websocket, mut server_websocket) = websocket_pair().await?; @@ -1096,6 +1270,7 @@ mod tests { &self, _harness_public_key: &NoiseChannelPublicKey, _authorization: &str, + _stream_id: Option, ) -> Result<(), ExecServerError> { Ok(()) } diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs index 48b52d9798..a63005f664 100644 --- a/codex-rs/exec-server/src/relay_noise_tests.rs +++ b/codex-rs/exec-server/src/relay_noise_tests.rs @@ -1,3 +1,6 @@ +use std::io::Read; +use std::io::Seek; +use std::io::SeekFrom; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -9,6 +12,10 @@ use codex_exec_server_protocol::JSONRPCResponse; use codex_exec_server_protocol::RequestId; use futures::SinkExt; use futures::StreamExt; +use opentelemetry::trace::SpanId; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::InMemorySpanExporter; +use opentelemetry_sdk::trace::SdkTracerProvider; use pretty_assertions::assert_eq; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; @@ -20,7 +27,13 @@ use tokio_tungstenite::accept_async; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; use tokio_util::task::AbortOnDropHandle; +use tracing::instrument::WithSubscriber; +use tracing_subscriber::Layer; +use tracing_subscriber::filter::filter_fn; +use tracing_subscriber::layer::SubscriberExt; +use uuid::Uuid; +use super::HARNESS_KEY_VALIDATION_TIMEOUT; use super::HarnessKeyValidator; use super::MAX_FAILED_NOISE_HANDSHAKES; use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES; @@ -220,6 +233,7 @@ impl HarnessKeyValidator for BlockingValidator { &self, _harness_public_key: &NoiseChannelPublicKey, _authorization: &str, + _stream_id: Option, ) -> impl std::future::Future> + Send { let calls = Arc::clone(&self.calls); let release = Arc::clone(&self.release); @@ -231,6 +245,41 @@ impl HarnessKeyValidator for BlockingValidator { } } +#[derive(Clone)] +struct DiagnosticValidator { + blocking: BlockingValidator, + timeout_started: Arc, +} + +impl HarnessKeyValidator for DiagnosticValidator { + async fn validate_harness_key( + &self, + harness_public_key: &NoiseChannelPublicKey, + authorization: &str, + stream_id: Option, + ) -> Result<(), ExecServerError> { + match authorization { + "reset-authorization-secret" | "timeout-authorization-secret" => { + assert!(stream_id.is_some()); + self.timeout_started.notify_one(); + std::future::pending().await + } + "http-authorization-secret" => Err(ExecServerError::EnvironmentRegistryHttp { + status: http::StatusCode::SERVICE_UNAVAILABLE, + code: None, + message: "registry-response-secret".to_string(), + }), + _ => { + // The existing healthy streams use non-UUID identifiers. + assert_eq!(stream_id, None); + self.blocking + .validate_harness_key(harness_public_key, authorization, stream_id) + .await + } + } + } +} + #[tokio::test] async fn processor_exit_resets_noise_harness_stream() -> Result<()> { let (registration_tx, mut registration_rx) = tokio::sync::mpsc::channel(1); @@ -539,17 +588,44 @@ async fn malformed_handshakes_preserve_active_streams_and_admission_recovers() - let harness_identity = NoiseChannelIdentity::generate()?; let calls = Arc::new(AtomicUsize::new(0)); let release = Arc::new(Notify::new()); - let environment_task = AbortOnDropHandle::new(tokio::spawn(run_multiplexed_environment( - environment_websocket, - EchoMessages, - ENVIRONMENT_ID.to_string(), - EXECUTOR_REGISTRATION_ID.to_string(), - environment_identity.clone(), - BlockingValidator { - calls: Arc::clone(&calls), - release: Arc::clone(&release), - }, - ))); + let timeout_started = Arc::new(Notify::new()); + let mut logs = tempfile::tempfile()?; + let writer = logs.try_clone()?; + let span_exporter = InMemorySpanExporter::default(); + let tracer_provider = SdkTracerProvider::builder() + .with_simple_exporter(span_exporter.clone()) + .build(); + let subscriber = tracing_subscriber::registry() + .with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(move || writer.try_clone().expect("clone log file")) + .with_filter( + tracing_subscriber::filter::Targets::new().with_default(tracing::Level::WARN), + ), + ) + .with( + tracing_opentelemetry::layer() + .with_tracer(tracer_provider.tracer("exec-server-noise-test")) + .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)), + ); + let environment_task = AbortOnDropHandle::new(tokio::spawn( + run_multiplexed_environment( + environment_websocket, + EchoMessages, + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + environment_identity.clone(), + DiagnosticValidator { + blocking: BlockingValidator { + calls: Arc::clone(&calls), + release: Arc::clone(&release), + }, + timeout_started: Arc::clone(&timeout_started), + }, + ) + .with_subscriber(subscriber), + )); let healthy_stream_id = "healthy"; let prologue = @@ -578,14 +654,19 @@ async fn malformed_handshakes_preserve_active_streams_and_admission_recovers() - assert_eq!(response.stream_id, healthy_stream_id); let mut healthy_transport = handshake.finish(&response.into_handshake_payload()?)?; + let valid_stream_id = Uuid::new_v4().to_string(); for attempt in 0..MAX_FAILED_NOISE_HANDSHAKES { - let stream_id = format!("malformed-{attempt}"); + let stream_id = if attempt == 0 { + valid_stream_id.to_uppercase() + } else { + format!("malformed-{attempt}") + }; let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, &stream_id); let (_handshake, mut request) = InitiatorHandshake::start( &harness_identity, &environment_identity.public_key(), &prologue, - b"authorization", + b"handshake-authorization-secret", )?; let last_byte = request.last_mut().expect("handshake request is not empty"); *last_byte ^= 1; @@ -676,11 +757,120 @@ async fn malformed_handshakes_preserve_active_streams_and_admission_recovers() - new_handshake.finish(&response.into_handshake_payload()?)?; assert_eq!(calls.load(Ordering::SeqCst), 2); + // Exercise an incoming reset, an HTTP failure and a validation timeout. + let reset_stream_id = Uuid::new_v4().to_string(); + for (stream_id, authorization) in [ + (reset_stream_id.clone(), "reset-authorization-secret"), + (Uuid::new_v4().to_string(), "http-authorization-secret"), + (Uuid::new_v4().to_string(), "timeout-authorization-secret"), + ] { + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, &stream_id); + let (_handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + authorization.as_bytes(), + )?; + let frame = RelayMessageFrame::handshake(stream_id.clone(), request); + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + if authorization != "http-authorization-secret" { + timeout_started.notified().await; + } + if authorization == "reset-authorization-secret" { + let frame = RelayMessageFrame::reset(stream_id, "untrusted-reset-secret".to_string()); + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + continue; + } + if authorization == "timeout-authorization-secret" { + // Keep servicing pings while the real validation timeout runs. + tokio::time::resume(); + } + let reset = timeout( + HARNESS_KEY_VALIDATION_TIMEOUT + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(reset.stream_id, stream_id); + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); + } + + // Diagnostic spans must export without waiting for the relay to close. + assert!(!environment_task.is_finished()); + let spans = span_exporter.get_finished_spans()?; + assert!(spans.iter().any(|span| { + span.name == "codex.exec_server.noise_handshake_diagnostic" + && span.parent_span_id == SpanId::INVALID + && [ + ("noise_outcome", "reset_received"), + ("stream_id", reset_stream_id.as_str()), + ] + .iter() + .all(|(key, value)| { + span.attributes.iter().any(|attribute| { + attribute.key.as_str() == *key && attribute.value.to_string() == *value + }) + }) + })); + harness_websocket.close(None).await?; assert_eq!( timeout(Duration::from_secs(1), environment_task).await??, RendezvousDisconnectReason::PeerClose, ); + logs.seek(SeekFrom::Start(0))?; + let mut output = String::new(); + logs.read_to_string(&mut output)?; + let events: Vec<_> = output + .lines() + .filter(|line| line.contains("Noise executor handshake diagnostic")) + .collect(); + assert_eq!(events.len(), MAX_FAILED_NOISE_HANDSHAKES + 5); + let rejection = events + .iter() + .find(|line| line.contains(&valid_stream_id)) + .expect("correlated rejection"); + for field in [ + "noise_outcome=\"rejected\"", + "handshake_request", + ENVIRONMENT_ID, + EXECUTOR_REGISTRATION_ID, + ] { + assert!(rejection.contains(field)); + } + for (reason, detail) in [ + ("handshake_request", "stream_id=\"invalid\""), + ("registry_http", "Some(503)"), + ("validation_timeout", "None"), + ( + "reset_pending_validation", + "noise_outcome=\"reset_received\"", + ), + ] { + assert!( + events + .iter() + .any(|line| line.contains(reason) && line.contains(detail)) + ); + } + assert!( + events.iter().any( + |line| line.contains("reset_pending_validation") && line.contains(&reset_stream_id) + ) + ); + let exported = format!("{spans:?}\n{output}"); + for secret in [ + "malformed-1", + "after-cooldown", + "authorization-secret", + "registry-response-secret", + "untrusted-reset-secret", + ] { + assert!(!exported.contains(secret)); + } Ok(()) } diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index aeb4467000..90eea34806 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -408,6 +408,7 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { &self, harness_public_key: &NoiseChannelPublicKey, authorization: &str, + stream_id: Option, ) -> Result<(), ExecServerError> { let environment_id = &self.environment_id; let url = endpoint_url( @@ -420,15 +421,17 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { harness_key_authorization: authorization.to_string(), }; let (auth_headers, executor_credential) = self.client.resolve_auth().await?; - let response = self + let mut request = self .client .http_for_auth() .post(url) .headers(auth_headers) .headers(current_trace_context_headers()) - .json(&body) - .send() - .await?; + .json(&body); + if let Some(stream_id) = stream_id { + request = request.header("x-codex-noise-stream-id", stream_id.to_string()); + } + let response = request.send().await?; let status = response.status(); if !status.is_success() { // The request contains the short-lived authorization. Do not include @@ -1235,7 +1238,11 @@ mod tests { executor_registration_id: "registration".to_string(), }; let validation = validator - .validate_harness_key(&public_key, "harness-authorization") + .validate_harness_key( + &public_key, + "harness-authorization", + /*stream_id*/ None, + ) .await .expect_err("invalid credentials must be rejected"); assert!(matches!( @@ -1317,7 +1324,11 @@ mod tests { executor_registration_id: "registration".to_string(), }; let error = validator - .validate_harness_key(&public_key, "harness-authorization") + .validate_harness_key( + &public_key, + "harness-authorization", + /*stream_id*/ None, + ) .await .expect_err("validation must reject a non-loopback HTTP registry"); assert!(matches!( @@ -1400,7 +1411,11 @@ mod tests { environment_id: "environment-requested".to_string(), executor_registration_id: "registration".to_string(), } - .validate_harness_key(&public_key, "harness-authorization") + .validate_harness_key( + &public_key, + "harness-authorization", + /*stream_id*/ None, + ) .await .expect("validate directly"); registry.verify().await; @@ -1488,7 +1503,11 @@ mod tests { executor_registration_id: "registration".to_string(), }; let validation = validator - .validate_harness_key(&public_key, "harness-authorization") + .validate_harness_key( + &public_key, + "harness-authorization", + /*stream_id*/ None, + ) .await .expect_err("invalid validation response"); assert!(!format!("{validation:?} {validation}").contains(TOKEN)); diff --git a/codex-rs/exec-server/src/remote/noise_tests.rs b/codex-rs/exec-server/src/remote/noise_tests.rs index 0e9dfb15c2..d6d5057336 100644 --- a/codex-rs/exec-server/src/remote/noise_tests.rs +++ b/codex-rs/exec-server/src/remote/noise_tests.rs @@ -116,7 +116,7 @@ async fn registry_requests_do_not_log_sensitive_urls_or_response_headers() -> Re ) .set_body_json(response), ) - .expect(1) + .expect(if operation == "validate" { 2 } else { 1 }) .mount(&server) .await; } @@ -134,13 +134,51 @@ async fn registry_requests_do_not_log_sensitive_urls_or_response_headers() -> Re client .connect_environment("environment-requested", harness_public_key.clone()) .await?; - RegistryHarnessKeyValidator { + let validator = RegistryHarnessKeyValidator { client, environment_id: "environment-requested".to_string(), executor_registration_id: "registration-1".to_string(), - } - .validate_harness_key(&harness_public_key, HARNESS_KEY_AUTHORIZATION) - .await?; + }; + let stream_id = uuid::Uuid::new_v4(); + validator + .validate_harness_key( + &harness_public_key, + HARNESS_KEY_AUTHORIZATION, + Some(stream_id), + ) + .await?; + validator + .validate_harness_key( + &harness_public_key, + HARNESS_KEY_AUTHORIZATION, + /*stream_id*/ None, + ) + .await?; + + let requests = server.received_requests().await.expect("recorded requests"); + let validation_requests: Vec<_> = requests + .iter() + .filter(|request| { + request + .url + .query_pairs() + .any(|(key, value)| key == "registry_token" && value.ends_with("/validate")) + }) + .collect(); + assert_eq!(validation_requests.len(), 2); + assert_eq!( + validation_requests[0] + .headers + .get("x-codex-noise-stream-id") + .map(|value| value.to_str().unwrap()), + Some(stream_id.to_string().as_str()) + ); + assert!( + validation_requests[1] + .headers + .get("x-codex-noise-stream-id") + .is_none() + ); let logs = String::from_utf8(log_buffer.lock().expect("log buffer lock").clone())?; assert!(logs.contains("registry log capture sentinel")); @@ -305,7 +343,11 @@ async fn validate_harness_key_requires_explicit_valid_response() { environment_id: "environment-requested".to_string(), executor_registration_id: "registration-1".to_string(), } - .validate_harness_key(&harness_public_key, HARNESS_KEY_AUTHORIZATION) + .validate_harness_key( + &harness_public_key, + HARNESS_KEY_AUTHORIZATION, + /*stream_id*/ None, + ) .await .expect_err("a false validation response must fail closed"); @@ -335,7 +377,11 @@ async fn validate_harness_key_does_not_expose_error_body() { environment_id: "environment-requested".to_string(), executor_registration_id: "registration-1".to_string(), } - .validate_harness_key(&harness_public_key, HARNESS_KEY_AUTHORIZATION) + .validate_harness_key( + &harness_public_key, + HARNESS_KEY_AUTHORIZATION, + /*stream_id*/ None, + ) .await .expect_err("validation HTTP error should fail closed");