diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index a250ba58866..e1f77a76feb 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 // bit i acknowledges peer seq = ack + 1 + i 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,17 +99,19 @@ 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` plus a fixed-width selective acknowledgement mask: ```text -ack = highest contiguous received segment seq +ack = 0 when no segment has been received contiguously +ack = highest contiguous received segment seq otherwise bit i in ack_bits acknowledges seq = ack + 1 + i ``` -Send `ack` and `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`. +The 32-segment receive window uses `u32 ack_bits`. +Bit zero is canonically clear because receiving `ack + 1` advances the cumulative `ack`. +Send `ack` and `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`. ## Lifecycle 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 89f33d938a4..07f42e47f16 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,9 @@ use std::sync::Arc; use std::sync::Mutex; +use std::time::Duration; +use tokio::sync::Notify; use tokio::sync::mpsc; use tokio::sync::watch; use tracing::warn; @@ -22,23 +24,35 @@ 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::RelayAckState; 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, } +#[derive(Default)] +struct InboundAckState { + latest: RelayAckState, + pending: Option, +} + /// One authenticated JSON-RPC stream carried by the executor's physical relay. /// /// Inbound delivery is intentionally nonblocking. An overloaded or abandoned @@ -48,6 +62,9 @@ pub(crate) struct NoiseVirtualStream { incoming_tx: mpsc::Sender, disconnected_tx: watch::Sender, transport: Arc>, + reliable_sender: Arc>, + inbound_ack_state: Arc>, + writer_wakeup: Arc, inbound_ciphertexts: OrderedCiphertextFrames, inbound_decoder: JsonRpcMessageDecoder, pub(crate) instance_id: u64, @@ -61,6 +78,19 @@ impl NoiseVirtualStream { .try_send(JsonRpcConnectionEvent::Disconnected { reason }); } + /// Apply ack metadata from one post-handshake peer frame and wake the + /// writer if it opened sequence or byte send capacity. + pub(crate) fn process_peer_ack(&self, ack_state: RelayAckState) -> Result<(), ExecServerError> { + let mut reliable_sender = self + .reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + reliable_sender.process_peer_ack(ack_state)?; + 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 +114,15 @@ impl NoiseVirtualStream { })?; } } + let ack_state = self.inbound_ciphertexts.ack_state(); + let mut inbound_ack_state = self + .inbound_ack_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + inbound_ack_state.latest = ack_state; + inbound_ack_state.pending = Some(ack_state); + drop(inbound_ack_state); + self.writer_wakeup.notify_one(); Ok(()) } } @@ -104,50 +143,156 @@ 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_state = Arc::new(Mutex::new(InboundAckState::default())); + 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_state = Arc::clone(&inbound_ack_state); + 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_inbound_ack_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .pending + .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 = { + let mut reliable_sender = writer_reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + 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; + } + } + }; + let ack_state = writer_inbound_ack_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .latest; + let frame = RelayMessageFrame::reliable_data( + writer_stream_id.clone(), + ack_state, + 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 = { + let mut reliable_sender = writer_reliable_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + reliable_sender.next_retry_due(tokio::time::Instant::now()) + }; + if let Some(outbound) = retry { + let ack_state = writer_inbound_ack_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .latest; + let frame = RelayMessageFrame::reliable_data( + writer_stream_id.clone(), + ack_state, + 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_state = writer_inbound_ack_state .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) + .pending + .take(); + let Some(ack_state) = ack_state else { + continue; + }; + let frame = RelayMessageFrame::ack(writer_stream_id.clone(), ack_state); + 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 +303,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 +330,9 @@ pub(crate) fn spawn_noise_virtual_stream( incoming_tx, disconnected_tx, transport, + reliable_sender, + inbound_ack_state, + 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 232dfa479eb..042fed20b7d 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 @@ -13,25 +13,18 @@ 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::RelayAckState; +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,98 @@ 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 first_ciphertext = harness_transport.encrypt(&framed[..1])?; + let second_ciphertext = harness_transport.encrypt(&framed[1..2])?; + stream.receive_data(RelayData { + seq: 2, + segment_index: 0, + segment_count: 1, + payload: second_ciphertext, + })?; + + assert_eq!( + stream + .inbound_ack_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .latest, + RelayAckState { + ack: 0, + ack_bits: 0b10, + } + ); + 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, 0); + assert_eq!(ack.ack_bits, 0b10); + + stream.receive_data(RelayData { + seq: 1, + segment_index: 0, + segment_count: 1, + payload: first_ciphertext, + })?; + assert_eq!( + stream + .inbound_ack_state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .latest, + RelayAckState { + ack: 2, + ack_bits: 0, + } + ); + let ack = timeout(Duration::from_secs(1), physical_outgoing_rx.recv()) + .await? + .expect("virtual stream writer should send the cumulative ack"); + let ack = decode_relay_message_frame(&ack)?; + assert_eq!(ack.validate()?, RelayFrameBodyKind::Ack); + assert_eq!(ack.ack, 2); + assert_eq!(ack.ack_bits, 0); + 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 9969a3e1300..02fb5f406df 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.ack_state(), + 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.ack_state(), + 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_state) = pending_ack.take() else { + continue; + }; + let frame = RelayMessageFrame::ack(stream_id.clone(), ack_state); + 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_state()) + { + 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.ack_state()); } - 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 348351e61cf..7805e07ee85 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,18 @@ 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; +use crate::relay::RelayAckState; 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 +57,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 +73,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 +88,115 @@ 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, + RelayAckState { + ack: 1, + ack_bits: 0, + }, + )) + .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, + RelayAckState { + ack: 0, + ack_bits: 0, + }, + /*seq*/ 2, + 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, 0); + assert_eq!(ack.ack_bits, 0b10); + 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 +253,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 +260,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 +354,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 +429,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 +436,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 +445,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 56fabd5cb5d..d45a7b2a5be 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 92bbd291d72..b2d18ff2b13 100644 --- a/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs +++ b/codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs @@ -1,20 +1,49 @@ use std::collections::BTreeMap; use crate::ExecServerError; +use crate::relay::RelayAckState; -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 { + /// Current cumulative/selective acknowledgement state. + /// + /// Pending ciphertexts are always later than `next_seq` and bounded to + /// the 32-record receive window, so they map directly to bits 1..31. + pub(crate) fn ack_state(&self) -> RelayAckState { + let ack_bits = self.pending.keys().fold(0, |ack_bits, seq| { + let bit = seq - self.next_seq; + debug_assert!(bit < u32::BITS); + ack_bits | (1u32 << bit) + }); + RelayAckState { + ack: self.next_seq - 1, + ack_bits, + } + } + /// 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 +53,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 6aa86fcdf05..a50c29c3775 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 @@ -2,19 +2,34 @@ use pretty_assertions::assert_eq; use super::MAX_PENDING_BYTES; use super::OrderedCiphertextFrames; +use crate::relay::RelayAckState; #[test] 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.push(/*seq*/ 0, b"first".to_vec()).unwrap(), + frames.ack_state(), + RelayAckState { + ack: 0, + ack_bits: 0b10, + } + ); + assert_eq!( + frames.push(/*seq*/ 1, b"first".to_vec()).unwrap(), vec![b"first".to_vec(), b"second".to_vec()] ); + assert_eq!( + frames.ack_state(), + RelayAckState { + ack: 2, + ack_bits: 0, + } + ); } #[test] @@ -22,31 +37,100 @@ 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.ack_state(), + RelayAckState { + ack: 0, + ack_bits: 0b10, + } + ); + assert_eq!( + frames.push(/*seq*/ 1, b"one".to_vec()).unwrap(), + vec![b"one".to_vec(), b"first copy".to_vec()] + ); + assert_eq!( + frames.push(/*seq*/ 1, b"duplicate".to_vec()).unwrap(), Vec::>::new() ); +} + +#[test] +fn selective_ack_bits_shift_after_a_gap_closes() { + let mut frames = OrderedCiphertextFrames::default(); + assert_eq!( - frames.push(/*seq*/ 0, b"zero".to_vec()).unwrap(), - vec![b"zero".to_vec(), b"first copy".to_vec()] + frames.push(/*seq*/ 2, b"two".to_vec()).unwrap(), + Vec::>::new() ); assert_eq!( - frames.push(/*seq*/ 0, b"duplicate".to_vec()).unwrap(), + frames.push(/*seq*/ 4, b"four".to_vec()).unwrap(), Vec::>::new() ); + assert_eq!( + frames.ack_state(), + RelayAckState { + ack: 0, + ack_bits: 0b1010, + } + ); + + assert_eq!( + frames.push(/*seq*/ 1, b"one".to_vec()).unwrap(), + vec![b"one".to_vec(), b"two".to_vec()] + ); + assert_eq!( + frames.ack_state(), + RelayAckState { + ack: 2, + ack_bits: 0b10, + } + ); + + assert_eq!( + frames.push(/*seq*/ 3, b"three".to_vec()).unwrap(), + vec![b"three".to_vec(), b"four".to_vec()] + ); + assert_eq!( + frames.ack_state(), + RelayAckState { + ack: 4, + ack_bits: 0, + } + ); } #[test] 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.ack_state(), + RelayAckState { + ack: 0, + ack_bits: u32::MAX - 1, + } + ); +} 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 00000000000..7c9ad7e498b --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/reliable_stream.rs @@ -0,0 +1,223 @@ +use std::collections::BTreeMap; +use std::time::Duration; + +use tokio::time::Instant; + +use crate::ExecServerError; +use crate::relay::RelayAckState; + +/// Maximum number of encrypted records retained awaiting peer acknowledgement. +pub(crate) const MAX_UNACKED_SEGMENTS: usize = 32; +/// Maximum encrypted bytes retained awaiting peer acknowledgement. +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 acknowledgement state, and retains +/// exact ciphertext for retries. +#[derive(Debug)] +pub(crate) struct ReliableSender { + next_seq: u32, + highest_sent_seq: u32, + peer_cumulative_ack: 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, + peer_cumulative_ack: 0, + resend_cursor: 1, + unacked: BTreeMap::new(), + unacked_bytes: 0, + } + } +} + +impl ReliableSender { + fn send_window_has_space(&self) -> bool { + self.next_seq + .checked_sub(self.peer_cumulative_ack) + .is_some_and(|span| span <= MAX_UNACKED_SEGMENTS as u32) + } + + /// 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.send_window_has_space() + && 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(), + )); + } + if !self.send_window_has_space() { + return Err(ExecServerError::Protocol( + "Noise reliable cumulative 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 and selective peer acknowledgement metadata. + /// + /// Selective acknowledgement frees cached ciphertext for retry and byte + /// accounting, but only the cumulative frontier slides the sequence window. + pub(crate) fn process_peer_ack( + &mut self, + ack_state: RelayAckState, + ) -> Result<(), ExecServerError> { + let RelayAckState { ack, ack_bits } = ack_state; + if ack > self.highest_sent_seq { + return Err(ExecServerError::Protocol(format!( + "Noise reliable peer ack {ack} exceeds highest sent sequence {}", + self.highest_sent_seq + ))); + } + if ack_bits & 1 != 0 { + return Err(ExecServerError::Protocol( + "Noise reliable selective ack bit zero is inconsistent with cumulative ack" + .to_string(), + )); + } + + let mut selective_acks = Vec::new(); + let mut remaining_ack_bits = ack_bits; + while remaining_ack_bits != 0 { + let bit = remaining_ack_bits.trailing_zeros(); + let seq = ack + .checked_add(1) + .and_then(|seq| seq.checked_add(bit)) + .ok_or_else(|| { + ExecServerError::Protocol( + "Noise reliable selective ack sequence overflow".to_string(), + ) + })?; + if seq > self.highest_sent_seq { + return Err(ExecServerError::Protocol(format!( + "Noise reliable selective ack seq {seq} exceeds highest sent sequence {}", + self.highest_sent_seq + ))); + } + selective_acks.push(seq); + remaining_ack_bits &= remaining_ack_bits - 1; + } + + self.peer_cumulative_ack = self.peer_cumulative_ack.max(ack); + let acknowledged = self + .unacked + .range(..=ack) + .map(|(seq, _pending)| *seq) + .collect::>(); + for seq in acknowledged.into_iter().chain(selective_acks) { + 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 00000000000..49ed78a2325 --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/reliable_stream_tests.rs @@ -0,0 +1,211 @@ +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; +use crate::relay::RelayAckState; + +#[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(RelayAckState { + ack: 0, + ack_bits: 0, + }) + .unwrap(); + assert_eq!(sender.unacked.len(), 2); + sender + .process_peer_ack(RelayAckState { + ack: 1, + ack_bits: 0, + }) + .unwrap(); + assert_eq!(sender.unacked.len(), 1); + assert_eq!(sender.unacked_bytes, b"second".len()); + sender + .process_peer_ack(RelayAckState { + ack: 2, + ack_bits: 0, + }) + .unwrap(); + assert_eq!(sender.unacked.len(), 0); + assert_eq!(sender.unacked_bytes, 0); +} + +#[test] +fn selective_ack_releases_out_of_order_cached_ciphertext() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + let first = sender.admit_ciphertext(b"first".to_vec(), now).unwrap(); + sender.admit_ciphertext(b"second".to_vec(), now).unwrap(); + let third = sender.admit_ciphertext(b"third".to_vec(), now).unwrap(); + sender.admit_ciphertext(b"fourth".to_vec(), now).unwrap(); + let ack_state = RelayAckState { + ack: 0, + ack_bits: 0b1010, + }; + + sender.process_peer_ack(ack_state).unwrap(); + sender.process_peer_ack(ack_state).unwrap(); + + assert_eq!( + sender.unacked.keys().copied().collect::>(), + vec![1, 3] + ); + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), Some(first)); + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), Some(third)); + assert_eq!(sender.next_retry_due(now + RESEND_AFTER), None); +} + +#[test] +fn rejects_inconsistent_or_unsent_ack_metadata() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + sender + .admit_ciphertext(b"ciphertext".to_vec(), now) + .unwrap(); + + assert!( + sender + .process_peer_ack(RelayAckState { + ack: 0, + ack_bits: 1, + }) + .is_err() + ); + assert!( + sender + .process_peer_ack(RelayAckState { + ack: 0, + ack_bits: 0b100, + }) + .is_err() + ); + assert!( + sender + .process_peer_ack(RelayAckState { + ack: 2, + ack_bits: 0, + }) + .is_err() + ); + assert_eq!(sender.unacked.len(), 1); +} + +#[test] +fn selective_acks_do_not_slide_the_cumulative_send_window() { + let now = Instant::now(); + let mut sender = ReliableSender::default(); + for _ in 0..MAX_UNACKED_SEGMENTS { + sender.admit_ciphertext(vec![0x5a], now).unwrap(); + } + + sender + .process_peer_ack(RelayAckState { + ack: 0, + ack_bits: u32::MAX - 1, + }) + .unwrap(); + + assert_eq!(sender.unacked.keys().copied().collect::>(), vec![1]); + assert!(!sender.can_admit_ciphertext(/*ciphertext_len*/ 1)); + assert!(sender.admit_ciphertext(vec![0x5a], now).is_err()); + + sender + .process_peer_ack(RelayAckState { + ack: 1, + ack_bits: 0, + }) + .unwrap(); + assert!(sender.can_admit_ciphertext(/*ciphertext_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 b8d477f5dbb..a57da4a6811 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)] @@ -83,6 +87,17 @@ pub(crate) enum RelayFrameBodyKind { Handshake, } +/// One coherent cumulative/selective acknowledgement snapshot. +/// +/// Bit `i` acknowledges sequence `ack + 1 + i`. Reliable Noise endpoints +/// pass this value as one unit so concurrent readers and writers cannot pair a +/// newer cumulative frontier with stale selective bits. +#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)] +pub(crate) struct RelayAckState { + pub(crate) ack: u32, + pub(crate) ack_bits: u32, +} + impl RelayMessageFrame { pub(crate) fn data(stream_id: String, seq: u32, payload: Vec) -> Self { Self { @@ -99,6 +114,41 @@ impl RelayMessageFrame { } } + /// Build one reliable Noise data frame with the receiver's ack state. + /// + /// 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_state: RelayAckState, + seq: u32, + payload: Vec, + ) -> Self { + let mut frame = Self::data(stream_id, seq, payload); + frame.ack = ack_state.ack; + frame.ack_bits = ack_state.ack_bits; + frame + } + + /// Build an unsequenced acknowledgement for a reliable stream. + pub(crate) fn ack(stream_id: String, ack_state: RelayAckState) -> Self { + Self { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + ack: ack_state.ack, + ack_bits: ack_state.ack_bits, + body: Some(relay_message_frame::Body::AckFrame(RelayAck {})), + } + } + + pub(crate) fn ack_state(&self) -> RelayAckState { + RelayAckState { + ack: self.ack, + ack_bits: self.ack_bits, + } + } + pub(crate) fn resume(stream_id: String) -> Self { Self { version: RELAY_MESSAGE_FRAME_VERSION, @@ -553,6 +603,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 +628,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 +785,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 +894,7 @@ where }); } RelayFrameBodyKind::Data => { + let peer_ack = frame.ack_state(); // 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 +908,27 @@ where } continue; }; + if let Err(error) = stream.process_peer_ack(peer_ack) { + 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 +936,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_state()) + { + 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 +1008,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 da7cb5296d4..b31d11163a4 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 918cba7393b..6d3b4de4853 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(()) +}