diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 40afbf0e28..3fc1324eeb 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -12,6 +12,7 @@ use prost::Message as ProstMessage; use tokio::sync::mpsc; use tokio::sync::watch; use tokio::task::JoinSet; +use tokio::time::Instant; use tokio::time::timeout; use tokio_tungstenite::tungstenite::Message; use tracing::debug; @@ -50,6 +51,7 @@ use crate::websocket_pong_watchdog::WebSocketPongWatchdog; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128; const MAX_FAILED_NOISE_HANDSHAKES: usize = 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); @@ -575,6 +577,7 @@ where let mut pending_handshakes: HashMap = HashMap::new(); let mut validation_tasks: JoinSet = JoinSet::new(); let mut failed_handshakes = 0usize; + let mut handshake_cooldown_until = None; let mut next_validation_id = 0u64; let mut disconnect_reason = RendezvousDisconnectReason::LocalShutdown; @@ -633,10 +636,10 @@ where "Noise harness authorization failure details" ); send_reset(&physical_outgoing_tx, validation_result.stream_id); - if failed_handshake_budget_exhausted(&mut failed_handshakes) { - warn!("closing Noise relay after repeated handshake failures"); - break; - } + record_failed_handshake( + &mut failed_handshakes, + &mut handshake_cooldown_until, + ); continue; } if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { @@ -653,10 +656,10 @@ where Err(error) => { warn!("failed to complete Noise relay handshake: {error}"); send_reset(&physical_outgoing_tx, validation_result.stream_id); - if failed_handshake_budget_exhausted(&mut failed_handshakes) { - warn!("closing Noise relay after repeated handshake failures"); - break; - } + record_failed_handshake( + &mut failed_handshakes, + &mut handshake_cooldown_until, + ); continue; } }; @@ -746,6 +749,16 @@ where let stream_id = frame.stream_id.clone(); match kind { RelayFrameBodyKind::Handshake => { + // Stop admitting work before parsing another hybrid handshake. + // Existing streams and already-admitted validations keep running. + if let Some(deadline) = handshake_cooldown_until { + if Instant::now() < deadline { + send_reset(&physical_outgoing_tx, stream_id); + continue; + } + failed_handshakes = 0; + handshake_cooldown_until = None; + } // Reject duplicate or busy streams before paying for a hybrid // handshake. Malformed attempts that reach cryptography are // covered by the connection-wide failure budget below. @@ -756,10 +769,7 @@ where // Removing pending state makes the in-flight validation result stale. if pending_handshakes.remove(&stream_id).is_some() { send_reset(&physical_outgoing_tx, stream_id); - if failed_handshake_budget_exhausted(&mut failed_handshakes) { - warn!("closing Noise relay after repeated handshake failures"); - break; - } + record_failed_handshake(&mut failed_handshakes, &mut handshake_cooldown_until); continue; } if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { @@ -788,10 +798,10 @@ where Err(error) => { warn!("failed to read Noise relay handshake request: {error}"); send_reset(&physical_outgoing_tx, stream_id); - if failed_handshake_budget_exhausted(&mut failed_handshakes) { - warn!("closing Noise relay after repeated handshake failures"); - break; - } + record_failed_handshake( + &mut failed_handshakes, + &mut handshake_cooldown_until, + ); continue; } }; @@ -815,10 +825,7 @@ where }; let Some(authorization) = authorization else { send_reset(&physical_outgoing_tx, stream_id); - if failed_handshake_budget_exhausted(&mut failed_handshakes) { - warn!("closing Noise relay after repeated handshake failures"); - break; - } + record_failed_handshake(&mut failed_handshakes, &mut handshake_cooldown_until); continue; }; let harness_public_key = pending.initiator_public_key.clone(); @@ -860,11 +867,11 @@ where let canceled_pending_handshake = pending_handshakes.remove(&stream_id).is_some(); send_reset(&physical_outgoing_tx, stream_id); - if canceled_pending_handshake - && failed_handshake_budget_exhausted(&mut failed_handshakes) - { - warn!("closing Noise relay after repeated handshake failures"); - break; + if canceled_pending_handshake { + record_failed_handshake( + &mut failed_handshakes, + &mut handshake_cooldown_until, + ); } continue; }; @@ -909,11 +916,21 @@ where /// Charge one failed authenticated-channel attempt to this physical relay. /// -/// Closing after a small fixed budget prevents a peer that has not been -/// authorized from triggering unbounded hybrid handshakes or registry checks. -fn failed_handshake_budget_exhausted(failed_handshakes: &mut usize) -> bool { +/// Pause new handshakes after a small fixed budget without disconnecting +/// authenticated streams. In-flight validation failures do not extend the +/// cooldown, so they cannot indefinitely prevent new streams from connecting. +fn record_failed_handshake( + failed_handshakes: &mut usize, + handshake_cooldown_until: &mut Option, +) { + if handshake_cooldown_until.is_some() { + return; + } *failed_handshakes += 1; - *failed_handshakes >= MAX_FAILED_NOISE_HANDSHAKES + if *failed_handshakes >= MAX_FAILED_NOISE_HANDSHAKES { + *handshake_cooldown_until = Some(Instant::now() + NOISE_HANDSHAKE_FAILURE_COOLDOWN); + warn!("pausing Noise relay handshakes after repeated failures"); + } } /// Responder state held while registry authorization is pending. diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs index b92d22b9bf..e066e1eb5e 100644 --- a/codex-rs/exec-server/src/relay_noise_tests.rs +++ b/codex-rs/exec-server/src/relay_noise_tests.rs @@ -10,9 +10,12 @@ use codex_exec_server_protocol::RequestId; use futures::SinkExt; use futures::StreamExt; use pretty_assertions::assert_eq; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; use tokio::net::TcpListener; use tokio::sync::Notify; use tokio::time::timeout; +use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; @@ -21,6 +24,7 @@ use tokio_util::task::AbortOnDropHandle; use super::HarnessKeyValidator; use super::MAX_FAILED_NOISE_HANDSHAKES; use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES; +use super::NOISE_HANDSHAKE_FAILURE_COOLDOWN; use super::RendezvousDisconnectReason; use super::run_multiplexed_environment; use crate::ExecServerError; @@ -31,6 +35,8 @@ use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; use crate::noise_channel::noise_channel_prologue; use crate::noise_relay::NoiseHarnessConnectionArgs; +use crate::noise_relay::message_framing::JsonRpcMessageDecoder; +use crate::noise_relay::message_framing::frame_jsonrpc_message; use crate::noise_relay::noise_harness_connection_from_websocket_with_readiness; use crate::noise_relay::stream_handler::NoiseOutboundMessage; use crate::noise_relay::stream_handler::NoiseStreamConnection; @@ -44,6 +50,51 @@ use crate::server::ConnectionProcessor; const ENVIRONMENT_ID: &str = "environment-1"; const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; +async fn next_relay_frame(websocket: &mut WebSocketStream) -> Result +where + S: AsyncRead + AsyncWrite + Unpin, +{ + loop { + match websocket.next().await { + Some(Ok(Message::Binary(payload))) => { + return Ok(decode_relay_message_frame(payload.as_ref())?); + } + Some(Ok(Message::Ping(payload))) => websocket.send(Message::Pong(payload)).await?, + Some(Ok(Message::Pong(_) | Message::Frame(_))) => {} + Some(Ok(message)) => anyhow::bail!("expected relay frame, got {message:?}"), + Some(Err(error)) => return Err(error.into()), + None => anyhow::bail!("environment closed before sending relay frame"), + } + } +} + +#[derive(Clone)] +struct EchoMessages; + +impl NoiseStreamHandler for EchoMessages { + type Incoming = JSONRPCMessage; + type Outgoing = JSONRPCMessage; + + fn decode(payload: bytes::Bytes) -> Result { + Ok(serde_json::from_slice(&payload)?) + } + + fn encode(message: Self::Outgoing) -> Result { + ConnectionProcessor::encode(message) + } + + async fn run_connection( + self, + mut connection: NoiseStreamConnection, + ) { + while let Some(message) = connection.incoming_rx.recv().await { + if connection.outgoing_tx.send(message).await.is_err() { + break; + } + } + } +} + #[derive(Clone)] struct ObservedRegistration( ConnectionProcessor, @@ -304,7 +355,7 @@ async fn pending_harness_key_validation_does_not_block_new_handshakes() -> Resul } #[tokio::test] -async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> { +async fn duplicate_handshakes_pause_admission_without_disconnecting() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let websocket_url = format!("ws://{}", listener.local_addr()?); let harness_connection = tokio::spawn(connect_async(websocket_url)); @@ -366,19 +417,11 @@ async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> { harness_websocket .send(Message::Binary(encoded.clone().into())) .await?; - let payload = timeout(Duration::from_secs(1), async { - loop { - match harness_websocket.next().await { - Some(Ok(Message::Binary(payload))) => break Ok(payload), - Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {} - Some(Ok(message)) => anyhow::bail!("expected reset frame, got {message:?}"), - Some(Err(error)) => break Err(error.into()), - None => anyhow::bail!("environment closed before sending reset"), - } - } - }) + let reset = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) .await??; - let reset = decode_relay_message_frame(payload.as_ref())?; assert_eq!(reset.stream_id, stream_id); assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); } @@ -392,10 +435,31 @@ async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> { } }) .await?; + harness_websocket + .send(Message::Binary(encoded.clone().into())) + .await?; + let reset = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); + // The next valid request is rejected before starting another registry check. harness_websocket .send(Message::Binary(encoded.into())) .await?; - timeout(Duration::from_secs(1), environment_task).await??; + let reset = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); + assert_eq!(calls.load(Ordering::SeqCst), MAX_FAILED_NOISE_HANDSHAKES); + harness_websocket.close(None).await?; + assert_eq!( + timeout(Duration::from_secs(1), environment_task).await??, + RendezvousDisconnectReason::PeerClose, + ); release.notify_waiters(); Ok(()) } @@ -457,7 +521,13 @@ async fn oversized_harness_authorization_is_rejected_before_validation() -> Resu } #[tokio::test] -async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()> { +async fn malformed_handshakes_preserve_active_streams_and_admission_recovers() -> Result<()> { + tokio::time::pause(); + // Socket I/O must not auto-advance the paused cooldown or keepalive clocks. + let (_clock_guard, clock_released) = std::sync::mpsc::channel::<()>(); + tokio::task::spawn_blocking(move || { + let _ = clock_released.recv(); + }); let listener = TcpListener::bind("127.0.0.1:0").await?; let websocket_url = format!("ws://{}", listener.local_addr()?); let harness_connection = tokio::spawn(connect_async(websocket_url)); @@ -467,20 +537,46 @@ async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()> let environment_identity = NoiseChannelIdentity::generate()?; let harness_identity = NoiseChannelIdentity::generate()?; - let environment_task = tokio::spawn(run_multiplexed_environment( + 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, - ConnectionProcessor::new(ExecServerRuntimePaths::new( - std::env::current_exe()?, - /*codex_linux_sandbox_exe*/ None, - )?), + EchoMessages, ENVIRONMENT_ID.to_string(), EXECUTOR_REGISTRATION_ID.to_string(), environment_identity.clone(), BlockingValidator { - calls: Arc::new(AtomicUsize::new(0)), - release: Arc::new(Notify::new()), + calls: Arc::clone(&calls), + release: Arc::clone(&release), }, - )); + ))); + + let healthy_stream_id = "healthy"; + let prologue = + noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, healthy_stream_id); + let (handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + b"authorization", + )?; + release.notify_one(); + harness_websocket + .send(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame::handshake( + healthy_stream_id.to_string(), + request, + )) + .into(), + )) + .await?; + let response = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(response.stream_id, healthy_stream_id); + let mut healthy_transport = handshake.finish(&response.into_handshake_payload()?)?; for attempt in 0..MAX_FAILED_NOISE_HANDSHAKES { let stream_id = format!("malformed-{attempt}"); @@ -497,14 +593,99 @@ async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()> harness_websocket .send(Message::Binary(encode_relay_message_frame(&frame).into())) .await?; + let reset = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!( + reset, + RelayMessageFrame::reset( + frame.stream_id, + crate::noise_relay::NOISE_RELAY_RESET_REASON.to_string(), + ) + ); } - timeout(Duration::from_secs(1), environment_task).await??; + let new_stream_id = "after-cooldown"; + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, new_stream_id); + let (new_handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + b"authorization", + )?; + let new_frame = RelayMessageFrame::handshake(new_stream_id.to_string(), request); + // Check both immediate rejection and a later attempt that must not extend + // the cooldown. Encrypted traffic on the healthy stream still round-trips. + for seq in 0..2 { + harness_websocket + .send(Message::Binary( + encode_relay_message_frame(&new_frame).into(), + )) + .await?; + let reset = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(reset.stream_id, new_stream_id); + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + let message = JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(i64::from(seq)), + result: serde_json::json!({"healthy": true}), + }); + let frame = RelayMessageFrame::data( + healthy_stream_id.to_string(), + seq, + healthy_transport.encrypt(&frame_jsonrpc_message(&message)?)?, + /*trace*/ None, + ); + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + let response = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(response.stream_id, healthy_stream_id); + let data = response.into_data()?; + assert_eq!(data.seq, seq); + let messages = + JsonRpcMessageDecoder::default().push(&healthy_transport.decrypt(&data.payload)?)?; + assert_eq!(messages, vec![message]); + + tokio::time::advance(NOISE_HANDSHAKE_FAILURE_COOLDOWN / 2).await; + } + + release.notify_one(); + harness_websocket + .send(Message::Binary( + encode_relay_message_frame(&new_frame).into(), + )) + .await?; + let response = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(response.stream_id, new_stream_id); + new_handshake.finish(&response.into_handshake_payload()?)?; + assert_eq!(calls.load(Ordering::SeqCst), 2); + + harness_websocket.close(None).await?; + assert_eq!( + timeout(Duration::from_secs(1), environment_task).await??, + RendezvousDisconnectReason::PeerClose, + ); Ok(()) } #[tokio::test] -async fn repeated_early_data_during_validation_closes_the_physical_relay() -> Result<()> { +async fn early_data_during_validation_pauses_admission_without_disconnecting() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let websocket_url = format!("ws://{}", listener.local_addr()?); let harness_connection = tokio::spawn(connect_async(websocket_url)); @@ -546,8 +727,18 @@ async fn repeated_early_data_during_validation_closes_the_physical_relay() -> Re .send(Message::Binary(encode_relay_message_frame(&frame).into())) .await?; } + let reset = timeout( + Duration::from_secs(1), + next_relay_frame(&mut harness_websocket), + ) + .await??; + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); } - timeout(Duration::from_secs(1), environment_task).await??; + harness_websocket.close(None).await?; + assert_eq!( + timeout(Duration::from_secs(1), environment_task).await??, + RendezvousDisconnectReason::PeerClose, + ); Ok(()) }