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

feat(exec-server): add ACK-based Noise relay retries

Browse Source
This commit is contained in:
Anton Panasenko
2026-07-07 12:04:28 -07:00
parent cca16a1087
commit e99be8cf65
13 changed files with 1048 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/README.md
+15 -8
View File
@@ -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`.
codex-rs/exec-server/src/noise_relay/executor_stream.rs
+174 -37
View File
@@ -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<JsonRpcConnectionEvent>,
disconnected_tx: watch::Sender<bool>,
transport: Arc<Mutex<NoiseTransport>>,
reliable_sender: Arc<Mutex<ReliableSender>>,
inbound_ack: Arc<AtomicU32>,
pending_ack: Arc<Mutex<Option<u32>>>,
writer_wakeup: Arc<Notify>,
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<u8>, 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,
codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs
+64 -13
View File
@@ -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))
}
codex-rs/exec-server/src/noise_relay/harness.rs
+106 -33
View File
@@ -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<u8>, 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(),
codex-rs/exec-server/src/noise_relay/harness_tests.rs
+163 -26
View File
@@ -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<Message>,
control: &ControlledWebSocketHandle,
) -> Result<RelayData> {
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<Message>,
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<Message>,
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<Message>,
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 {
codex-rs/exec-server/src/noise_relay/mod.rs
+1 -12
View File
@@ -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<u32, ExecServerError> {
// 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)
}
codex-rs/exec-server/src/noise_relay/ordered_ciphertext.rs
+25 -3
View File
@@ -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<u32, Vec<u8>>,
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<u8>,
) -> Result<Vec<Vec<u8>>, 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());
codex-rs/exec-server/src/noise_relay/ordered_ciphertext_tests.rs
+22 -9
View File
@@ -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::<Vec<u8>>::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::<Vec<u8>>::new()
);
assert_eq!(
frames.push(/*seq*/ 1, b"replacement".to_vec()).unwrap(),
frames.push(/*seq*/ 2, b"replacement".to_vec()).unwrap(),
Vec::<Vec<u8>>::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::<Vec<u8>>::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);
}
codex-rs/exec-server/src/noise_relay/reliable_stream.rs
+189
View File
@@ -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<u8>,
}
#[derive(Debug)]
struct UnackedCiphertext {
payload: Vec<u8>,
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<u32, UnackedCiphertext>,
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<u8>,
now: Instant,
) -> Result<OutboundCiphertext, ExecServerError> {
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::<Vec<_>>();
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<OutboundCiphertext> {
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;
codex-rs/exec-server/src/noise_relay/reliable_stream_tests.rs
+120
View File
@@ -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());
}
codex-rs/exec-server/src/relay.rs
+100 -7
View File
@@ -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<u8>) -> 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<String, NoiseVirtualStream> = 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<String> = HashSet::new();
let mut retired_stream_order: VecDeque<String> = VecDeque::new();
let mut pending_handshakes: HashMap<String, PendingHandshake> = HashMap::new();
let mut validation_tasks: JoinSet<HarnessKeyValidationResult> = 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<String>,
retired_stream_order: &mut VecDeque<String>,
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<Vec<u8>>, stream_id: String) {
codex-rs/exec-server/src/relay_proto.rs
+1
View File
@@ -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;
codex-rs/exec-server/tests/relay.rs
+68 -1
View File
@@ -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<Mutex<Option<(u32, Vec<u8>)>>>;
#[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<TcpStream>,
mut harness: WebSocketStream<TcpStream>,
captured_frames: Arc<Mutex<Vec<Vec<u8>>>>,
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<bool> {
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<Vec<Vec<u8>>>, message: &Message) {
if let Message::Binary(bytes) = message {
captured_frames
@@ -273,3 +311,32 @@ fn assert_relay_data_is_encrypted(captured_frames: &Mutex<Vec<Vec<u8>>>) -> Resu
);
Ok(())
}
fn assert_dropped_ciphertext_was_retried(
captured_frames: &Mutex<Vec<Vec<u8>>>,
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(())
}