From e99be8cf65d41392d8b9320724965fae7d71a133 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Tue, 7 Jul 2026 12:04:28 -0700 Subject: [PATCH] feat(exec-server): add ACK-based Noise relay retries --- codex-rs/exec-server/README.md | 23 +- .../src/noise_relay/executor_stream.rs | 211 +++++++++++++++--- .../src/noise_relay/executor_stream_tests.rs | 77 +++++-- .../exec-server/src/noise_relay/harness.rs | 139 +++++++++--- .../src/noise_relay/harness_tests.rs | 189 +++++++++++++--- codex-rs/exec-server/src/noise_relay/mod.rs | 13 +- .../src/noise_relay/ordered_ciphertext.rs | 28 ++- .../noise_relay/ordered_ciphertext_tests.rs | 31 ++- .../src/noise_relay/reliable_stream.rs | 189 ++++++++++++++++ .../src/noise_relay/reliable_stream_tests.rs | 120 ++++++++++ codex-rs/exec-server/src/relay.rs | 107 ++++++++- codex-rs/exec-server/src/relay_proto.rs | 1 + codex-rs/exec-server/tests/relay.rs | 69 +++++- 13 files changed, 1048 insertions(+), 149 deletions(-) create mode 100644 codex-rs/exec-server/src/noise_relay/reliable_stream.rs create mode 100644 codex-rs/exec-server/src/noise_relay/reliable_stream_tests.rs diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index a250ba5886..f362d8db02 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -60,8 +60,8 @@ identity plus endpoint-owned reliability metadata: version stream_id body // handshake | data | ack_frame | resume | reset | heartbeat -ack // highest contiguous peer segment seq received -ack_bits // bitset for peer segment seqs after ack +ack // highest contiguous peer segment seq received; 0 means none +ack_bits // reserved in v1; must be 0 seq // data only: segment sequence number segment_index // data only: 0-based index within message segment_count // data only: number of segments in message @@ -74,13 +74,18 @@ reason // reset only: reset reason environment websocket. The harness generates a UUIDv4 `stream_id`; the environment demuxes frames by `stream_id` and runs an independent `ConnectionProcessor` per stream. +After a Noise session ends, do not reuse its `stream_id` on the same physical +relay connection; delayed cached ciphertext may still be draining. Use segment-level sequence numbers for reliability: ```text -seq = 0, 1, 2, 3, ... +seq = 1, 2, 3, 4, ... ``` +Sequence zero is reserved so `ack = 0` unambiguously means that no segment has +been received contiguously yet. + Use contiguous segment sequence ranges to identify and stitch a segmented application message: @@ -94,15 +99,17 @@ segment_count = 1 unsplit messages, `message_start_seq == seq`, `segment_index == 0`, and `segment_count == 1`. -Use cumulative `ack` plus fixed-size `ack_bits` instead of variable ack ranges: +V1 uses cumulative `ack` and reserves `ack_bits` for a later selective-ack +extension: ```text -ack = highest contiguous received segment seq -bit i in ack_bits acknowledges seq = ack + 1 + i +ack = 0 when no segment has been received contiguously +ack = highest contiguous received segment seq otherwise +ack_bits = 0 ``` -Send `ack` and `ack_bits` redundantly on every outbound frame. Acks are not -themselves acked. Acks, retries, duplicate suppression, segmentation, and +Send `ack` and zero `ack_bits` redundantly on every outbound frame. Acks are +not themselves acked. Acks, retries, duplicate suppression, segmentation, and reassembly are endpoint responsibilities; rendezvous only routes relay frames by `stream_id`. diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 89f33d938a..82df5b1d3a 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -6,7 +6,11 @@ use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicU32; +use std::sync::atomic::Ordering; +use std::time::Duration; +use tokio::sync::Notify; use tokio::sync::mpsc; use tokio::sync::watch; use tracing::warn; @@ -22,18 +26,23 @@ use crate::noise_relay::message_framing::JsonRpcMessageDecoder; use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN; use crate::noise_relay::message_framing::frame_jsonrpc_message; use crate::noise_relay::ordered_ciphertext::OrderedCiphertextFrames; -use crate::noise_relay::take_next_sequence; +use crate::noise_relay::reliable_stream::MAX_UNACKED_BYTES; +use crate::noise_relay::reliable_stream::MAX_UNACKED_SEGMENTS; +use crate::noise_relay::reliable_stream::ReliableSender; use crate::relay::encode_relay_message_frame; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; use crate::server::ConnectionProcessor; use crate::telemetry::ConnectionTransport; +const RELIABLE_RETRY_SCAN_INTERVAL: Duration = Duration::from_millis(50); +const MAX_RELIABLE_CIPHERTEXT_BYTES: usize = MAX_UNACKED_BYTES / MAX_UNACKED_SEGMENTS; + /// Identifies one completed virtual-stream instance. /// -/// Stream IDs are supplied by the untrusted relay peer and may be reused. The -/// instance ID prevents a delayed writer notification from removing a newer -/// stream that happens to use the same routing ID. +/// Stream IDs are supplied by the untrusted relay peer. The instance ID lets +/// the environment verify that a delayed writer notification still belongs to +/// the active stream before retiring its routing ID. pub(crate) struct ClosedNoiseVirtualStream { pub(crate) stream_id: String, pub(crate) instance_id: u64, @@ -48,6 +57,10 @@ pub(crate) struct NoiseVirtualStream { incoming_tx: mpsc::Sender, disconnected_tx: watch::Sender, transport: Arc>, + reliable_sender: Arc>, + inbound_ack: Arc, + pending_ack: Arc>>, + writer_wakeup: Arc, inbound_ciphertexts: OrderedCiphertextFrames, inbound_decoder: JsonRpcMessageDecoder, pub(crate) instance_id: u64, @@ -61,6 +74,19 @@ impl NoiseVirtualStream { .try_send(JsonRpcConnectionEvent::Disconnected { reason }); } + /// Apply ack metadata from one post-handshake peer frame and wake the + /// writer if the cumulative prefix opened send-window capacity. + pub(crate) fn process_peer_ack(&self, ack: u32, ack_bits: u32) -> Result<(), ExecServerError> { + let mut reliable_sender = self + .reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + reliable_sender.process_peer_ack(ack, ack_bits)?; + drop(reliable_sender); + self.writer_wakeup.notify_one(); + Ok(()) + } + /// Reorder and decrypt one inbound record, then queue complete JSON-RPC messages. /// This must stay nonblocking because all virtual streams share the read loop. pub(crate) fn receive_data(&mut self, data: RelayData) -> Result<(), ExecServerError> { @@ -84,6 +110,13 @@ impl NoiseVirtualStream { })?; } } + let ack = self.inbound_ciphertexts.cumulative_ack(); + self.inbound_ack.store(ack, Ordering::Relaxed); + *self + .pending_ack + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(ack); + self.writer_wakeup.notify_one(); Ok(()) } } @@ -104,50 +137,150 @@ pub(crate) fn spawn_noise_virtual_stream( let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); let (disconnected_tx, disconnected_rx) = watch::channel(false); let transport = Arc::new(Mutex::new(transport)); + let reliable_sender = Arc::new(Mutex::new(ReliableSender::default())); + let inbound_ack = Arc::new(AtomicU32::new(0)); + let pending_ack = Arc::new(Mutex::new(None)); + let writer_wakeup = Arc::new(Notify::new()); let writer_transport = Arc::clone(&transport); + let writer_reliable_sender = Arc::clone(&reliable_sender); + let writer_inbound_ack = Arc::clone(&inbound_ack); + let writer_pending_ack = Arc::clone(&pending_ack); + let writer_wakeup_task = Arc::clone(&writer_wakeup); + let writer_physical_outgoing_tx = physical_outgoing_tx; let processor_stream_id = stream_id.clone(); let processor_closed_stream_tx = closed_stream_tx.clone(); let writer_stream_id = stream_id; let writer_task = tokio::spawn(async move { - let mut next_seq = 0u32; - 'writer: while let Some(message) = json_outgoing_rx.recv().await { - // Each chunk becomes one Noise record and consumes one nonce. - let framed = match frame_jsonrpc_message(&message) { - Ok(framed) => framed, - Err(error) => { - warn!("failed to frame Noise virtual stream JSON-RPC payload: {error}"); - break; + let mut pending_outbound: Option<(Vec, usize)> = None; + let mut retry_tick = tokio::time::interval_at( + tokio::time::Instant::now() + RELIABLE_RETRY_SCAN_INTERVAL, + RELIABLE_RETRY_SCAN_INTERVAL, + ); + retry_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + 'writer: loop { + let can_send_pending = pending_outbound.is_some() + && writer_reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .can_admit_ciphertext(MAX_RELIABLE_CIPHERTEXT_BYTES); + let has_pending_ack = writer_pending_ack + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_some(); + tokio::select! { + maybe_message = json_outgoing_rx.recv(), if pending_outbound.is_none() => { + let Some(message) = maybe_message else { + break; + }; + pending_outbound = Some(match frame_jsonrpc_message(&message) { + Ok(framed) => (framed, 0), + Err(error) => { + warn!("failed to frame Noise virtual stream JSON-RPC payload: {error}"); + break; + } + }); } - }; - for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) { - let seq = match take_next_sequence(&mut next_seq) { - Ok(seq) => seq, - Err(error) => { - warn!("Noise virtual stream sequence exhausted: {error}"); + _ = std::future::ready(()), if can_send_pending => { + let (ciphertext, next_offset, message_complete) = { + let Some((framed, offset)) = pending_outbound.as_ref() else { + continue; + }; + let next_offset = (*offset + NOISE_RECORD_PLAINTEXT_LEN).min(framed.len()); + let ciphertext = { + let mut transport = writer_transport + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + transport.encrypt(&framed[*offset..next_offset]) + }; + let ciphertext = match ciphertext { + Ok(ciphertext) => ciphertext, + Err(error) => { + warn!("failed to encrypt Noise virtual stream payload: {error}"); + break 'writer; + } + }; + if ciphertext.len() > MAX_RELIABLE_CIPHERTEXT_BYTES { + warn!("Noise virtual stream ciphertext exceeds reliable record budget"); + break 'writer; + } + (ciphertext, next_offset, next_offset == framed.len()) + }; + let (outbound, ack) = { + let mut reliable_sender = writer_reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let outbound = match reliable_sender + .admit_ciphertext(ciphertext, tokio::time::Instant::now()) + { + Ok(outbound) => outbound, + Err(error) => { + warn!("failed to admit Noise reliable ciphertext: {error}"); + break 'writer; + } + }; + (outbound, writer_inbound_ack.load(Ordering::Relaxed)) + }; + let frame = RelayMessageFrame::reliable_data( + writer_stream_id.clone(), + ack, + outbound.seq, + outbound.payload, + ); + if writer_physical_outgoing_tx + .send(encode_relay_message_frame(&frame)) + .await + .is_err() + { break 'writer; } - }; - let ciphertext = { - let mut transport = writer_transport + if message_complete { + pending_outbound = None; + } else if let Some((_framed, offset)) = pending_outbound.as_mut() { + *offset = next_offset; + } + } + _ = retry_tick.tick() => { + let (retry, ack) = { + let mut reliable_sender = writer_reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let retry = reliable_sender.next_retry_due(tokio::time::Instant::now()); + (retry, writer_inbound_ack.load(Ordering::Relaxed)) + }; + if let Some(outbound) = retry { + let frame = RelayMessageFrame::reliable_data( + writer_stream_id.clone(), + ack, + outbound.seq, + outbound.payload, + ); + if writer_physical_outgoing_tx + .send(encode_relay_message_frame(&frame)) + .await + .is_err() + { + break 'writer; + } + } + } + _ = std::future::ready(()), if has_pending_ack => { + let ack = writer_pending_ack .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - transport.encrypt(plaintext_record) - }; - let ciphertext = match ciphertext { - Ok(ciphertext) => ciphertext, - Err(error) => { - warn!("failed to encrypt Noise virtual stream payload: {error}"); + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + let Some(ack) = ack else { + continue; + }; + let frame = RelayMessageFrame::ack(writer_stream_id.clone(), ack); + if writer_physical_outgoing_tx + .send(encode_relay_message_frame(&frame)) + .await + .is_err() + { break 'writer; } - }; - let frame = RelayMessageFrame::data(writer_stream_id.clone(), seq, ciphertext); - if physical_outgoing_tx - .send(encode_relay_message_frame(&frame)) - .await - .is_err() - { - break 'writer; } + _ = writer_wakeup_task.notified() => {} } } @@ -158,7 +291,7 @@ pub(crate) fn spawn_noise_virtual_stream( }; let reset = RelayMessageFrame::reset(writer_stream_id, NOISE_RELAY_RESET_REASON.to_string()); - let _ = physical_outgoing_tx.try_send(encode_relay_message_frame(&reset)); + let _ = writer_physical_outgoing_tx.try_send(encode_relay_message_frame(&reset)); let _ = closed_stream_tx.send(closed_stream).await; }); @@ -185,6 +318,10 @@ pub(crate) fn spawn_noise_virtual_stream( incoming_tx, disconnected_tx, transport, + reliable_sender, + inbound_ack, + pending_ack, + writer_wakeup, inbound_ciphertexts: OrderedCiphertextFrames::default(), inbound_decoder: JsonRpcMessageDecoder::default(), instance_id, diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs index 232dfa479e..cdf39f688f 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs @@ -1,3 +1,4 @@ +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Result; @@ -13,25 +14,17 @@ use crate::ExecServerRuntimePaths; use crate::connection::CHANNEL_CAPACITY; use crate::noise_channel::InitiatorHandshake; use crate::noise_channel::NoiseChannelIdentity; +use crate::noise_channel::NoiseTransport; use crate::noise_channel::PendingResponderHandshake; use crate::noise_relay::message_framing::frame_jsonrpc_message; +use crate::relay::RelayFrameBodyKind; +use crate::relay::decode_relay_message_frame; use crate::relay_proto::RelayData; use crate::server::ConnectionProcessor; #[tokio::test] async fn processor_exit_reports_closed_virtual_stream() -> Result<()> { - let executor_identity = NoiseChannelIdentity::generate()?; - let harness_identity = NoiseChannelIdentity::generate()?; - let prologue = b"test-prologue"; - let (initiator, request) = InitiatorHandshake::start( - &harness_identity, - &executor_identity.public_key(), - prologue, - b"authorization", - )?; - let pending = PendingResponderHandshake::read_request(&executor_identity, prologue, &request)?; - let (executor_transport, response) = pending.complete()?; - let mut harness_transport = initiator.finish(&response)?; + let (executor_transport, mut harness_transport) = completed_handshake()?; let (physical_outgoing_tx, _physical_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (closed_stream_tx, mut closed_stream_rx) = mpsc::channel(1); @@ -53,7 +46,7 @@ async fn processor_exit_reports_closed_virtual_stream() -> Result<()> { }); let ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message)?)?; stream.receive_data(RelayData { - seq: 0, + seq: 1, segment_index: 0, segment_count: 1, payload: ciphertext, @@ -68,3 +61,61 @@ async fn processor_exit_reports_closed_virtual_stream() -> Result<()> { )); Ok(()) } + +#[tokio::test] +async fn full_physical_queue_defers_ack_without_resetting_stream() -> Result<()> { + let (executor_transport, mut harness_transport) = completed_handshake()?; + let (physical_outgoing_tx, mut physical_outgoing_rx) = mpsc::channel(1); + physical_outgoing_tx.send(vec![0x5a]).await?; + let (closed_stream_tx, _closed_stream_rx) = mpsc::channel(1); + let mut stream = spawn_noise_virtual_stream( + "stream-1".to_string(), + /*instance_id*/ 7, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + physical_outgoing_tx, + closed_stream_tx, + executor_transport, + ); + + let message = JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result: serde_json::Value::Null, + }); + let framed = frame_jsonrpc_message(&message)?; + let ciphertext = harness_transport.encrypt(&framed[..1])?; + stream.receive_data(RelayData { + seq: 1, + segment_index: 0, + segment_count: 1, + payload: ciphertext, + })?; + + assert_eq!(stream.inbound_ack.load(Ordering::Relaxed), 1); + assert_eq!(physical_outgoing_rx.try_recv()?, vec![0x5a]); + let ack = timeout(Duration::from_secs(1), physical_outgoing_rx.recv()) + .await? + .expect("virtual stream writer should send the deferred ack"); + let ack = decode_relay_message_frame(&ack)?; + assert_eq!(ack.validate()?, RelayFrameBodyKind::Ack); + assert_eq!(ack.ack, 1); + Ok(()) +} + +fn completed_handshake() -> Result<(NoiseTransport, NoiseTransport)> { + let executor_identity = NoiseChannelIdentity::generate()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let prologue = b"test-prologue"; + let (initiator, request) = InitiatorHandshake::start( + &harness_identity, + &executor_identity.public_key(), + prologue, + b"authorization", + )?; + let pending = PendingResponderHandshake::read_request(&executor_identity, prologue, &request)?; + let (executor_transport, response) = pending.complete()?; + let harness_transport = initiator.finish(&response)?; + Ok((executor_transport, harness_transport)) +} diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 9969a3e130..6c5dc52611 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -6,6 +6,8 @@ //! normal `JsonRpcConnection`. Outbound JSON-RPC is framed and split into Noise //! records; inbound records are reordered before decryption and reassembly. +use std::time::Duration; + use futures::FutureExt; use futures::Sink; use futures::SinkExt; @@ -35,7 +37,9 @@ use crate::noise_relay::message_framing::JsonRpcMessageDecoder; use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN; use crate::noise_relay::message_framing::frame_jsonrpc_message; use crate::noise_relay::ordered_ciphertext::OrderedCiphertextFrames; -use crate::noise_relay::take_next_sequence; +use crate::noise_relay::reliable_stream::MAX_UNACKED_BYTES; +use crate::noise_relay::reliable_stream::MAX_UNACKED_SEGMENTS; +use crate::noise_relay::reliable_stream::ReliableSender; use crate::relay::RelayFrameBodyKind; use crate::relay::decode_relay_message_frame; use crate::relay::encode_relay_message_frame; @@ -65,6 +69,12 @@ pub(crate) struct NoiseHarnessConnectionArgs { const NOISE_RELAY_RESET_DISCONNECT_REASON: &str = "Noise relay stream reset"; // Give a Pong already queued behind data a bounded chance to reach the reader. const MAX_FRAMES_DRAINED_AFTER_PONG_DEADLINE: usize = 32; +// Poll frequently enough that 500ms resend deadlines do not wait for another +// application or websocket event, while still bounding retry wakeups. +const RELIABLE_RETRY_SCAN_INTERVAL: Duration = Duration::from_millis(50); +// The current Noise plaintext record cap plus transport overhead fits under one +// 64KiB slot. Check this conservative slot before consuming another send nonce. +const MAX_RELIABLE_CIPHERTEXT_BYTES: usize = MAX_UNACKED_BYTES / MAX_UNACKED_SEGMENTS; /// Adapt one harness rendezvous websocket into an authenticated JSON-RPC connection. /// @@ -124,21 +134,16 @@ where } }; - // Resume claims the stream ID at rendezvous; Handshake carries the - // opaque first IK message. No JSON-RPC data is sent before the - // responder proves possession of the pinned static key. - let resume = RelayMessageFrame::resume(stream_id.clone()); + // Handshake carries the opaque first IK message. New logical streams do + // not claim resume state before the responder proves possession of the + // pinned static key. let handshake = RelayMessageFrame::handshake(stream_id.clone(), request); if websocket - .send(Message::Binary(encode_relay_message_frame(&resume).into())) + .send(Message::Binary( + encode_relay_message_frame(&handshake).into(), + )) .await .is_err() - || websocket - .send(Message::Binary( - encode_relay_message_frame(&handshake).into(), - )) - .await - .is_err() { let _ = disconnected_tx.send(true); return; @@ -261,7 +266,7 @@ where // transport record. Outbound records are encrypted once; inbound // records are reordered and deduplicated before decryption. let mut websocket = websocket.peekable(); - let mut next_outbound_seq = 0u32; + let mut reliable_sender = ReliableSender::default(); let mut inbound_ciphertexts = OrderedCiphertextFrames::default(); let mut inbound_decoder = JsonRpcMessageDecoder::default(); let mut keepalive = tokio::time::interval_at( @@ -269,6 +274,11 @@ where WEBSOCKET_KEEPALIVE_INTERVAL, ); keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut retry_tick = tokio::time::interval_at( + tokio::time::Instant::now() + RELIABLE_RETRY_SCAN_INTERVAL, + RELIABLE_RETRY_SCAN_INTERVAL, + ); + retry_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut pong_watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT); let pong_deadline = tokio::time::sleep(WEBSOCKET_PONG_TIMEOUT); tokio::pin!(pong_deadline); @@ -276,9 +286,12 @@ where // creates a scheduling point for keepalive and inbound control frames // without splitting the WebSocket reader and writer. let mut pending_outbound: Option<(Vec, usize)> = None; + let mut pending_ack = None; let mut force_incoming = false; let mut frames_drained_after_pong_deadline = 0usize; 'relay: loop { + let can_send_pending = pending_outbound.is_some() + && reliable_sender.can_admit_ciphertext(MAX_RELIABLE_CIPHERTEXT_BYTES); // Consume a due tick before the always-ready record arm below can win // another select iteration and postpone the keepalive. if pong_watchdog.deadline().is_none() @@ -327,10 +340,10 @@ where } // While a Pong is outstanding, drain already-queued inbound traffic - // before the next fragment so a queued Pong cannot sit behind writes. + // before any write so a queued Pong cannot sit behind fragments, + // retries, or standalone acknowledgements. if !force_incoming && pong_watchdog.deadline().is_some() - && pending_outbound.is_some() && std::pin::Pin::new(&mut websocket) .peek() .now_or_never() @@ -352,14 +365,7 @@ where } }); } - _ = std::future::ready(()), if pending_outbound.is_some() && !force_incoming && !pong_deadline_expired => { - let seq = match take_next_sequence(&mut next_outbound_seq) { - Ok(seq) => seq, - Err(error) => { - warn!("Noise relay sequence exhausted: {error}"); - break 'relay; - } - }; + _ = std::future::ready(()), if can_send_pending && !force_incoming && !pong_deadline_expired => { let (ciphertext, next_offset, message_complete) = { let Some((framed, offset)) = pending_outbound.as_ref() else { continue; @@ -372,9 +378,27 @@ where break 'relay; } }; + if ciphertext.len() > MAX_RELIABLE_CIPHERTEXT_BYTES { + warn!("Noise relay ciphertext exceeds reliable record budget"); + break 'relay; + } (ciphertext, next_offset, next_offset == framed.len()) }; - let frame = RelayMessageFrame::data(stream_id.clone(), seq, ciphertext); + let outbound = match reliable_sender + .admit_ciphertext(ciphertext, tokio::time::Instant::now()) + { + Ok(outbound) => outbound, + Err(error) => { + warn!("failed to admit Noise reliable ciphertext: {error}"); + break 'relay; + } + }; + let frame = RelayMessageFrame::reliable_data( + stream_id.clone(), + inbound_ciphertexts.cumulative_ack(), + outbound.seq, + outbound.payload, + ); // A Pong can arrive after the readiness check while this write owns the // combined sink and stream. A single bounded record can therefore hit the // deadline and disconnect with that Pong queued. Treat that as write @@ -395,6 +419,42 @@ where *offset = next_offset; } } + _ = retry_tick.tick(), if !force_incoming && !pong_deadline_expired => { + if let Some(outbound) = reliable_sender.next_retry_due(tokio::time::Instant::now()) { + let frame = RelayMessageFrame::reliable_data( + stream_id.clone(), + inbound_ciphertexts.cumulative_ack(), + outbound.seq, + outbound.payload, + ); + if let Err(error) = send_websocket_message( + &mut websocket, + Message::Binary(encode_relay_message_frame(&frame).into()), + pong_watchdog.write_deadline(tokio::time::Instant::now()), + ) + .await + { + warn!("failed to retry Noise relay websocket frame: {error}"); + break 'relay; + } + } + } + _ = std::future::ready(()), if pending_ack.is_some() && !force_incoming && !pong_deadline_expired => { + let Some(ack) = pending_ack.take() else { + continue; + }; + let frame = RelayMessageFrame::ack(stream_id.clone(), ack); + if let Err(error) = send_websocket_message( + &mut websocket, + Message::Binary(encode_relay_message_frame(&frame).into()), + pong_watchdog.write_deadline(tokio::time::Instant::now()), + ) + .await + { + warn!("failed to write Noise relay ack: {error}"); + break 'relay; + } + } _ = &mut pong_deadline, if pong_watchdog.deadline().is_some() && !force_incoming => { continue; } @@ -437,8 +497,22 @@ where if frame.stream_id != stream_id { continue; } - match frame.validate() { - Ok(RelayFrameBodyKind::Data) => { + let kind = match frame.validate() { + Ok(kind) => kind, + Err(error) => { + send_malformed(&incoming_tx, error.to_string()); + break; + } + }; + if !matches!(kind, RelayFrameBodyKind::Handshake) + && let Err(error) = + reliable_sender.process_peer_ack(frame.ack, frame.ack_bits) + { + send_malformed(&incoming_tx, error.to_string()); + break; + } + match kind { + RelayFrameBodyKind::Data => { let data = match frame.into_data() { Ok(data) => data, Err(error) => { @@ -462,8 +536,9 @@ where send_malformed(&incoming_tx, error.to_string()); break; } + pending_ack = Some(inbound_ciphertexts.cumulative_ack()); } - Ok(RelayFrameBodyKind::Reset) => { + RelayFrameBodyKind::Reset => { let _ = incoming_tx.try_send( JsonRpcConnectionEvent::Disconnected { reason: Some( @@ -473,12 +548,10 @@ where ); break; } - Ok( - RelayFrameBodyKind::Ack - | RelayFrameBodyKind::Resume - | RelayFrameBodyKind::Heartbeat, - ) => {} - Ok(RelayFrameBodyKind::Handshake) | Err(_) => { + RelayFrameBodyKind::Ack + | RelayFrameBodyKind::Resume + | RelayFrameBodyKind::Heartbeat => {} + RelayFrameBodyKind::Handshake => { send_malformed( &incoming_tx, "Noise relay received invalid post-handshake frame".to_string(), 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 348351e61c..68946abc93 100644 --- a/codex-rs/exec-server/src/noise_relay/harness_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/harness_tests.rs @@ -10,7 +10,9 @@ use anyhow::Context; use anyhow::Result; use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCRequest; +use codex_exec_server_protocol::JSONRPCResponse; use codex_exec_server_protocol::RequestId; +use futures::FutureExt; use futures::Sink; use futures::SinkExt; use futures::StreamExt; @@ -26,14 +28,17 @@ use tokio_tungstenite::tungstenite::Message; use super::*; use crate::connection::JsonRpcConnectionEvent; +use crate::noise_channel::NoiseTransport; use crate::noise_channel::PendingResponderHandshake; +use crate::noise_relay::reliable_stream::RESEND_AFTER; const ENVIRONMENT_ID: &str = "environment-1"; const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; #[tokio::test(start_paused = true)] async fn fragmented_writes_yield_to_keepalive_and_queued_pong() -> Result<()> { - let (connection, mut control, mut outbound_rx) = connected_controlled_harness().await?; + let (connection, mut control, mut outbound_rx, _stream_id, _transport) = + connected_controlled_harness().await?; connection .outgoing_tx @@ -51,7 +56,7 @@ async fn fragmented_writes_yield_to_keepalive_and_queued_pong() -> Result<()> { tokio::time::advance(WEBSOCKET_KEEPALIVE_INTERVAL + Duration::from_millis(10)).await; control.grant_writes(/*count*/ 1); let first_data = read_outbound_data(&mut outbound_rx).await?; - assert_eq!(first_data.seq, 0); + assert_eq!(first_data.seq, 1); control.wait_for_blocked_write(/*expected*/ 2).await?; control.grant_writes(/*count*/ 1); @@ -67,7 +72,7 @@ async fn fragmented_writes_yield_to_keepalive_and_queued_pong() -> Result<()> { tokio::time::advance(WEBSOCKET_KEEPALIVE_INTERVAL + Duration::from_millis(10)).await; control.grant_writes(/*count*/ 1); let second_data = read_outbound_data(&mut outbound_rx).await?; - assert_eq!(second_data.seq, 1); + assert_eq!(second_data.seq, 2); control.wait_for_blocked_write(/*expected*/ 4).await?; control.grant_writes(/*count*/ 1); @@ -82,9 +87,100 @@ async fn fragmented_writes_yield_to_keepalive_and_queued_pong() -> Result<()> { Ok(()) } +#[tokio::test(start_paused = true)] +async fn dropped_data_is_retried_with_identical_ciphertext_until_cumulative_ack() -> Result<()> { + let (connection, control, mut outbound_rx, stream_id, _transport) = + connected_controlled_harness_with_write_permits(/*write_permits*/ 64).await?; + + connection + .outgoing_tx + .send(JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "retry".to_string(), + params: None, + trace: None, + })) + .await?; + + let first_data = read_outbound_data_with_pongs(&mut outbound_rx, &control).await?; + assert_eq!(first_data.seq, 1); + + tokio::time::advance(RESEND_AFTER + RELIABLE_RETRY_SCAN_INTERVAL).await; + let retry_data = read_outbound_data_with_pongs(&mut outbound_rx, &control).await?; + assert_eq!(retry_data, first_data); + drain_outbound_control(&mut outbound_rx, &control).await?; + + control.send_inbound(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame::ack(stream_id, /*ack*/ 1)).into(), + ))?; + tokio::task::yield_now().await; + drain_outbound_control(&mut outbound_rx, &control).await?; + tokio::time::advance(RESEND_AFTER + RELIABLE_RETRY_SCAN_INTERVAL).await; + drain_outbound_control(&mut outbound_rx, &control).await?; + + for task in &connection.task_handles { + task.abort(); + } + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn queued_pong_is_drained_before_deferred_ack_write() -> Result<()> { + let (connection, mut control, mut outbound_rx, stream_id, mut executor_transport) = + connected_controlled_harness().await?; + + control.wait_for_blocked_write(/*expected*/ 1).await?; + control.grant_writes(/*count*/ 1); + let Message::Ping(ping_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending keepalive")? + else { + anyhow::bail!("expected keepalive ping"); + }; + let reads_before_deadline = control.inbound_reads(); + + let framed = frame_jsonrpc_message(&JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result: serde_json::Value::Null, + }))?; + let ciphertext = executor_transport.encrypt(&framed[..1])?; + let data = + RelayMessageFrame::reliable_data(stream_id, /*ack*/ 0, /*seq*/ 1, ciphertext); + control.send_inbound(Message::Binary(encode_relay_message_frame(&data).into()))?; + control.send_inbound(Message::Pong(ping_payload))?; + + // Hold the current-thread runtime until the Pong deadline passes so Data + // and Pong are both already queued when the grace drain starts. + std::thread::sleep(WEBSOCKET_PONG_TIMEOUT + Duration::from_millis(10)); + for _ in 0..8 { + tokio::task::yield_now().await; + if control.inbound_reads() - reads_before_deadline == 2 { + break; + } + } + + assert_eq!(control.inbound_reads() - reads_before_deadline, 2); + assert!(!*connection.disconnected_rx.borrow()); + control.grant_writes(/*count*/ 1); + let Message::Binary(ack_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending deferred ack")? + else { + anyhow::bail!("expected deferred ack frame"); + }; + let ack = decode_relay_message_frame(ack_payload.as_ref())?; + assert_eq!(ack.validate()?, RelayFrameBodyKind::Ack); + assert_eq!(ack.ack, 1); + for task in &connection.task_handles { + task.abort(); + } + Ok(()) +} + #[tokio::test(flavor = "current_thread")] async fn post_deadline_drain_stops_before_frame_33() -> Result<()> { - let (mut connection, mut control, mut outbound_rx) = connected_controlled_harness().await?; + let (mut connection, mut control, mut outbound_rx, _stream_id, _transport) = + connected_controlled_harness().await?; control.wait_for_blocked_write(/*expected*/ 1).await?; control.grant_writes(/*count*/ 1); @@ -141,15 +237,6 @@ async fn pong_keeps_harness_alive_until_peer_stops_responding() -> Result<()> { }, ); - let resume_message = timeout(Duration::from_secs(1), executor_websocket.next()) - .await? - .context("harness closed before sending resume")??; - let Message::Binary(resume_payload) = resume_message else { - anyhow::bail!("expected resume frame, got {resume_message:?}"); - }; - let resume = decode_relay_message_frame(resume_payload.as_ref())?; - assert_eq!(resume.validate()?, RelayFrameBodyKind::Resume); - let handshake_message = timeout(Duration::from_secs(1), executor_websocket.next()) .await? .context("harness closed before sending handshake")??; @@ -157,7 +244,7 @@ async fn pong_keeps_harness_alive_until_peer_stops_responding() -> Result<()> { anyhow::bail!("expected handshake frame, got {handshake_message:?}"); }; let handshake = decode_relay_message_frame(handshake_payload.as_ref())?; - assert_eq!(handshake.stream_id, resume.stream_id); + assert_eq!(handshake.validate()?, RelayFrameBodyKind::Handshake); let stream_id = handshake.stream_id.clone(); let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id.as_str()); @@ -251,12 +338,68 @@ async fn read_outbound_data( frame.into_data().map_err(anyhow::Error::from) } +async fn read_outbound_data_with_pongs( + outbound_rx: &mut futures_mpsc::UnboundedReceiver, + control: &ControlledWebSocketHandle, +) -> Result { + loop { + let message = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .context("harness closed before sending data")?; + match message { + Message::Binary(payload) => { + let frame = decode_relay_message_frame(payload.as_ref())?; + assert_eq!(frame.validate()?, RelayFrameBodyKind::Data); + return frame.into_data().map_err(anyhow::Error::from); + } + Message::Ping(payload) => control.send_inbound(Message::Pong(payload))?, + Message::Pong(_) | Message::Frame(_) => {} + message => anyhow::bail!("expected relay data frame, got {message:?}"), + } + } +} + +async fn drain_outbound_control( + outbound_rx: &mut futures_mpsc::UnboundedReceiver, + control: &ControlledWebSocketHandle, +) -> Result<()> { + for _ in 0..3 { + tokio::task::yield_now().await; + while let Some(message) = outbound_rx.next().now_or_never().flatten() { + match message { + Message::Ping(payload) => control.send_inbound(Message::Pong(payload))?, + Message::Binary(payload) => { + let frame = decode_relay_message_frame(payload.as_ref())?; + assert_ne!(frame.validate()?, RelayFrameBodyKind::Data); + } + Message::Pong(_) | Message::Frame(_) => {} + message => anyhow::bail!("unexpected outbound message after ack: {message:?}"), + } + } + } + Ok(()) +} + async fn connected_controlled_harness() -> Result<( JsonRpcConnection, ControlledWebSocketHandle, futures_mpsc::UnboundedReceiver, + String, + NoiseTransport, )> { - let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(/*write_permits*/ 2); + connected_controlled_harness_with_write_permits(/*write_permits*/ 1).await +} + +async fn connected_controlled_harness_with_write_permits( + write_permits: usize, +) -> Result<( + JsonRpcConnection, + ControlledWebSocketHandle, + futures_mpsc::UnboundedReceiver, + String, + NoiseTransport, +)> { + let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(write_permits); let executor_identity = NoiseChannelIdentity::generate()?; let connection = noise_harness_connection_from_websocket( websocket, @@ -270,13 +413,6 @@ async fn connected_controlled_harness() -> Result<( }, ); - let Message::Binary(resume_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) - .await? - .context("harness closed before sending resume")? - else { - anyhow::bail!("expected resume frame"); - }; - let resume = decode_relay_message_frame(resume_payload.as_ref())?; let Message::Binary(handshake_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) .await? .context("harness closed before sending handshake")? @@ -284,8 +420,8 @@ async fn connected_controlled_harness() -> Result<( anyhow::bail!("expected handshake frame"); }; let handshake = decode_relay_message_frame(handshake_payload.as_ref())?; + assert_eq!(handshake.validate()?, RelayFrameBodyKind::Handshake); let stream_id = handshake.stream_id.clone(); - assert_eq!(stream_id, resume.stream_id); let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id.as_str()); let pending = PendingResponderHandshake::read_request( @@ -293,11 +429,12 @@ async fn connected_controlled_harness() -> Result<( &prologue, &handshake.into_handshake_payload()?, )?; - let (_transport, response) = pending.complete()?; + let (transport, response) = pending.complete()?; control.send_inbound(Message::Binary( - encode_relay_message_frame(&RelayMessageFrame::handshake(stream_id, response)).into(), + encode_relay_message_frame(&RelayMessageFrame::handshake(stream_id.clone(), response)) + .into(), ))?; - Ok((connection, control, outbound_rx)) + Ok((connection, control, outbound_rx, stream_id, transport)) } struct ControlledWebSocket { diff --git a/codex-rs/exec-server/src/noise_relay/mod.rs b/codex-rs/exec-server/src/noise_relay/mod.rs index 56fabd5cb5..d45a7b2a5b 100644 --- a/codex-rs/exec-server/src/noise_relay/mod.rs +++ b/codex-rs/exec-server/src/noise_relay/mod.rs @@ -2,11 +2,10 @@ pub(crate) mod executor_stream; mod harness; mod message_framing; mod ordered_ciphertext; +mod reliable_stream; use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; -use crate::ExecServerError; - pub(crate) use harness::NoiseHarnessConnectionArgs; pub(crate) use harness::noise_harness_connection_from_websocket; @@ -22,13 +21,3 @@ pub(crate) fn noise_relay_websocket_config() -> WebSocketConfig { .max_frame_size(Some(MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE)) .max_message_size(Some(MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE)) } - -fn take_next_sequence(next_seq: &mut u32) -> Result { - // Never wrap: relay sequence is the explicit ordering key for an implicit - // Noise nonce. Reusing zero after u32::MAX would be ambiguous and unsafe. - let seq = *next_seq; - *next_seq = next_seq.checked_add(1).ok_or_else(|| { - ExecServerError::Protocol("Noise relay sequence number exhausted".to_string()) - })?; - Ok(seq) -} diff --git a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs index 92bbd291d7..5bd7301245 100644 --- a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs +++ b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs @@ -2,19 +2,36 @@ use std::collections::BTreeMap; use crate::ExecServerError; -const MAX_REORDER_DISTANCE: u32 = 64; -const MAX_PENDING_BYTES: usize = 1024 * 1024; +// A receive window of 32 includes the next expected sequence itself, so only +// 31 later records can wait behind one gap. +const MAX_REORDER_DISTANCE: u32 = 31; +const MAX_PENDING_BYTES: usize = 2 * 1024 * 1024; /// Reorders relay records before they reach Noise's implicit receive nonce. /// The window is bounded, and each sequence number is released at most once. -#[derive(Default)] pub(crate) struct OrderedCiphertextFrames { next_seq: u32, pending: BTreeMap>, pending_bytes: usize, } +impl Default for OrderedCiphertextFrames { + fn default() -> Self { + Self { + // Reliable Noise streams reserve zero as the initial cumulative ack. + next_seq: 1, + pending: BTreeMap::new(), + pending_bytes: 0, + } + } +} + impl OrderedCiphertextFrames { + /// Highest contiguous sequence released for decryption. + pub(crate) fn cumulative_ack(&self) -> u32 { + self.next_seq - 1 + } + /// Accept one relay record and return the newly contiguous ciphertext run. /// /// Returns nothing for duplicates or while a gap remains. Closing a gap also @@ -24,6 +41,11 @@ impl OrderedCiphertextFrames { seq: u32, payload: Vec, ) -> Result>, ExecServerError> { + if seq == 0 { + return Err(ExecServerError::Protocol( + "Noise reliable data sequence zero is reserved".to_string(), + )); + } // Keep the first ciphertext for a sequence. Later copies are duplicates. if seq < self.next_seq || self.pending.contains_key(&seq) { return Ok(Vec::new()); diff --git a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs index 6aa86fcdf0..db45c11080 100644 --- a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs +++ b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs @@ -8,13 +8,15 @@ fn releases_ciphertexts_only_in_nonce_order() { let mut frames = OrderedCiphertextFrames::default(); assert_eq!( - frames.push(/*seq*/ 1, b"second".to_vec()).unwrap(), + frames.push(/*seq*/ 2, b"second".to_vec()).unwrap(), Vec::>::new() ); + assert_eq!(frames.cumulative_ack(), 0); assert_eq!( - frames.push(/*seq*/ 0, b"first".to_vec()).unwrap(), + frames.push(/*seq*/ 1, b"first".to_vec()).unwrap(), vec![b"first".to_vec(), b"second".to_vec()] ); + assert_eq!(frames.cumulative_ack(), 2); } #[test] @@ -22,19 +24,19 @@ fn ignores_duplicate_ciphertexts_without_replacing_buffered_record() { let mut frames = OrderedCiphertextFrames::default(); assert_eq!( - frames.push(/*seq*/ 1, b"first copy".to_vec()).unwrap(), + frames.push(/*seq*/ 2, b"first copy".to_vec()).unwrap(), Vec::>::new() ); assert_eq!( - frames.push(/*seq*/ 1, b"replacement".to_vec()).unwrap(), + frames.push(/*seq*/ 2, b"replacement".to_vec()).unwrap(), Vec::>::new() ); assert_eq!( - frames.push(/*seq*/ 0, b"zero".to_vec()).unwrap(), - vec![b"zero".to_vec(), b"first copy".to_vec()] + frames.push(/*seq*/ 1, b"one".to_vec()).unwrap(), + vec![b"one".to_vec(), b"first copy".to_vec()] ); assert_eq!( - frames.push(/*seq*/ 0, b"duplicate".to_vec()).unwrap(), + frames.push(/*seq*/ 1, b"duplicate".to_vec()).unwrap(), Vec::>::new() ); } @@ -43,10 +45,21 @@ fn ignores_duplicate_ciphertexts_without_replacing_buffered_record() { fn rejects_unbounded_reordering() { let mut frames = OrderedCiphertextFrames::default(); - assert!(frames.push(/*seq*/ 65, Vec::new()).is_err()); + assert!(frames.push(/*seq*/ 0, Vec::new()).is_err()); + assert!(frames.push(/*seq*/ 33, Vec::new()).is_err()); assert!( frames - .push(/*seq*/ 1, vec![0; MAX_PENDING_BYTES + 1]) + .push(/*seq*/ 2, vec![0; MAX_PENDING_BYTES + 1]) .is_err() ); } + +#[test] +fn buffers_the_full_receive_window_behind_one_gap() { + let mut frames = OrderedCiphertextFrames::default(); + + for seq in 2..=32 { + assert!(frames.push(seq, vec![0; 64 * 1024]).is_ok()); + } + assert_eq!(frames.cumulative_ack(), 0); +} diff --git a/codex-rs/exec-server/src/noise_relay/reliable_stream.rs b/codex-rs/exec-server/src/noise_relay/reliable_stream.rs new file mode 100644 index 0000000000..624f8024ae --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/reliable_stream.rs @@ -0,0 +1,189 @@ +use std::collections::BTreeMap; +use std::time::Duration; + +use tokio::time::Instant; + +use crate::ExecServerError; + +/// Maximum number of encrypted records retained until the peer cumulatively +/// acknowledges them. +pub(crate) const MAX_UNACKED_SEGMENTS: usize = 32; +/// Maximum encrypted bytes retained until the peer cumulatively acknowledges +/// them. +pub(crate) const MAX_UNACKED_BYTES: usize = 2 * 1024 * 1024; +/// How long an encrypted record may remain unacknowledged before it is retried. +pub(crate) const RESEND_AFTER: Duration = Duration::from_millis(500); + +/// One encrypted record ready for an initial send or retry. +/// +/// The payload is already Noise-encrypted. Retries clone these exact bytes +/// rather than asking Noise to encrypt the logical record again. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct OutboundCiphertext { + pub(crate) seq: u32, + pub(crate) payload: Vec, +} + +#[derive(Debug)] +struct UnackedCiphertext { + payload: Vec, + last_sent_at: Instant, +} + +/// Sender-side reliability state for one logical Noise relay stream. +/// +/// The receive frontier stays in `OrderedCiphertextFrames`, which already owns +/// the bounded reorder buffer needed before Noise decryption. This type only +/// allocates send sequences, applies peer cumulative acks, and retains exact +/// ciphertext for retries. +#[derive(Debug)] +pub(crate) struct ReliableSender { + next_seq: u32, + highest_sent_seq: u32, + resend_cursor: u32, + unacked: BTreeMap, + unacked_bytes: usize, +} + +impl Default for ReliableSender { + fn default() -> Self { + Self { + // Sequence zero is reserved so ack = 0 unambiguously means that + // nothing has been received contiguously yet. + next_seq: 1, + highest_sent_seq: 0, + resend_cursor: 1, + unacked: BTreeMap::new(), + unacked_bytes: 0, + } + } +} + +impl ReliableSender { + /// Whether one newly encrypted payload can be retained in the send window. + /// + /// Callers must check this before consuming another Noise send nonce. + pub(crate) fn can_admit_ciphertext(&self, ciphertext_len: usize) -> bool { + ciphertext_len > 0 + && ciphertext_len <= MAX_UNACKED_BYTES + && self.unacked.len() < MAX_UNACKED_SEGMENTS + && self + .unacked_bytes + .checked_add(ciphertext_len) + .is_some_and(|bytes| bytes <= MAX_UNACKED_BYTES) + } + + /// Allocate the next sequence number and retain an already-encrypted record. + pub(crate) fn admit_ciphertext( + &mut self, + payload: Vec, + now: Instant, + ) -> Result { + if payload.is_empty() { + return Err(ExecServerError::Protocol( + "Noise reliable ciphertext payload is empty".to_string(), + )); + } + if payload.len() > MAX_UNACKED_BYTES { + return Err(ExecServerError::Protocol(format!( + "Noise reliable ciphertext exceeds unacked byte limit: {} > {MAX_UNACKED_BYTES}", + payload.len() + ))); + } + if self.unacked.len() >= MAX_UNACKED_SEGMENTS { + return Err(ExecServerError::Protocol( + "Noise reliable segment send window is full".to_string(), + )); + } + let unacked_bytes = self + .unacked_bytes + .checked_add(payload.len()) + .filter(|bytes| *bytes <= MAX_UNACKED_BYTES) + .ok_or_else(|| { + ExecServerError::Protocol("Noise reliable byte send window is full".to_string()) + })?; + let seq = self.next_seq; + self.next_seq = self.next_seq.checked_add(1).ok_or_else(|| { + ExecServerError::Protocol("Noise reliable sequence number exhausted".to_string()) + })?; + self.highest_sent_seq = seq; + self.unacked_bytes = unacked_bytes; + self.unacked.insert( + seq, + UnackedCiphertext { + payload: payload.clone(), + last_sent_at: now, + }, + ); + Ok(OutboundCiphertext { seq, payload }) + } + + /// Apply cumulative peer acknowledgement metadata. + /// + /// ack = 0 is the reserved empty state. PR1 rejects nonzero selective bits + /// so they cannot silently alter delivery semantics before PR2. + pub(crate) fn process_peer_ack( + &mut self, + ack: u32, + ack_bits: u32, + ) -> Result<(), ExecServerError> { + if ack_bits != 0 { + return Err(ExecServerError::Protocol(format!( + "Noise reliable selective ack bits are unsupported in PR1: {ack_bits}" + ))); + } + if ack == 0 { + return Ok(()); + } + if ack > self.highest_sent_seq { + return Err(ExecServerError::Protocol(format!( + "Noise reliable peer ack {ack} exceeds highest sent sequence {}", + self.highest_sent_seq + ))); + } + + let acknowledged = self + .unacked + .range(..=ack) + .map(|(seq, _pending)| *seq) + .collect::>(); + for seq in acknowledged { + if let Some(pending) = self.unacked.remove(&seq) { + self.unacked_bytes -= pending.payload.len(); + } + } + Ok(()) + } + + /// Return the next cached record whose retry deadline has elapsed. + /// + /// Returning one record at a time preserves a scheduling point between + /// retries so websocket control traffic cannot sit behind a full-window + /// resend burst. The cursor gives every retained record a chance before + /// wrapping to an older record that became due again. + pub(crate) fn next_retry_due(&mut self, now: Instant) -> Option { + let due_seq = self + .unacked + .range(self.resend_cursor..) + .chain(self.unacked.range(..self.resend_cursor)) + .find_map(|(seq, pending)| { + now.checked_duration_since(pending.last_sent_at) + .is_some_and(|elapsed| elapsed >= RESEND_AFTER) + .then_some(*seq) + })?; + let payload = { + let pending = self.unacked.get_mut(&due_seq)?; + pending.last_sent_at = now; + pending.payload.clone() + }; + self.resend_cursor = due_seq.checked_add(1).unwrap_or(1); + Some(OutboundCiphertext { + seq: due_seq, + payload, + }) + } +} + +#[cfg(test)] +#[path = "reliable_stream_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/noise_relay/reliable_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/reliable_stream_tests.rs new file mode 100644 index 0000000000..5d0fac9be7 --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/reliable_stream_tests.rs @@ -0,0 +1,120 @@ +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::time::Instant; + +use super::MAX_UNACKED_BYTES; +use super::MAX_UNACKED_SEGMENTS; +use super::OutboundCiphertext; +use super::RESEND_AFTER; +use super::ReliableSender; + +#[test] +fn starts_at_sequence_one() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + + assert_eq!( + sender + .admit_ciphertext(b"ciphertext".to_vec(), now) + .unwrap(), + OutboundCiphertext { + seq: 1, + payload: b"ciphertext".to_vec(), + } + ); +} + +#[test] +fn cumulative_ack_zero_clears_nothing_and_positive_ack_releases_prefix() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + sender.admit_ciphertext(b"first".to_vec(), now).unwrap(); + sender.admit_ciphertext(b"second".to_vec(), now).unwrap(); + + sender.process_peer_ack(/*ack*/ 0, /*ack_bits*/ 0).unwrap(); + assert_eq!(sender.unacked.len(), 2); + sender.process_peer_ack(/*ack*/ 1, /*ack_bits*/ 0).unwrap(); + assert_eq!(sender.unacked.len(), 1); + assert_eq!(sender.unacked_bytes, b"second".len()); + sender.process_peer_ack(/*ack*/ 2, /*ack_bits*/ 0).unwrap(); + assert_eq!(sender.unacked.len(), 0); + assert_eq!(sender.unacked_bytes, 0); +} + +#[test] +fn rejects_selective_bits_and_ack_beyond_highest_sent_sequence() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + sender + .admit_ciphertext(b"ciphertext".to_vec(), now) + .unwrap(); + + assert!(sender.process_peer_ack(/*ack*/ 0, /*ack_bits*/ 1).is_err()); + assert!(sender.process_peer_ack(/*ack*/ 2, /*ack_bits*/ 0).is_err()); + assert_eq!(sender.unacked.len(), 1); +} + +#[test] +fn retries_exact_cached_ciphertext_after_deadline() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + let first = sender + .admit_ciphertext(b"encrypted-once".to_vec(), now) + .unwrap(); + + assert_eq!( + sender.next_retry_due(now + RESEND_AFTER - Duration::from_millis(1)), + None + ); + assert_eq!( + sender.next_retry_due(now + RESEND_AFTER), + Some(first.clone()) + ); + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), None); + assert_eq!(first.payload, b"encrypted-once".to_vec()); +} + +#[test] +fn returns_one_due_retry_per_scan() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + let first = sender.admit_ciphertext(b"first".to_vec(), now).unwrap(); + let second = sender.admit_ciphertext(b"second".to_vec(), now).unwrap(); + + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), Some(first)); + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), Some(second)); + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), None); +} + +#[test] +fn retry_cursor_does_not_starve_later_due_records() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + let first = sender.admit_ciphertext(b"first".to_vec(), now).unwrap(); + let second = sender.admit_ciphertext(b"second".to_vec(), now).unwrap(); + + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), Some(first)); + assert_eq!( + sender.next_retry_due(now + RESEND_AFTER + RESEND_AFTER), + Some(second) + ); +} + +#[test] +fn enforces_segment_and_byte_send_windows() { + let now = Instant::now(); + let mut segment_window = ReliableSender::default(); + for _ in 0..MAX_UNACKED_SEGMENTS { + segment_window.admit_ciphertext(vec![0x5a], now).unwrap(); + } + assert!(!segment_window.can_admit_ciphertext(/*ciphertext_len*/ 1)); + assert!(segment_window.admit_ciphertext(vec![0x5a], now).is_err()); + + let mut byte_window = ReliableSender::default(); + byte_window + .admit_ciphertext(vec![0x5a; MAX_UNACKED_BYTES], now) + .unwrap(); + assert!(!byte_window.can_admit_ciphertext(/*ciphertext_len*/ 1)); + assert!(byte_window.admit_ciphertext(vec![0x5a], now).is_err()); +} diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index b8d477f5db..03d91993e3 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,4 +1,6 @@ use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; use std::time::Duration; use codex_exec_server_protocol::JSONRPCMessage; @@ -34,6 +36,7 @@ use crate::noise_relay::NOISE_RELAY_RESET_REASON; use crate::noise_relay::executor_stream::ClosedNoiseVirtualStream; use crate::noise_relay::executor_stream::NoiseVirtualStream; use crate::noise_relay::executor_stream::spawn_noise_virtual_stream; +use crate::relay_proto::RelayAck; use crate::relay_proto::RelayData; use crate::relay_proto::RelayHandshake; use crate::relay_proto::RelayMessageFrame; @@ -50,6 +53,7 @@ const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128; const MAX_FAILED_NOISE_HANDSHAKES: usize = 8; const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096; const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32; +const MAX_RECENT_RETIRED_NOISE_RELAY_STREAMS: usize = 1024; const HARNESS_KEY_VALIDATION_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Debug, Clone, Copy, Eq, PartialEq)] @@ -99,6 +103,28 @@ impl RelayMessageFrame { } } + /// Build one reliable Noise data frame with the receiver's cumulative ack. + /// + /// The legacy plaintext relay still uses [`Self::data`] and keeps its old + /// sequence behavior. Noise reliability owns its own sequence and ack + /// invariants, so keep that policy at the Noise call sites. + pub(crate) fn reliable_data(stream_id: String, ack: u32, seq: u32, payload: Vec) -> Self { + let mut frame = Self::data(stream_id, seq, payload); + frame.ack = ack; + frame + } + + /// Build an unsequenced cumulative acknowledgement for a reliable stream. + pub(crate) fn ack(stream_id: String, ack: u32) -> Self { + Self { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + ack, + ack_bits: 0, + body: Some(relay_message_frame::Body::AckFrame(RelayAck {})), + } + } + pub(crate) fn resume(stream_id: String) -> Self { Self { version: RELAY_MESSAGE_FRAME_VERSION, @@ -553,6 +579,11 @@ where } }); let mut streams: HashMap = HashMap::new(); + // Stream IDs are single-use by contract. Keep a bounded recent tombstone + // set so an accidental near-term reuse cannot receive ciphertext that an + // old writer still has queued in the shared physical channel. + let mut retired_stream_ids: HashSet = HashSet::new(); + let mut retired_stream_order: VecDeque = VecDeque::new(); let mut pending_handshakes: HashMap = HashMap::new(); let mut validation_tasks: JoinSet = JoinSet::new(); let mut failed_handshakes = 0usize; @@ -573,13 +604,18 @@ where break; } Some(closed_stream) = closed_stream_rx.recv() => { - // A stream ID may have been reused before this writer exits. - // Remove only the instance that sent the notification. + // A delayed close can race another teardown path. Remove only + // the instance that sent the notification. let is_current = streams .get(&closed_stream.stream_id) .is_some_and(|stream| stream.instance_id == closed_stream.instance_id); if is_current { streams.remove(&closed_stream.stream_id); + retire_noise_stream_id( + &mut retired_stream_ids, + &mut retired_stream_order, + closed_stream.stream_id, + ); } continue; } @@ -725,10 +761,10 @@ where let stream_id = frame.stream_id.clone(); match kind { RelayFrameBodyKind::Handshake => { - // 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. - if streams.contains_key(&stream_id) { + // Reject duplicate, retired, or busy streams before paying for + // a hybrid handshake. Malformed attempts that reach cryptography + // are covered by the connection-wide failure budget below. + if streams.contains_key(&stream_id) || retired_stream_ids.contains(&stream_id) { send_reset(&physical_outgoing_tx, stream_id); continue; } @@ -834,6 +870,8 @@ where }); } RelayFrameBodyKind::Data => { + let peer_ack = frame.ack; + let peer_ack_bits = frame.ack_bits; // Removing pending state also makes any in-flight validation stale. let Some(stream) = streams.get_mut(&stream_id) else { let canceled_pending_handshake = @@ -847,11 +885,27 @@ where } continue; }; + if let Err(error) = stream.process_peer_ack(peer_ack, peer_ack_bits) { + warn!("failed to process Noise relay peer ack: {error}"); + streams.remove(&stream_id); + retire_noise_stream_id( + &mut retired_stream_ids, + &mut retired_stream_order, + stream_id.clone(), + ); + send_reset(&physical_outgoing_tx, stream_id); + continue; + } let data = match frame.into_data() { Ok(data) => data, Err(error) => { warn!("dropping malformed Noise relay data frame: {error}"); streams.remove(&stream_id); + retire_noise_stream_id( + &mut retired_stream_ids, + &mut retired_stream_order, + stream_id.clone(), + ); send_reset(&physical_outgoing_tx, stream_id); continue; } @@ -859,19 +913,42 @@ where if let Err(error) = stream.receive_data(data) { warn!("failed to process Noise relay payload: {error}"); streams.remove(&stream_id); + retire_noise_stream_id( + &mut retired_stream_ids, + &mut retired_stream_order, + stream_id.clone(), + ); send_reset(&physical_outgoing_tx, stream_id); } } RelayFrameBodyKind::Reset => { pending_handshakes.remove(&stream_id); if let Some(stream) = streams.remove(&stream_id) { + retire_noise_stream_id( + &mut retired_stream_ids, + &mut retired_stream_order, + stream_id, + ); // The reset reason is unauthenticated, so do not log it. stream.disconnect(/*reason*/ None); } } RelayFrameBodyKind::Ack | RelayFrameBodyKind::Resume - | RelayFrameBodyKind::Heartbeat => {} + | RelayFrameBodyKind::Heartbeat => { + if let Some(stream) = streams.get(&stream_id) + && let Err(error) = stream.process_peer_ack(frame.ack, frame.ack_bits) + { + warn!("failed to process Noise relay peer ack: {error}"); + streams.remove(&stream_id); + retire_noise_stream_id( + &mut retired_stream_ids, + &mut retired_stream_order, + stream_id.clone(), + ); + send_reset(&physical_outgoing_tx, stream_id); + } + } } } @@ -908,6 +985,22 @@ struct HarnessKeyValidationResult { result: Result<(), ExecServerError>, } +fn retire_noise_stream_id( + retired_stream_ids: &mut HashSet, + retired_stream_order: &mut VecDeque, + stream_id: String, +) { + if !retired_stream_ids.insert(stream_id.clone()) { + return; + } + retired_stream_order.push_back(stream_id); + if retired_stream_order.len() > MAX_RECENT_RETIRED_NOISE_RELAY_STREAMS + && let Some(oldest_stream_id) = retired_stream_order.pop_front() + { + retired_stream_ids.remove(&oldest_stream_id); + } +} + /// Queue a best-effort reset without blocking the shared websocket loop. /// Reset reasons are relay control data and are not treated as trusted text. fn send_reset(physical_outgoing_tx: &mpsc::Sender>, stream_id: String) { diff --git a/codex-rs/exec-server/src/relay_proto.rs b/codex-rs/exec-server/src/relay_proto.rs index da7cb5296d..b31d11163a 100644 --- a/codex-rs/exec-server/src/relay_proto.rs +++ b/codex-rs/exec-server/src/relay_proto.rs @@ -1,6 +1,7 @@ #[path = "proto/codex.exec_server.relay.v1.rs"] mod generated; +pub(crate) use generated::RelayAck; pub(crate) use generated::RelayData; pub(crate) use generated::RelayHandshake; pub(crate) use generated::RelayMessageFrame; diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index 918cba7393..6d3b4de485 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -53,6 +53,8 @@ const HARNESS_KEY_AUTHORIZATION: &str = "harness-key-authorization"; const REGISTRY_TOKEN: &str = "registry-token"; const TEST_TIMEOUT: Duration = Duration::from_secs(10); +type DroppedRelayData = Arc)>>>; + #[derive(Debug)] struct StaticRegistryAuthProvider; @@ -70,7 +72,7 @@ fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { +async fn remote_environment_retries_dropped_encrypted_frames() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let rendezvous_url = format!("ws://{}", listener.local_addr()?); let registry = MockServer::start().await; @@ -131,10 +133,14 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { tokio::spawn(async move { ExecServerClient::connect_noise_rendezvous(client_args).await }); let harness_websocket = accept_websocket(&listener, "harness").await?; let captured_frames = Arc::new(Mutex::new(Vec::new())); + let dropped_harness_data = Arc::new(Mutex::new(None)); + let dropped_environment_data = Arc::new(Mutex::new(None)); let relay_task = tokio::spawn(proxy_relay_frames( environment_websocket, harness_websocket, Arc::clone(&captured_frames), + Arc::clone(&dropped_harness_data), + Arc::clone(&dropped_environment_data), )); let client = timeout(TEST_TIMEOUT, client_task) .await @@ -178,6 +184,12 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { ); assert_relay_data_is_encrypted(&captured_frames)?; + assert_dropped_ciphertext_was_retried(&captured_frames, "harness", &dropped_harness_data)?; + assert_dropped_ciphertext_was_retried( + &captured_frames, + "environment", + &dropped_environment_data, + )?; drop(client); relay_task.abort(); @@ -218,6 +230,8 @@ async fn proxy_relay_frames( mut environment: WebSocketStream, mut harness: WebSocketStream, captured_frames: Arc>>>, + dropped_harness_data: DroppedRelayData, + dropped_environment_data: DroppedRelayData, ) -> Result<()> { loop { tokio::select! { @@ -227,6 +241,9 @@ async fn proxy_relay_frames( }; let message = message?; capture_binary_frame(&captured_frames, &message); + if drop_first_data(&dropped_environment_data, &message)? { + continue; + } harness.send(message).await?; } message = harness.next() => { @@ -235,6 +252,9 @@ async fn proxy_relay_frames( }; let message = message?; capture_binary_frame(&captured_frames, &message); + if drop_first_data(&dropped_harness_data, &message)? { + continue; + } environment.send(message).await?; } } @@ -242,6 +262,24 @@ async fn proxy_relay_frames( Ok(()) } +fn drop_first_data(dropped_data: &DroppedRelayData, message: &Message) -> Result { + let Message::Binary(bytes) = message else { + return Ok(false); + }; + let frame = RelayMessageFrame::decode(bytes.as_ref())?; + let Some(relay_message_frame::Body::Data(data)) = frame.body else { + return Ok(false); + }; + let mut dropped_data = dropped_data + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if dropped_data.is_some() { + return Ok(false); + } + *dropped_data = Some((data.seq, data.payload)); + Ok(true) +} + fn capture_binary_frame(captured_frames: &Mutex>>, message: &Message) { if let Message::Binary(bytes) = message { captured_frames @@ -273,3 +311,32 @@ fn assert_relay_data_is_encrypted(captured_frames: &Mutex>>) -> Resu ); Ok(()) } + +fn assert_dropped_ciphertext_was_retried( + captured_frames: &Mutex>>, + direction: &str, + dropped_data: &DroppedRelayData, +) -> Result<()> { + let dropped_data = dropped_data + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .with_context(|| format!("fake rendezvous should drop one {direction} data frame"))?; + let captured_frames = captured_frames + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let matching_frames = captured_frames + .iter() + .filter_map(|encoded| RelayMessageFrame::decode(encoded.as_slice()).ok()) + .filter_map(|frame| match frame.body { + Some(relay_message_frame::Body::Data(data)) => Some(data), + _ => None, + }) + .filter(|data| data.seq == dropped_data.0 && data.payload == dropped_data.1) + .count(); + assert!( + matching_frames >= 2, + "expected dropped {direction} ciphertext to be retried without re-encryption" + ); + Ok(()) +}