diff --git a/codex-rs/exec-server/src/noise_relay/environment.rs b/codex-rs/exec-server/src/noise_relay/environment.rs index 6ae96e2380..cb17ad199e 100644 --- a/codex-rs/exec-server/src/noise_relay/environment.rs +++ b/codex-rs/exec-server/src/noise_relay/environment.rs @@ -1,6 +1,4 @@ use std::collections::HashMap; -use std::sync::Arc; -use std::sync::Mutex; use std::time::Duration; use futures::SinkExt; @@ -8,7 +6,6 @@ use futures::StreamExt; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::sync::mpsc; -use tokio::sync::watch; use tokio::task::JoinSet; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; @@ -18,29 +15,20 @@ use tracing::warn; use crate::ExecServerError; use crate::connection::CHANNEL_CAPACITY; -use crate::connection::JsonRpcConnection; -use crate::connection::JsonRpcConnectionEvent; -use crate::connection::JsonRpcTransport; use crate::noise_channel::NoiseChannelIdentity; use crate::noise_channel::NoiseChannelPublicKey; -use crate::noise_channel::NoiseTransport; use crate::noise_channel::PendingResponderHandshake; use crate::noise_channel::noise_channel_prologue; -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::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::RelayFrameBodyKind; use crate::relay::decode_relay_message_frame; use crate::relay::encode_relay_message_frame; -use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; use crate::server::ConnectionProcessor; -// This value is already part of the relay wire contract. Keep it stable even -// though the source module now uses the more precise Noise terminology. -const NOISE_RELAY_RESET_REASON: &str = "secure_relay_protocol_error"; const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128; const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096; const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32; @@ -76,9 +64,23 @@ pub(crate) async fn run_noise_multiplexed_environment( S: AsyncRead + AsyncWrite + Unpin + Send + 'static, V: HarnessKeyValidator + Clone + 'static, { - let mut websocket = stream; + let (mut websocket_sink, mut websocket_stream) = stream.split(); let (physical_outgoing_tx, mut physical_outgoing_rx) = mpsc::channel::>(CHANNEL_CAPACITY); + let (closed_stream_tx, mut closed_stream_rx) = + mpsc::channel::(MAX_ACTIVE_NOISE_RELAY_STREAMS); + // A separate writer task is required because this state machine also + // produces resets and handshake responses. If the same task both sent into + // and drained the bounded outgoing channel, backpressure could make it wait + // on itself and stop servicing the physical websocket. + let mut physical_writer_task = tokio::spawn(async move { + while let Some(encoded) = physical_outgoing_rx.recv().await { + if let Err(error) = websocket_sink.send(Message::Binary(encoded.into())).await { + debug!("Noise multiplexed environment websocket write failed: {error}"); + break; + } + } + }); let mut streams: HashMap = HashMap::new(); let mut pending_handshakes: HashMap = HashMap::new(); let mut validation_tasks: JoinSet = JoinSet::new(); @@ -89,12 +91,21 @@ pub(crate) async fn run_noise_multiplexed_environment( // malicious authorization request must not block existing streams or // prevent other handshakes from being received and bounded. let frame = tokio::select! { - maybe_encoded = physical_outgoing_rx.recv() => { - let Some(encoded) = maybe_encoded else { - break; - }; - if websocket.send(Message::Binary(encoded.into())).await.is_err() { - break; + writer_result = &mut physical_writer_task => { + if let Err(error) = writer_result { + warn!("Noise multiplexed environment websocket writer failed: {error}"); + } + break; + } + Some(closed_stream) = closed_stream_rx.recv() => { + // A writer can finish after its peer resets and reuses the same + // routing ID. Remove only the exact authenticated stream + // instance that produced this close notification. + let is_current = streams + .get(&closed_stream.stream_id) + .is_some_and(|stream| stream.is_instance(closed_stream.instance_id)); + if is_current { + streams.remove(&closed_stream.stream_id); } continue; } @@ -117,14 +128,18 @@ pub(crate) async fn run_noise_multiplexed_environment( else { continue; }; - if let Err(error) = validation_result.result { - warn!("Noise relay harness key validation failed: {error}"); - send_reset(&physical_outgoing_tx, validation_result.stream_id).await; + if validation_result.result.is_err() { + // Validators receive the short-lived authorization. + // Keep their error text out of logs even though the + // registry implementation below also sanitizes + // response bodies. + warn!("Noise relay harness key validation failed"); + send_reset(&physical_outgoing_tx, validation_result.stream_id); continue; } if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { warn!("Noise relay has too many active streams"); - send_reset(&physical_outgoing_tx, validation_result.stream_id).await; + send_reset(&physical_outgoing_tx, validation_result.stream_id); continue; } @@ -135,7 +150,7 @@ pub(crate) async fn run_noise_multiplexed_environment( Ok(completed) => completed, Err(error) => { warn!("failed to complete Noise relay handshake: {error}"); - send_reset(&physical_outgoing_tx, validation_result.stream_id).await; + send_reset(&physical_outgoing_tx, validation_result.stream_id); continue; } }; @@ -143,9 +158,13 @@ pub(crate) async fn run_noise_multiplexed_environment( validation_result.stream_id.clone(), response, ); + // The shared state machine must never wait behind an + // overloaded writer queue. If a successful handshake + // response cannot be queued immediately, close this + // physical connection rather than expose a half-open + // virtual stream. if physical_outgoing_tx - .send(encode_relay_message_frame(&response)) - .await + .try_send(encode_relay_message_frame(&response)) .is_err() { break; @@ -154,8 +173,10 @@ pub(crate) async fn run_noise_multiplexed_environment( validation_result.stream_id.clone(), spawn_noise_virtual_stream( validation_result.stream_id, + validation_result.validation_id, processor.clone(), physical_outgoing_tx.clone(), + closed_stream_tx.clone(), transport, ), ); @@ -165,14 +186,14 @@ pub(crate) async fn run_noise_multiplexed_environment( let stream_ids = pending_handshakes.keys().cloned().collect::>(); pending_handshakes.clear(); for stream_id in stream_ids { - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); } } None => {} } continue; } - incoming_message = websocket.next() => match incoming_message { + incoming_message = websocket_stream.next() => match incoming_message { Some(Ok(Message::Binary(payload))) => match decode_relay_message_frame(payload.as_ref()) { Ok(frame) => frame, Err(error) => { @@ -206,17 +227,17 @@ pub(crate) async fn run_noise_multiplexed_environment( // Bound all pre-authentication state before doing expensive // hybrid cryptography or starting an external validation. if streams.contains_key(&stream_id) || pending_handshakes.contains_key(&stream_id) { - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { warn!("Noise relay has too many active streams"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } if validation_tasks.len() >= MAX_PENDING_HANDSHAKE_VALIDATIONS { warn!("Noise relay has too many pending harness key validations"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } let prologue = match noise_channel_prologue( @@ -227,7 +248,7 @@ pub(crate) async fn run_noise_multiplexed_environment( Ok(prologue) => prologue, Err(error) => { warn!("failed to build Noise relay prologue: {error}"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } }; @@ -235,7 +256,7 @@ pub(crate) async fn run_noise_multiplexed_environment( Ok(request) => request, Err(error) => { warn!("failed to read Noise relay handshake frame: {error}"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } }; @@ -244,7 +265,7 @@ pub(crate) async fn run_noise_multiplexed_environment( Ok(pending) => pending, Err(error) => { warn!("failed to read Noise relay handshake request: {error}"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } }; @@ -260,12 +281,12 @@ pub(crate) async fn run_noise_multiplexed_environment( } Ok(_) => { warn!("Noise relay handshake authorization is too long"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } Err(_) => { warn!("Noise relay handshake authorization is not UTF-8"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } }; @@ -273,7 +294,7 @@ pub(crate) async fn run_noise_multiplexed_environment( let validation_id = next_validation_id; let Some(next_id) = next_validation_id.checked_add(1) else { warn!("Noise relay harness key validation id exhausted"); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; }; next_validation_id = next_id; @@ -310,11 +331,11 @@ pub(crate) async fn run_noise_multiplexed_environment( } RelayFrameBodyKind::Data => { // Data before handshake completion is always invalid. Removing - // a pending handshake also ensures a peer cannot keep its - // authorization task alive while sending application records. + // pending state makes the time-bounded validation result stale, + // so it can never complete a stream after this protocol error. let Some(stream) = streams.get_mut(&stream_id) else { pending_handshakes.remove(&stream_id); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; }; let data = match frame.into_data() { @@ -322,20 +343,23 @@ pub(crate) async fn run_noise_multiplexed_environment( Err(error) => { warn!("dropping malformed Noise relay data frame: {error}"); streams.remove(&stream_id); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); continue; } }; - if let Err(error) = stream.receive_data(data).await { + if let Err(error) = stream.receive_data(data) { warn!("failed to process Noise relay payload: {error}"); streams.remove(&stream_id); - send_reset(&physical_outgoing_tx, stream_id).await; + send_reset(&physical_outgoing_tx, stream_id); } } RelayFrameBodyKind::Reset => { pending_handshakes.remove(&stream_id); if let Some(stream) = streams.remove(&stream_id) { - stream.disconnect(frame.into_reset_reason()).await; + // Reset is cleartext relay control and is not authenticated + // by Noise. Honor its availability effect, but never + // forward attacker-controlled reason text into logs. + stream.disconnect(/*reason*/ None); } } RelayFrameBodyKind::Ack @@ -345,7 +369,14 @@ pub(crate) async fn run_noise_multiplexed_environment( } for (_stream_id, stream) in streams { - stream.disconnect(/*reason*/ None).await; + stream.disconnect(/*reason*/ None); + } + // Dropping the JoinSet below aborts any still-running registry validations. + // Await an abort only when the select loop did not already consume the + // writer result. + if !physical_writer_task.is_finished() { + physical_writer_task.abort(); + let _ = physical_writer_task.await; } } @@ -360,138 +391,10 @@ struct HarnessKeyValidationResult { result: Result<(), ExecServerError>, } -struct NoiseVirtualStream { - incoming_tx: mpsc::Sender, - disconnected_tx: watch::Sender, - transport: Arc>, - inbound_ciphertexts: OrderedCiphertextFrames, - inbound_decoder: JsonRpcMessageDecoder, -} - -impl NoiseVirtualStream { - async fn disconnect(self, reason: Option) { - let _ = self.disconnected_tx.send(true); - let _ = self - .incoming_tx - .send(JsonRpcConnectionEvent::Disconnected { reason }) - .await; - } - - async fn receive_data(&mut self, data: RelayData) -> Result<(), ExecServerError> { - // Relay sequence ordering is enforced before taking the transport lock - // and decrypting. Each virtual stream owns one ordered Noise nonce - // space shared by its reader and writer transport halves. - for ciphertext in self.inbound_ciphertexts.push(data.seq, data.payload)? { - let plaintext = { - let mut transport = self - .transport - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - transport.decrypt(&ciphertext).map_err(|error| { - ExecServerError::Protocol(format!("Noise relay decryption failed: {error}")) - })? - }; - for message in self.inbound_decoder.push(&plaintext)? { - self.incoming_tx - .send(JsonRpcConnectionEvent::Message(message)) - .await - .map_err(|_| ExecServerError::Closed)?; - } - } - Ok(()) - } -} - -fn spawn_noise_virtual_stream( - stream_id: String, - processor: ConnectionProcessor, - physical_outgoing_tx: mpsc::Sender>, - transport: NoiseTransport, -) -> NoiseVirtualStream { - let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); - 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 writer_transport = Arc::clone(&transport); - 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 { - // Frame first, then split into bounded Noise records. Each record - // receives one checked relay sequence and is encrypted exactly - // once, preserving the implicit Noise sending 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; - } - }; - 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}"); - break 'writer; - } - }; - let ciphertext = { - let mut transport = writer_transport - .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}"); - 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; - } - } - } - - // Tell the harness to discard this virtual stream whenever its writer - // exits, including processor shutdown or a cryptographic/send failure. - // Otherwise the peer could wait indefinitely on a dead stream. - let reset = - RelayMessageFrame::reset(writer_stream_id, NOISE_RELAY_RESET_REASON.to_string()); - let _ = physical_outgoing_tx - .send(encode_relay_message_frame(&reset)) - .await; - }); - - let connection = JsonRpcConnection { - outgoing_tx: json_outgoing_tx, - incoming_rx, - disconnected_rx, - task_handles: vec![writer_task], - transport: JsonRpcTransport::External, - }; - tokio::spawn(async move { - processor.run_connection(connection).await; - }); - - NoiseVirtualStream { - incoming_tx, - disconnected_tx, - transport, - inbound_ciphertexts: OrderedCiphertextFrames::default(), - inbound_decoder: JsonRpcMessageDecoder::default(), - } -} - -async fn send_reset(physical_outgoing_tx: &mpsc::Sender>, stream_id: String) { +fn send_reset(physical_outgoing_tx: &mpsc::Sender>, stream_id: String) { let reset = RelayMessageFrame::reset(stream_id, NOISE_RELAY_RESET_REASON.to_string()); + // Resets are best effort. Untrusted relay input must never block the shared + // state machine behind an overloaded physical writer queue. let _ = physical_outgoing_tx - .send(encode_relay_message_frame(&reset)) - .await; + .try_send(encode_relay_message_frame(&reset)); } diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs new file mode 100644 index 0000000000..82e36c05af --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -0,0 +1,181 @@ +use std::sync::Arc; +use std::sync::Mutex; + +use tokio::sync::mpsc; +use tokio::sync::watch; +use tracing::warn; + +use crate::ExecServerError; +use crate::connection::CHANNEL_CAPACITY; +use crate::connection::JsonRpcConnection; +use crate::connection::JsonRpcConnectionEvent; +use crate::connection::JsonRpcTransport; +use crate::noise_channel::NoiseTransport; +use crate::noise_relay::NOISE_RELAY_RESET_REASON; +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::relay::encode_relay_message_frame; +use crate::relay_proto::RelayData; +use crate::relay_proto::RelayMessageFrame; +use crate::server::ConnectionProcessor; + +/// 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. +pub(super) struct ClosedNoiseVirtualStream { + pub(super) stream_id: String, + pub(super) instance_id: u64, +} + +/// One authenticated JSON-RPC stream carried by the executor's physical relay. +/// +/// Inbound delivery is intentionally nonblocking. An overloaded or abandoned +/// stream fails independently instead of stalling every stream multiplexed over +/// the same physical websocket. +pub(super) struct NoiseVirtualStream { + incoming_tx: mpsc::Sender, + disconnected_tx: watch::Sender, + transport: Arc>, + inbound_ciphertexts: OrderedCiphertextFrames, + inbound_decoder: JsonRpcMessageDecoder, + instance_id: u64, +} + +impl NoiseVirtualStream { + pub(super) fn disconnect(self, reason: Option) { + let _ = self.disconnected_tx.send(true); + let _ = self + .incoming_tx + .try_send(JsonRpcConnectionEvent::Disconnected { reason }); + } + + pub(super) fn is_instance(&self, instance_id: u64) -> bool { + self.instance_id == instance_id + } + + pub(super) fn receive_data(&mut self, data: RelayData) -> Result<(), ExecServerError> { + // Relay sequence ordering is enforced before taking the transport lock + // and decrypting. Each virtual stream owns one ordered Noise nonce + // space shared by its reader and writer transport halves. + for ciphertext in self.inbound_ciphertexts.push(data.seq, data.payload)? { + let plaintext = { + let mut transport = self + .transport + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + transport.decrypt(&ciphertext).map_err(|error| { + ExecServerError::Protocol(format!("Noise relay decryption failed: {error}")) + })? + }; + for message in self.inbound_decoder.push(&plaintext)? { + self.incoming_tx + .try_send(JsonRpcConnectionEvent::Message(message)) + .map_err(|_| { + ExecServerError::Protocol( + "Noise virtual stream inbound queue is full or closed".to_string(), + ) + })?; + } + } + Ok(()) + } +} + +pub(super) fn spawn_noise_virtual_stream( + stream_id: String, + instance_id: u64, + processor: ConnectionProcessor, + physical_outgoing_tx: mpsc::Sender>, + closed_stream_tx: mpsc::Sender, + transport: NoiseTransport, +) -> NoiseVirtualStream { + let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + 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 writer_transport = Arc::clone(&transport); + 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 { + // Frame first, then split into bounded Noise records. Each record + // receives one checked relay sequence and is encrypted exactly + // once, preserving the implicit Noise sending 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; + } + }; + 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}"); + break 'writer; + } + }; + let ciphertext = { + let mut transport = writer_transport + .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}"); + 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; + } + } + } + + // Reset is best effort because an overloaded physical writer must not + // keep this dead stream alive. The reliable local close notification + // below lets the shared state machine reap the exact stream instance. + let closed_stream = ClosedNoiseVirtualStream { + stream_id: writer_stream_id.clone(), + instance_id, + }; + 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 _ = closed_stream_tx.send(closed_stream).await; + }); + + let connection = JsonRpcConnection { + outgoing_tx: json_outgoing_tx, + incoming_rx, + disconnected_rx, + task_handles: vec![writer_task], + transport: JsonRpcTransport::External, + }; + tokio::spawn(async move { + processor.run_connection(connection).await; + }); + + NoiseVirtualStream { + incoming_tx, + disconnected_tx, + transport, + inbound_ciphertexts: OrderedCiphertextFrames::default(), + inbound_decoder: JsonRpcMessageDecoder::default(), + instance_id, + } +} diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 8e8b5ce796..ce09bfb916 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -30,6 +30,11 @@ use crate::relay::encode_relay_message_frame; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; +// Reset frames are cleartext relay control and are not authenticated by Noise. +// Preserve the availability signal while replacing attacker-controlled reason +// text before it reaches disconnect diagnostics. +const NOISE_RELAY_RESET_DISCONNECT_REASON: &str = "Noise relay stream reset"; + /// Adapt one harness rendezvous websocket into an authenticated JSON-RPC connection. /// /// The returned connection is not usable until the background task completes @@ -199,9 +204,7 @@ where send_disconnected( &incoming_tx, &disconnected_tx, - frame - .into_reset_reason() - .unwrap_or_else(|| "Noise relay reset during handshake".to_string()), + NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(), ) .await; return; @@ -306,9 +309,12 @@ where } } Ok(RelayFrameBodyKind::Reset) => { - let reason = frame.into_reset_reason(); let _ = incoming_tx - .send(JsonRpcConnectionEvent::Disconnected { reason }) + .send(JsonRpcConnectionEvent::Disconnected { + reason: Some( + NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(), + ), + }) .await; break; } diff --git a/codex-rs/exec-server/src/noise_relay/mod.rs b/codex-rs/exec-server/src/noise_relay/mod.rs index 2399a8cc8c..d171d8aadc 100644 --- a/codex-rs/exec-server/src/noise_relay/mod.rs +++ b/codex-rs/exec-server/src/noise_relay/mod.rs @@ -1,3 +1,4 @@ +mod executor_stream; mod environment; mod harness; mod message_framing; @@ -11,6 +12,10 @@ pub(crate) use environment::HarnessKeyValidator; pub(crate) use environment::run_noise_multiplexed_environment; pub(crate) use harness::noise_harness_connection_from_websocket; +// This value is already part of the relay wire contract. Keep it stable even +// though the source module now uses the more precise Noise terminology. +const NOISE_RELAY_RESET_REASON: &str = "secure_relay_protocol_error"; + // This bounds allocation in tungstenite before protobuf and Noise record // validation run. It comfortably fits one maximum Noise record plus metadata. const MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE: usize = 256 * 1024; diff --git a/codex-rs/exec-server/src/relay_proto.rs b/codex-rs/exec-server/src/relay_proto.rs index 4da5be5e69..da7cb5296d 100644 --- a/codex-rs/exec-server/src/relay_proto.rs +++ b/codex-rs/exec-server/src/relay_proto.rs @@ -4,5 +4,6 @@ mod generated; pub(crate) use generated::RelayData; pub(crate) use generated::RelayHandshake; pub(crate) use generated::RelayMessageFrame; +pub(crate) use generated::RelayReset; pub(crate) use generated::RelayResume; pub(crate) use generated::relay_message_frame; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 86e0b12439..2504e3d632 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -3,29 +3,20 @@ use std::time::Duration; use codex_api::SharedAuthProvider; use reqwest::StatusCode; use serde::Deserialize; -use serde::Serialize; use tokio::time::sleep; -use tokio::time::timeout; -use tokio_tungstenite::connect_async_with_config; -use tracing::info; +use tokio_tungstenite::connect_async; use tracing::warn; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use crate::ExecServerError; use crate::ExecServerRuntimePaths; -use crate::NoiseChannelIdentity; -use crate::NoiseChannelPublicKey; -use crate::noise_relay::HarnessKeyValidator; -use crate::noise_relay::noise_relay_websocket_config; -use crate::noise_relay::run_noise_multiplexed_environment; use crate::relay::run_multiplexed_environment; use crate::server::ConnectionProcessor; +mod noise; + const ERROR_BODY_PREVIEW_BYTES: usize = 4096; -const ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); -const MAX_REMOTE_ENVIRONMENT_ID_LEN: usize = 256; -const REMOTE_RENDEZVOUS_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Clone)] struct EnvironmentRegistryClient { @@ -51,16 +42,16 @@ impl EnvironmentRegistryClient { auth_provider, http: reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) - .timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT) .build()?, }) } /// Register using the original body-less registry contract. /// - /// This remains the default so a new exec-server can still connect through - /// registries and harnesses that have not rolled out Noise relay support. - async fn register_legacy_environment( + /// This method intentionally preserves the legacy request shape and timeout + /// behavior. Noise support is opt-in and must not silently alter existing + /// remote exec-server registrations. + async fn register_environment( &self, environment_id: &str, ) -> Result { @@ -76,64 +67,6 @@ impl EnvironmentRegistryClient { self.parse_json_response(response).await } - /// Register the exec-server's static Noise identity with a Noise-aware registry. - /// - /// Supplying this request body is the protocol-level opt in. Registries can - /// therefore distinguish Noise registrations from the legacy body-less - /// contract without guessing based on binary version or rollout state. - async fn register_noise_environment( - &self, - environment_id: &str, - executor_public_key: &NoiseChannelPublicKey, - ) -> Result { - let response = self - .http - .post(endpoint_url( - &self.base_url, - &format!("/cloud/environment/{environment_id}/register"), - )) - .headers(self.auth_provider.to_auth_headers()) - .json(&EnvironmentRegistryRegistrationRequest { - security_profile: NOISE_RELAY_SECURITY_PROFILE, - executor_public_key, - }) - .send() - .await?; - self.parse_json_response(response).await - } - - async fn validate_harness_key( - &self, - environment_id: &str, - executor_registration_id: &str, - harness_public_key: &NoiseChannelPublicKey, - harness_key_authorization: &str, - ) -> Result<(), ExecServerError> { - let response = self - .http - .post(endpoint_url( - &self.base_url, - &format!("/cloud/environment/{environment_id}/validate"), - )) - .headers(self.auth_provider.to_auth_headers()) - .json(&EnvironmentRegistryHarnessKeyValidationRequest { - executor_registration_id, - harness_public_key, - harness_key_authorization, - }) - .send() - .await?; - let response = self - .parse_json_response::(response) - .await?; - if !response.valid { - return Err(ExecServerError::Protocol( - "environment registry rejected Noise relay harness key".to_string(), - )); - } - Ok(()) - } - async fn parse_json_response( &self, response: reqwest::Response, @@ -155,64 +88,12 @@ impl EnvironmentRegistryClient { } } -const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; - -#[derive(Serialize)] -struct EnvironmentRegistryRegistrationRequest<'a> { - security_profile: &'static str, - executor_public_key: &'a NoiseChannelPublicKey, -} - #[derive(Debug, Clone, Eq, PartialEq, Deserialize)] struct EnvironmentRegistryRegistrationResponse { environment_id: String, url: String, } -#[derive(Debug, Clone, Eq, PartialEq, Deserialize)] -struct EnvironmentRegistryNoiseRegistrationResponse { - environment_id: String, - url: String, - security_profile: String, - executor_registration_id: String, -} - -#[derive(Serialize)] -struct EnvironmentRegistryHarnessKeyValidationRequest<'a> { - executor_registration_id: &'a str, - harness_public_key: &'a NoiseChannelPublicKey, - harness_key_authorization: &'a str, -} - -#[derive(Deserialize)] -struct EnvironmentRegistryHarnessKeyValidationResponse { - valid: bool, -} - -#[derive(Clone)] -struct RegistryHarnessKeyValidator { - client: EnvironmentRegistryClient, - environment_id: String, - executor_registration_id: String, -} - -impl HarnessKeyValidator for RegistryHarnessKeyValidator { - async fn validate_harness_key( - &self, - harness_public_key: &NoiseChannelPublicKey, - authorization: &str, - ) -> Result<(), ExecServerError> { - self.client - .validate_harness_key( - &self.environment_id, - &self.executor_registration_id, - harness_public_key, - authorization, - ) - .await - } -} - /// Protocol used for an exec-server's registered remote relay. /// /// Legacy is intentionally the default during rollout. Noise must be selected @@ -220,40 +101,13 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { /// contract until both endpoints are ready for Noise. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] pub enum RemoteRelayProtocol { + /// Original cleartext JSON-RPC relay protocol. #[default] Legacy, + /// Authenticated, end-to-end encrypted Noise relay protocol. Noise, } -enum RemoteRelayProtocolState { - Legacy, - Noise(NoiseChannelIdentity), -} - -enum RegisteredRemoteEnvironment { - Legacy(EnvironmentRegistryRegistrationResponse), - Noise { - response: EnvironmentRegistryNoiseRegistrationResponse, - identity: NoiseChannelIdentity, - }, -} - -impl RegisteredRemoteEnvironment { - fn environment_id(&self) -> &str { - match self { - Self::Legacy(response) => &response.environment_id, - Self::Noise { response, .. } => &response.environment_id, - } - } - - fn websocket_url(&self) -> &str { - match self { - Self::Legacy(response) => &response.url, - Self::Noise { response, .. } => &response.url, - } - } -} - /// Configuration for registering an exec-server for remote use. #[derive(Clone)] pub struct RemoteEnvironmentConfig { @@ -303,92 +157,27 @@ pub async fn run_remote_environment( let client = EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; let processor = ConnectionProcessor::new(runtime_paths); - let protocol_state = match config.relay_protocol { - RemoteRelayProtocol::Legacy => RemoteRelayProtocolState::Legacy, - RemoteRelayProtocol::Noise => RemoteRelayProtocolState::Noise( - NoiseChannelIdentity::generate().map_err(|error| { - ExecServerError::Protocol(format!( - "failed to generate Noise relay identity: {error}" - )) - })?, - ), - }; + if config.relay_protocol == RemoteRelayProtocol::Noise { + return noise::run_remote_environment(&config, &client, processor).await; + } + let mut backoff = Duration::from_secs(1); loop { - let registration = match &protocol_state { - RemoteRelayProtocolState::Legacy => RegisteredRemoteEnvironment::Legacy( - client - .register_legacy_environment(&config.environment_id) - .await?, - ), - RemoteRelayProtocolState::Noise(identity) => { - let response = client - .register_noise_environment(&config.environment_id, &identity.public_key()) - .await?; - if response.security_profile != NOISE_RELAY_SECURITY_PROFILE { - return Err(ExecServerError::Protocol(format!( - "environment registry returned unsupported security profile `{}`", - response.security_profile - ))); - } - RegisteredRemoteEnvironment::Noise { - response, - identity: identity.clone(), - } - } - }; - if registration.environment_id() != config.environment_id { - return Err(ExecServerError::Protocol( - "environment registry returned a different environment id".to_string(), - )); - } - let environment_id = registration.environment_id(); - info!( - "codex exec-server remote environment registered with environment_id {environment_id}" + let response = client.register_environment(&config.environment_id).await?; + eprintln!( + "codex exec-server remote environment registered with environment_id {}", + response.environment_id ); - let websocket_config = match ®istration { - RegisteredRemoteEnvironment::Legacy(_) => None, - RegisteredRemoteEnvironment::Noise { .. } => Some(noise_relay_websocket_config()), - }; - match timeout( - REMOTE_RENDEZVOUS_CONNECT_TIMEOUT, - connect_async_with_config( - registration.websocket_url(), - websocket_config, - /*disable_nagle*/ false, - ), - ) - .await - { - Ok(Ok((websocket, _))) => { + match connect_async(response.url.as_str()).await { + Ok((websocket, _)) => { backoff = Duration::from_secs(1); - match registration { - RegisteredRemoteEnvironment::Legacy(_) => { - run_multiplexed_environment(websocket, processor.clone()).await; - } - RegisteredRemoteEnvironment::Noise { response, identity } => { - run_noise_multiplexed_environment( - websocket, - processor.clone(), - response.environment_id, - response.executor_registration_id.clone(), - identity, - RegistryHarnessKeyValidator { - client: client.clone(), - environment_id: config.environment_id.clone(), - executor_registration_id: response.executor_registration_id, - }, - ) - .await; - } - } + run_multiplexed_environment(websocket, processor.clone()).await; } - Ok(Err(err)) => { + Err(err) => { warn!("failed to connect remote exec-server websocket: {err}"); } - Err(_) => warn!("timed out connecting remote exec-server websocket"), } sleep(backoff).await; @@ -397,32 +186,12 @@ pub async fn run_remote_environment( } fn normalize_environment_id(environment_id: String) -> Result { + let environment_id = environment_id.trim().to_string(); if environment_id.is_empty() { return Err(ExecServerError::EnvironmentRegistryConfig( "environment id is required for remote exec-server registration".to_string(), )); } - if environment_id.trim() != environment_id { - return Err(ExecServerError::EnvironmentRegistryConfig( - "environment id must not contain surrounding whitespace".to_string(), - )); - } - if environment_id.len() > MAX_REMOTE_ENVIRONMENT_ID_LEN { - return Err(ExecServerError::EnvironmentRegistryConfig(format!( - "environment id cannot be longer than {MAX_REMOTE_ENVIRONMENT_ID_LEN} characters" - ))); - } - // The ID is interpolated into authenticated registry request paths below. - // Keep it to one URL path component so a caller cannot use a delimiter to - // redirect the exec-server's registration credential to another endpoint. - if !environment_id - .chars() - .all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_') - { - return Err(ExecServerError::EnvironmentRegistryConfig( - "environment id must contain only ASCII letters, numbers, '-' or '_'".to_string(), - )); - } Ok(environment_id) } diff --git a/codex-rs/exec-server/src/remote/noise.rs b/codex-rs/exec-server/src/remote/noise.rs new file mode 100644 index 0000000000..5b34c88d5c --- /dev/null +++ b/codex-rs/exec-server/src/remote/noise.rs @@ -0,0 +1,270 @@ +//! Explicitly selected Noise registration and encrypted relay runtime. +//! +//! The legacy remote-exec path stays in the parent module. Keeping the Noise +//! path here makes the opt-in boundary visible and keeps its stricter registry, +//! identifier, timeout, and websocket requirements from changing legacy +//! behavior accidentally. + +use std::time::Duration; + +use reqwest::StatusCode; +use serde::Deserialize; +use serde::Serialize; +use tokio::time::sleep; +use tokio::time::timeout; +use tokio_tungstenite::connect_async_with_config; +use tracing::info; +use tracing::warn; + +use super::EnvironmentRegistryClient; +use super::RemoteEnvironmentConfig; +use super::endpoint_url; +use crate::ExecServerError; +use crate::NoiseChannelIdentity; +use crate::NoiseChannelPublicKey; +use crate::noise_relay::HarnessKeyValidator; +use crate::noise_relay::noise_relay_websocket_config; +use crate::noise_relay::run_noise_multiplexed_environment; +use crate::server::ConnectionProcessor; + +const ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_EXECUTOR_REGISTRATION_ID_LEN: usize = 256; +const MAX_REMOTE_ENVIRONMENT_ID_LEN: usize = 256; +const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; +const REMOTE_RENDEZVOUS_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); + +impl EnvironmentRegistryClient { + /// Register the exec-server's static Noise identity with a Noise-aware registry. + /// + /// Supplying this request body is the protocol-level opt in. Registries can + /// therefore distinguish Noise registrations from the legacy body-less + /// contract without guessing based on binary version or rollout state. + async fn register_noise_environment( + &self, + environment_id: &str, + executor_public_key: &NoiseChannelPublicKey, + ) -> Result { + let response = self + .http + .post(endpoint_url( + &self.base_url, + &format!("/cloud/environment/{environment_id}/register"), + )) + .headers(self.auth_provider.to_auth_headers()) + .timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT) + .json(&EnvironmentRegistryRegistrationRequest { + security_profile: NOISE_RELAY_SECURITY_PROFILE, + executor_public_key, + }) + .send() + .await?; + self.parse_json_response(response).await + } + + /// Validate the authenticated harness key without exposing its authorization. + async fn validate_harness_key( + &self, + environment_id: &str, + executor_registration_id: &str, + harness_public_key: &NoiseChannelPublicKey, + harness_key_authorization: &str, + ) -> Result<(), ExecServerError> { + let response = self + .http + .post(endpoint_url( + &self.base_url, + &format!("/cloud/environment/{environment_id}/validate"), + )) + .headers(self.auth_provider.to_auth_headers()) + .timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT) + .json(&EnvironmentRegistryHarnessKeyValidationRequest { + executor_registration_id, + harness_public_key, + harness_key_authorization, + }) + .send() + .await?; + let status = response.status(); + if !status.is_success() { + // This request contains the short-lived harness authorization. + // Never propagate a response body that might echo it into logs or + // user-visible error chains. + if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) { + return Err(ExecServerError::EnvironmentRegistryAuth(format!( + "environment registry harness key validation authentication failed ({status})" + ))); + } + return Err(ExecServerError::EnvironmentRegistryHttp { + status, + code: None, + message: "environment registry harness key validation failed".to_string(), + }); + } + let response = response + .json::() + .await?; + if !response.valid { + return Err(ExecServerError::Protocol( + "environment registry rejected Noise relay harness key".to_string(), + )); + } + Ok(()) + } +} + +#[derive(Serialize)] +struct EnvironmentRegistryRegistrationRequest<'a> { + security_profile: &'static str, + executor_public_key: &'a NoiseChannelPublicKey, +} + +#[derive(Deserialize)] +struct EnvironmentRegistryNoiseRegistrationResponse { + environment_id: String, + url: String, + security_profile: String, + executor_registration_id: String, +} + +#[derive(Serialize)] +struct EnvironmentRegistryHarnessKeyValidationRequest<'a> { + executor_registration_id: &'a str, + harness_public_key: &'a NoiseChannelPublicKey, + harness_key_authorization: &'a str, +} + +#[derive(Deserialize)] +struct EnvironmentRegistryHarnessKeyValidationResponse { + valid: bool, +} + +#[derive(Clone)] +struct RegistryHarnessKeyValidator { + client: EnvironmentRegistryClient, + environment_id: String, + executor_registration_id: String, +} + +impl HarnessKeyValidator for RegistryHarnessKeyValidator { + async fn validate_harness_key( + &self, + harness_public_key: &NoiseChannelPublicKey, + authorization: &str, + ) -> Result<(), ExecServerError> { + self.client + .validate_harness_key( + &self.environment_id, + &self.executor_registration_id, + harness_public_key, + authorization, + ) + .await + } +} + +/// Run the Noise registration and encrypted relay loop. +/// +/// A new executor identity is generated once per process invocation and reused +/// across physical reconnects. Every registry registration still receives a new +/// registration ID, which is bound into each virtual stream's Noise prologue. +pub(super) async fn run_remote_environment( + config: &RemoteEnvironmentConfig, + client: &EnvironmentRegistryClient, + processor: ConnectionProcessor, +) -> Result<(), ExecServerError> { + validate_environment_id(&config.environment_id)?; + let identity = NoiseChannelIdentity::generate().map_err(|error| { + ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}")) + })?; + let mut backoff = Duration::from_secs(1); + + loop { + let response = client + .register_noise_environment(&config.environment_id, &identity.public_key()) + .await?; + if response.environment_id != config.environment_id { + return Err(ExecServerError::Protocol( + "environment registry returned a different environment id".to_string(), + )); + } + if response.security_profile != NOISE_RELAY_SECURITY_PROFILE { + return Err(ExecServerError::Protocol(format!( + "environment registry returned unsupported security profile `{}`", + response.security_profile + ))); + } + validate_executor_registration_id(&response.executor_registration_id)?; + let environment_id = &response.environment_id; + info!( + "codex exec-server Noise environment registered with environment_id {environment_id}" + ); + + match timeout( + REMOTE_RENDEZVOUS_CONNECT_TIMEOUT, + connect_async_with_config( + response.url.as_str(), + Some(noise_relay_websocket_config()), + /*disable_nagle*/ false, + ), + ) + .await + { + Ok(Ok((websocket, _))) => { + backoff = Duration::from_secs(1); + let executor_registration_id = response.executor_registration_id; + run_noise_multiplexed_environment( + websocket, + processor.clone(), + response.environment_id, + executor_registration_id.clone(), + identity.clone(), + RegistryHarnessKeyValidator { + client: client.clone(), + environment_id: config.environment_id.clone(), + executor_registration_id, + }, + ) + .await; + } + Ok(Err(err)) => warn!("failed to connect Noise remote exec-server websocket: {err}"), + Err(_) => warn!("timed out connecting Noise remote exec-server websocket"), + } + + sleep(backoff).await; + backoff = (backoff * 2).min(Duration::from_secs(30)); + } +} + +fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> { + if environment_id.len() > MAX_REMOTE_ENVIRONMENT_ID_LEN { + return Err(ExecServerError::EnvironmentRegistryConfig(format!( + "environment id cannot be longer than {MAX_REMOTE_ENVIRONMENT_ID_LEN} characters" + ))); + } + // The ID is interpolated into authenticated registry request paths below. + // Keep it to one URL path component so a caller cannot use a delimiter to + // redirect the exec-server's registration credential to another endpoint. + if !environment_id + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_') + { + return Err(ExecServerError::EnvironmentRegistryConfig( + "environment id must contain only ASCII letters, numbers, '-' or '_'".to_string(), + )); + } + Ok(()) +} + +fn validate_executor_registration_id( + executor_registration_id: &str, +) -> Result<(), ExecServerError> { + if executor_registration_id.is_empty() + || executor_registration_id.trim() != executor_registration_id + || executor_registration_id.len() > MAX_EXECUTOR_REGISTRATION_ID_LEN + { + return Err(ExecServerError::Protocol( + "environment registry returned an invalid executor registration id".to_string(), + )); + } + Ok(()) +}