openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add correlated diagnostics for Noise relay handshakes (#55498)

Browse Source 
## 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
This commit is contained in:
virajrai-oai
2026-10-20 07:56:07 +00:00
committed by copyberry
parent 014489466a
commit e18a20d6eb
9 changed files with 586 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/client_refresh_tests.rs
+1
View File
@@ -111,6 +111,7 @@ impl HarnessKeyValidator for Validator {
&self,
_: &NoiseChannelPublicKey,
_: &str,
_: Option<uuid::Uuid>,
) -> Result<(), ExecServerError> {
if let Some(handshake) = &self.handshake {
self.started.notify_one();
codex-rs/exec-server/src/client_transport.rs
+15 -5
View File
@@ -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,
codex-rs/exec-server/src/client_transport_tests.rs
+29 -3
View File
@@ -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(())
}
codex-rs/exec-server/src/noise_relay/harness.rs
+43 -6
View File
@@ -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<T, E>(
pub(crate) fn noise_harness_connection_from_websocket_with_readiness<T>(
stream: T,
args: NoiseHarnessConnectionArgs,
) -> NoiseHarnessConnection
where
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
T: Sink<Message, Error = WebSocketError>
+ Stream<Item = Result<Message, WebSocketError>>
+ 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",
}
}
codex-rs/exec-server/src/noise_relay/harness_tests.rs
+28 -7
View File
@@ -29,13 +29,16 @@ use super::*;
use crate::connection::JsonRpcConnectionEvent;
use crate::noise_channel::PendingResponderHandshake;
fn noise_harness_connection_from_websocket<T, E>(
fn noise_harness_connection_from_websocket<T>(
stream: T,
args: NoiseHarnessConnectionArgs,
) -> JsonRpcConnection
where
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
T: Sink<Message, Error = WebSocketError>
+ Stream<Item = Result<Message, WebSocketError>>
+ 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<Result<Message, std::convert::Infallible>>,
inbound_rx: futures_mpsc::UnboundedReceiver<Result<Message, WebSocketError>>,
outbound_tx: futures_mpsc::UnboundedSender<Message>,
write_permit_rx: futures_mpsc::UnboundedReceiver<()>,
blocked_write_tx: futures_mpsc::UnboundedSender<usize>,
@@ -378,7 +399,7 @@ struct ControlledWebSocket {
}
struct ControlledWebSocketHandle {
inbound_tx: futures_mpsc::UnboundedSender<Result<Message, std::convert::Infallible>>,
inbound_tx: futures_mpsc::UnboundedSender<Result<Message, WebSocketError>>,
write_permit_tx: futures_mpsc::UnboundedSender<()>,
blocked_write_rx: futures_mpsc::UnboundedReceiver<usize>,
inbound_reads: Arc<AtomicUsize>,
@@ -452,7 +473,7 @@ impl ControlledWebSocketHandle {
}
impl Sink<Message> for ControlledWebSocket {
type Error = std::convert::Infallible;
type Error = WebSocketError;
fn poll_ready(self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Result<(), Self::Error>> {
let this = self.get_mut();
@@ -497,7 +518,7 @@ impl Sink<Message> for ControlledWebSocket {
}
impl futures::Stream for ControlledWebSocket {
type Item = Result<Message, std::convert::Infallible>;
type Item = Result<Message, WebSocketError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut TaskContext<'_>) -> Poll<Option<Self::Item>> {
let result = Pin::new(&mut self.inbound_rx).poll_next(cx);
codex-rs/exec-server/src/relay.rs
+187 -12
View File
@@ -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<Uuid>,
) -> impl std::future::Future<Output = Result<(), ExecServerError>> + 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::<Vec<_>>();
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<u16>,
) {
// 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<u16>)>,
}
fn classify_harness_key_validation_error(error: &ExecServerError) -> (&'static str, Option<u16>) {
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<Uuid>,
) -> Result<(), ExecServerError> {
Ok(())
}
codex-rs/exec-server/src/relay_noise_tests.rs
+203 -13
View File
@@ -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<Uuid>,
) -> impl std::future::Future<Output = Result<(), ExecServerError>> + 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<Notify>,
}
impl HarnessKeyValidator for DiagnosticValidator {
async fn validate_harness_key(
&self,
harness_public_key: &NoiseChannelPublicKey,
authorization: &str,
stream_id: Option<Uuid>,
) -> 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(())
}
codex-rs/exec-server/src/remote.rs
+27 -8
View File
@@ -408,6 +408,7 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator {
&self,
harness_public_key: &NoiseChannelPublicKey,
authorization: &str,
stream_id: Option<uuid::Uuid>,
) -> 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));
codex-rs/exec-server/src/remote/noise_tests.rs
+53 -7
View File
@@ -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");