From e18a20d6eb6f5ede32a69fb56e3fe276d5e9d914 Mon Sep 17 00:00:00 2001 From: virajrai-oai Date: Tue, 20 Oct 2026 07:49:19 +0000 Subject: [PATCH] Add correlated diagnostics for Noise relay handshakes (#55498) ## What changed - Record executor handshake rejections and pending-validation resets in promptly completed diagnostic spans, with classified reasons and registry HTTP status when available. Normalize stream IDs as UUIDs and cap diagnostics at 100 per minute per relay connection, followed by a suppression marker. - Add elapsed time and connection identifiers to harness handshake diagnostics, including initialization timeouts, and classify WebSocket read failures from typed errors. - Forward valid stream UUIDs to harness key validation requests through the `x-codex-noise-stream-id` header. ## Testing Extend tests to cover diagnostic rate limits, timeout correlation, prompt span export, typed WebSocket error classification, and optional stream ID headers. Verify that diagnostic output excludes authorization secrets, registry response text, reset reasons, and invalid stream IDs. GitOrigin-RevId: 2c224e2d2828b300164868b10806add410e148b2 --- .../exec-server/src/client_refresh_tests.rs | 1 + codex-rs/exec-server/src/client_transport.rs | 20 +- .../exec-server/src/client_transport_tests.rs | 32 ++- .../exec-server/src/noise_relay/harness.rs | 49 +++- .../src/noise_relay/harness_tests.rs | 35 ++- codex-rs/exec-server/src/relay.rs | 199 +++++++++++++++- codex-rs/exec-server/src/relay_noise_tests.rs | 216 ++++++++++++++++-- codex-rs/exec-server/src/remote.rs | 35 ++- .../exec-server/src/remote/noise_tests.rs | 60 ++++- 9 files changed, 586 insertions(+), 61 deletions(-) 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");