From 6e50b22e55cfc986b54d33979acb1b7511efddff Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Mon, 15 Jun 2026 17:39:00 -0700 Subject: [PATCH] exec-server: default remote transport to Noise (#26245) ## Why The transport in [openai/codex#26242](https://github.com/openai/codex/pull/26242) needs to be used by every remote orchestrator-to-executor connection before JSON-RPC traffic starts. ## Changes - Generates one executor Noise identity when remote exec-server starts and registers its public key. - Creates a harness identity for each physical remote environment connection. - Fetches a fresh registry bundle before connecting and validates the authenticated harness key before completing the executor handshake. - Multiplexes encrypted logical streams over the existing executor WebSocket. - Adds bounded stream, handshake-failure, and reassembly state. - Adds safe lifecycle diagnostics without logging keys, authorizations, plaintext, or ciphertext. - Covers reconnects, replay rejection, validation failure, framing limits, and encrypted JSON-RPC tool traffic. ## Stack 1. [openai/codex#26242](https://github.com/openai/codex/pull/26242): Noise channel and relay transport 2. **[openai/codex#26245](https://github.com/openai/codex/pull/26245)**: remote registration and runtime activation ## Verification - `just test -p codex-exec-server` - `just fix -p codex-exec-server` - `just bazel-lock-check` - `cargo shear` --------- Co-authored-by: Codex --- codex-rs/cli/Cargo.toml | 2 +- codex-rs/exec-server/README.md | 8 +- codex-rs/exec-server/src/noise_channel.rs | 56 ++ .../exec-server/src/noise_channel_tests.rs | 175 ++++++ .../src/noise_relay/executor_stream.rs | 193 ++++++ .../src/noise_relay/executor_stream_tests.rs | 70 +++ .../exec-server/src/noise_relay/harness.rs | 37 +- codex-rs/exec-server/src/noise_relay/mod.rs | 3 + codex-rs/exec-server/src/relay.rs | 590 ++++++++++++------ codex-rs/exec-server/src/relay_noise_tests.rs | 358 +++++++++++ codex-rs/exec-server/src/relay_proto.rs | 1 + codex-rs/exec-server/src/remote.rs | 227 ++++++- .../exec-server/src/remote/noise_tests.rs | 163 +++++ codex-rs/exec-server/tests/relay.rs | 510 +++++++-------- 14 files changed, 1902 insertions(+), 491 deletions(-) create mode 100644 codex-rs/exec-server/src/noise_relay/executor_stream.rs create mode 100644 codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs create mode 100644 codex-rs/exec-server/src/relay_noise_tests.rs create mode 100644 codex-rs/exec-server/src/remote/noise_tests.rs diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index 8896759f9..7be1146df 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -19,7 +19,7 @@ workspace = true [dependencies] anyhow = { workspace = true } -clap = { workspace = true, features = ["derive"] } +clap = { workspace = true, features = ["derive", "env"] } clap_complete = { workspace = true } codex-app-server = { workspace = true } codex-app-server-daemon = { workspace = true } diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 54b00ea92..0f70b85df 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -26,6 +26,8 @@ The CLI entrypoint supports: Remote mode registers the local exec-server with the environment registry, then reconnects to the service-provided rendezvous websocket as the environment. +Remote communication uses the Noise relay contract; the registry and harness +must support it. It uses the standard Codex ChatGPT sign-in state; run `codex login` first when remote registration needs authentication. Containerized callers that receive an Agent Identity JWT in `CODEX_ACCESS_TOKEN` can opt into that auth path with @@ -45,7 +47,7 @@ codex exec-server \ Wire framing: - local websocket: one JSON-RPC message per websocket frame -- remote websocket: binary protobuf relay frames carrying JSON-RPC payloads +- Noise remote websocket: binary protobuf relay frames carrying encrypted payloads ## Remote Relay Message Format @@ -57,13 +59,13 @@ identity plus endpoint-owned reliability metadata: ```text version stream_id -body // data | ack_frame | resume | reset | heartbeat +body // handshake | data | ack_frame | resume | reset | heartbeat ack // highest contiguous peer segment seq received ack_bits // bitset for peer segment seqs after ack seq // data only: segment sequence number segment_index // data only: 0-based index within message segment_count // data only: number of segments in message -payload // data only: JSON-RPC message bytes or segment bytes +payload // handshake bytes or encrypted data record next_seq // resume only: next sender seq reason // reset only: reset reason ``` diff --git a/codex-rs/exec-server/src/noise_channel.rs b/codex-rs/exec-server/src/noise_channel.rs index 42cc2e6b3..062e54a68 100644 --- a/codex-rs/exec-server/src/noise_channel.rs +++ b/codex-rs/exec-server/src/noise_channel.rs @@ -167,6 +167,62 @@ impl InitiatorHandshake { } } +/// Executor-side handshake state while harness authorization is pending. +/// This is not a usable transport until the registry accepts the authenticated +/// harness key. +pub(crate) struct PendingResponderHandshake { + handshake: Handshake, + pub(crate) initiator_public_key: NoiseChannelPublicKey, + pub(crate) payload: Vec, +} + +impl PendingResponderHandshake { + /// Parse the first IK message and recover the authenticated harness key. + /// Callers must authorize that key before calling [`Self::complete`]. + pub(crate) fn read_request( + identity: &NoiseChannelIdentity, + prologue: &[u8], + request: &[u8], + ) -> Result { + ensure_noise_frame_len(request.len(), "handshake request is too large")?; + let params = HybridHandshakeParams::new(noise_hybrid_ik(), false) + .with_prologue(prologue) + .with_s(identity.dh.clone()) + .with_s_kem(identity.kem.clone()); + let mut handshake = Handshake::new(params)?; + let mut payload = [0u8; MAX_MESSAGE_LEN]; + let payload_len = handshake.read_message(request, &mut payload)?; + // Clatter exposes this key only after the first IK message authenticates. + let remote = handshake + .get_remote_static() + .ok_or(NoiseChannelError::InvalidMessage( + "handshake request is missing initiator static key", + ))?; + let initiator_public_key = NoiseChannelPublicKey { + suite: NOISE_CHANNEL_SUITE.to_string(), + x25519_public_key: STANDARD.encode(remote.dh()), + mlkem768_public_key: STANDARD.encode(remote.kem().as_slice()), + }; + Ok(Self { + handshake, + initiator_public_key, + payload: payload[..payload_len].to_vec(), + }) + } + + /// Finish the handshake after the registry authorizes the harness key. + pub(crate) fn complete(mut self) -> Result<(NoiseTransport, Vec), NoiseChannelError> { + let mut response = [0u8; MAX_MESSAGE_LEN]; + let response_len = self.handshake.write_message(&[], &mut response)?; + Ok(( + NoiseTransport { + transport: self.handshake.finalize()?, + }, + response[..response_len].to_vec(), + )) + } +} + /// Established channel with independent implicit send and receive nonces. /// Relay records must be ordered before decryption, and a logical record must /// not be encrypted again for retry. diff --git a/codex-rs/exec-server/src/noise_channel_tests.rs b/codex-rs/exec-server/src/noise_channel_tests.rs index 1638eef89..b298e8e58 100644 --- a/codex-rs/exec-server/src/noise_channel_tests.rs +++ b/codex-rs/exec-server/src/noise_channel_tests.rs @@ -6,6 +6,181 @@ use super::NOISE_CHANNEL_SUITE; use super::NoiseChannelError; use super::NoiseChannelIdentity; use super::NoiseChannelPublicKey; +use super::PendingResponderHandshake; +use super::noise_channel_prologue; + +#[test] +fn hybrid_ik_roundtrip_authenticates_both_endpoints() { + let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); + let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); + let prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + let authorization = b"harness-key-authorization"; + + let (initiator_handshake, request) = InitiatorHandshake::start( + &initiator, + &responder.public_key(), + &prologue, + authorization, + ) + .expect("start initiator handshake"); + let responder_handshake = + PendingResponderHandshake::read_request(&responder, &prologue, &request) + .expect("read responder handshake"); + + assert_eq!( + &responder_handshake.initiator_public_key, + &initiator.public_key() + ); + assert_eq!(responder_handshake.payload.as_slice(), authorization); + + let (mut responder_transport, response) = responder_handshake + .complete() + .expect("complete responder handshake"); + let mut initiator_transport = initiator_handshake + .finish(&response) + .expect("complete initiator handshake"); + + let request_ciphertext = initiator_transport + .encrypt(b"request") + .expect("encrypt request"); + assert_ne!(request_ciphertext, b"request"); + assert_eq!( + responder_transport + .decrypt(&request_ciphertext) + .expect("decrypt request"), + b"request" + ); + + let response_ciphertext = responder_transport + .encrypt(b"response") + .expect("encrypt response"); + assert_ne!(response_ciphertext, b"response"); + assert_eq!( + initiator_transport + .decrypt(&response_ciphertext) + .expect("decrypt response"), + b"response" + ); +} + +#[test] +fn initiator_rejects_wrong_responder_key() { + let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); + let expected_responder = NoiseChannelIdentity::generate().expect("generate expected identity"); + let actual_responder = NoiseChannelIdentity::generate().expect("generate actual identity"); + let prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + + let (_initiator_handshake, request) = InitiatorHandshake::start( + &initiator, + &expected_responder.public_key(), + &prologue, + b"authorization", + ) + .expect("start initiator handshake"); + + assert!( + PendingResponderHandshake::read_request(&actual_responder, &prologue, &request).is_err() + ); +} + +#[test] +fn responder_rejects_mismatched_prologue() { + let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); + let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); + let initiator_prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + let responder_prologue = noise_channel_prologue("env-1", "registration-1", "stream-2"); + let (_initiator_handshake, request) = InitiatorHandshake::start( + &initiator, + &responder.public_key(), + &initiator_prologue, + b"authorization", + ) + .expect("start initiator handshake"); + + assert!( + PendingResponderHandshake::read_request(&responder, &responder_prologue, &request).is_err() + ); +} + +#[test] +fn prologue_encoding_is_stable_and_unambiguous() { + let prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + + assert_eq!( + prologue, + b"\x00\x00\x00\x00\x00\x00\x00\x20codex-exec-server-relay-noise/v1\ + \x00\x00\x00\x00\x00\x00\x00\x05env-1\ + \x00\x00\x00\x00\x00\x00\x00\x0eregistration-1\ + \x00\x00\x00\x00\x00\x00\x00\x08stream-1" + .to_vec() + ); +} + +#[test] +fn transport_rejects_tampered_ciphertext() { + let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); + let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); + let prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + let (initiator_handshake, request) = InitiatorHandshake::start( + &initiator, + &responder.public_key(), + &prologue, + b"authorization", + ) + .expect("start initiator handshake"); + let responder_handshake = + PendingResponderHandshake::read_request(&responder, &prologue, &request) + .expect("read responder handshake"); + let (mut responder_transport, response) = responder_handshake + .complete() + .expect("complete responder handshake"); + let mut initiator_transport = initiator_handshake + .finish(&response) + .expect("complete initiator handshake"); + let mut ciphertext = initiator_transport + .encrypt(b"request") + .expect("encrypt request"); + ciphertext[0] ^= 1; + + assert!(responder_transport.decrypt(&ciphertext).is_err()); +} + +#[test] +fn transport_rejects_replayed_ciphertext() { + let initiator = NoiseChannelIdentity::generate().expect("generate initiator identity"); + let responder = NoiseChannelIdentity::generate().expect("generate responder identity"); + let prologue = noise_channel_prologue("env-1", "registration-1", "stream-1"); + let (initiator_handshake, request) = InitiatorHandshake::start( + &initiator, + &responder.public_key(), + &prologue, + b"authorization", + ) + .expect("start initiator handshake"); + let responder_handshake = + PendingResponderHandshake::read_request(&responder, &prologue, &request) + .expect("read responder handshake"); + let (mut responder_transport, response) = responder_handshake + .complete() + .expect("complete responder handshake"); + let mut initiator_transport = initiator_handshake + .finish(&response) + .expect("complete initiator handshake"); + let ciphertext = initiator_transport + .encrypt(b"request") + .expect("encrypt request"); + + assert_eq!( + responder_transport + .decrypt(&ciphertext) + .expect("decrypt request"), + b"request" + ); + assert!(matches!( + responder_transport.decrypt(&ciphertext), + Err(NoiseChannelError::Transport(_)) + )); +} #[test] fn public_key_validation_rejects_unknown_suite() { 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 000000000..57a05579d --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -0,0 +1,193 @@ +//! One executor-side virtual stream after the Noise handshake. +//! +//! The environment loop owns reads and a per-stream task owns writes. They share +//! `NoiseTransport` because its send and receive nonces live in the same value; +//! the mutex is never held across `.await`. + +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(crate) struct ClosedNoiseVirtualStream { + pub(crate) stream_id: String, + pub(crate) 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(crate) struct NoiseVirtualStream { + incoming_tx: mpsc::Sender, + disconnected_tx: watch::Sender, + transport: Arc>, + inbound_ciphertexts: OrderedCiphertextFrames, + inbound_decoder: JsonRpcMessageDecoder, + pub(crate) instance_id: u64, +} + +impl NoiseVirtualStream { + pub(crate) fn disconnect(self, reason: Option) { + let _ = self.disconnected_tx.send(true); + let _ = self + .incoming_tx + .try_send(JsonRpcConnectionEvent::Disconnected { reason }); + } + + /// 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> { + 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(()) + } +} + +/// Start JSON-RPC processing for a completed handshake. +/// +/// The returned value is the read half; the spawned task owns outbound framing +/// and reports its instance ID on exit so stream-ID reuse is safe. +pub(crate) 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 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; + } + }; + 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; + } + } + } + + // The reset is best effort; the local close notification is not. + 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::Plain, + }; + tokio::spawn(async move { + processor.run_connection(connection).await; + let _ = processor_closed_stream_tx + .send(ClosedNoiseVirtualStream { + stream_id: processor_stream_id, + instance_id, + }) + .await; + }); + + NoiseVirtualStream { + incoming_tx, + disconnected_tx, + transport, + inbound_ciphertexts: OrderedCiphertextFrames::default(), + inbound_decoder: JsonRpcMessageDecoder::default(), + instance_id, + } +} + +#[cfg(test)] +#[path = "executor_stream_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs new file mode 100644 index 000000000..9119236cf --- /dev/null +++ b/codex-rs/exec-server/src/noise_relay/executor_stream_tests.rs @@ -0,0 +1,70 @@ +use std::time::Duration; + +use anyhow::Result; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use tokio::sync::mpsc; +use tokio::time::timeout; + +use super::ClosedNoiseVirtualStream; +use super::spawn_noise_virtual_stream; +use crate::ExecServerRuntimePaths; +use crate::connection::CHANNEL_CAPACITY; +use crate::noise_channel::InitiatorHandshake; +use crate::noise_channel::NoiseChannelIdentity; +use crate::noise_channel::PendingResponderHandshake; +use crate::noise_relay::message_framing::frame_jsonrpc_message; +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 (physical_outgoing_tx, _physical_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (closed_stream_tx, mut 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 ciphertext = harness_transport.encrypt(&frame_jsonrpc_message(&message)?)?; + stream.receive_data(RelayData { + seq: 0, + segment_index: 0, + segment_count: 1, + payload: ciphertext, + })?; + + assert!(matches!( + timeout(Duration::from_secs(1), closed_stream_rx.recv()).await?, + Some(ClosedNoiseVirtualStream { + stream_id, + instance_id: 7, + }) if stream_id == "stream-1" + )); + Ok(()) +} diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index a857ebfbe..ef75f2526 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -13,7 +13,9 @@ use futures::StreamExt; use tokio::sync::mpsc; use tokio::sync::watch; use tokio_tungstenite::tungstenite::Message; +use tracing::Instrument; use tracing::debug; +use tracing::info; use tracing::warn; use uuid::Uuid; @@ -52,6 +54,11 @@ pub(crate) struct NoiseHarnessConnectionArgs { pub(crate) harness_key_authorization: String, } +// 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 @@ -79,6 +86,13 @@ where let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); let (disconnected_tx, disconnected_rx) = watch::channel(false); + let stream_span = tracing::debug_span!( + "noise_relay.stream", + noise_side = "harness", + environment_id = %environment_id, + executor_registration_id = %executor_registration_id, + stream_id = %stream_id, + ); let websocket_task = tokio::spawn(async move { let mut websocket = stream; @@ -203,7 +217,14 @@ where } }; match initiator_handshake.finish(&response) { - Ok(transport) => break transport, + Ok(transport) => { + info!( + noise_event = "handshake", + noise_outcome = "ok", + "Noise harness handshake completed" + ); + break transport; + } Err(error) => { send_disconnected( &incoming_tx, @@ -219,9 +240,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; @@ -326,9 +345,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; } @@ -366,7 +388,8 @@ where } } let _ = disconnected_tx.send(true); - }); + } + .instrument(stream_span)); JsonRpcConnection { outgoing_tx, diff --git a/codex-rs/exec-server/src/noise_relay/mod.rs b/codex-rs/exec-server/src/noise_relay/mod.rs index 5e2bddcd0..56fabd5cb 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 @@ +pub(crate) mod executor_stream; mod harness; mod message_framing; mod ordered_ciphertext; @@ -9,6 +10,8 @@ use crate::ExecServerError; pub(crate) use harness::NoiseHarnessConnectionArgs; pub(crate) use harness::noise_harness_connection_from_websocket; +pub(crate) const NOISE_RELAY_RESET_REASON: &str = "noise_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.rs b/codex-rs/exec-server/src/relay.rs index e0c4d7ccb..1202ec2d1 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::time::Duration; use codex_app_server_protocol::JSONRPCMessage; use futures::Sink; @@ -10,9 +11,12 @@ 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; use tokio_tungstenite::tungstenite::Message; use tracing::debug; +use tracing::info; use tracing::warn; use uuid::Uuid; @@ -22,14 +26,29 @@ use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; +use crate::noise_channel::NoiseChannelIdentity; +use crate::noise_channel::NoiseChannelPublicKey; +use crate::noise_channel::PendingResponderHandshake; +use crate::noise_channel::noise_channel_prologue; +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::RelayData; use crate::relay_proto::RelayHandshake; use crate::relay_proto::RelayMessageFrame; +use crate::relay_proto::RelayReset; use crate::relay_proto::RelayResume; use crate::relay_proto::relay_message_frame; use crate::server::ConnectionProcessor; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; +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 HARNESS_KEY_VALIDATION_TIMEOUT: Duration = Duration::from_secs(10); + #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub(crate) enum RelayFrameBodyKind { Data, @@ -80,6 +99,16 @@ impl RelayMessageFrame { } } + pub(crate) fn reset(stream_id: String, reason: String) -> Self { + Self { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Reset(RelayReset { reason })), + } + } + pub(crate) fn validate(&self) -> Result { if self.version != RELAY_MESSAGE_FRAME_VERSION { return Err(ExecServerError::Protocol(format!( @@ -394,57 +423,216 @@ where } } -pub(crate) async fn run_multiplexed_environment( +/// Validates that a Noise-authenticated harness public key is authorized. +/// +/// Implementations must consult an authority independent of rendezvous. The +/// exec-server invokes this after parsing the first IK message and before +/// completing the responder handshake. +pub(crate) trait HarnessKeyValidator: Send + Sync { + fn validate_harness_key( + &self, + harness_public_key: &NoiseChannelPublicKey, + authorization: &str, + ) -> impl std::future::Future> + Send; +} + +/// Serve authenticated virtual JSON-RPC streams over one executor websocket. +/// +/// Parsing the first Noise message authenticates the harness key. Only a +/// successful registry check turns that pending handshake into a virtual stream. +#[tracing::instrument( + level = "debug", + skip_all, + fields( + noise_side = "executor", + environment_id = %environment_id, + executor_registration_id = %executor_registration_id, + ) +)] +pub(crate) async fn run_multiplexed_environment( stream: WebSocketStream, processor: ConnectionProcessor, + environment_id: String, + executor_registration_id: String, + identity: NoiseChannelIdentity, + validator: V, ) where 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); + // Use a separate writer so this loop never waits on the channel it drains. + let mut physical_writer_task = tokio::spawn(async move { + let mut keepalive = tokio::time::interval_at( + tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, + WEBSOCKET_KEEPALIVE_INTERVAL, + ); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + let message = tokio::select! { + encoded = physical_outgoing_rx.recv() => { + let Some(encoded) = encoded else { + break; + }; + Message::Binary(encoded.into()) + } + _ = keepalive.tick() => Message::Ping(Vec::new().into()), + }; + if let Err(error) = websocket_sink.send(message).await { + warn!("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(); + let mut failed_handshakes = 0usize; + let mut next_validation_id = 0u64; - let mut keepalive = tokio::time::interval_at( - tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, - WEBSOCKET_KEEPALIVE_INTERVAL, - ); - keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut streams: HashMap = HashMap::new(); - 'websocket: loop { + loop { + // Registry calls run separately so a slow check does not block the relay. 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 stream ID may have been reused before this writer exits. + // 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); } continue; } - _ = keepalive.tick() => { - if websocket.send(Message::Ping(Vec::new().into())).await.is_err() { - break; - } - continue; - } - incoming_message = websocket.next() => match incoming_message { - Some(Ok(Message::Binary(payload))) => { - match decode_relay_message_frame(payload.as_ref()) { - Ok(frame) => frame, - Err(err) => { - warn!("dropping malformed relay message frame from harness: {err}"); + validation_result = validation_tasks.join_next(), if !validation_tasks.is_empty() => { + match validation_result { + Some(Ok(validation_result)) => { + // The stream ID may have been reused while validation ran. + let is_current = pending_handshakes + .get(&validation_result.stream_id) + .is_some_and(|pending| { + pending.validation_id == validation_result.validation_id + }); + if !is_current { continue; } + let Some(pending) = + pending_handshakes.remove(&validation_result.stream_id) + else { + continue; + }; + if validation_result.result.is_err() { + // Validator errors may contain authorization details. + warn!( + noise_event = "authorization", + noise_outcome = "error", + noise_reason = "authorization_failed", + "Noise harness authorization failed" + ); + debug!( + stream_id = validation_result.stream_id, + "Noise harness authorization failure details" + ); + send_reset(&physical_outgoing_tx, validation_result.stream_id); + if failed_handshake_budget_exhausted(&mut failed_handshakes) { + warn!("closing Noise relay after repeated handshake failures"); + break; + } + 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); + continue; + } + + // This is the only point where the responder completes + // IK and exposes a JSON-RPC stream: Noise authenticated + // the harness key and the registry authorized it. + let (transport, response) = match pending.handshake.complete() { + Ok(completed) => completed, + Err(error) => { + warn!("failed to complete Noise relay handshake: {error}"); + send_reset(&physical_outgoing_tx, validation_result.stream_id); + if failed_handshake_budget_exhausted(&mut failed_handshakes) { + warn!("closing Noise relay after repeated handshake failures"); + break; + } + continue; + } + }; + let response = RelayMessageFrame::handshake( + validation_result.stream_id.clone(), + response, + ); + // Do not leave a half-open stream if the handshake reply + // cannot be queued immediately. + if physical_outgoing_tx + .try_send(encode_relay_message_frame(&response)) + .is_err() + { + break; + } + info!( + noise_event = "handshake", + noise_outcome = "ok", + "Noise executor handshake completed" + ); + debug!( + stream_id = validation_result.stream_id, + active_streams = streams.len() + 1, + "Noise executor stream activated" + ); + streams.insert( + 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, + ), + ); } + Some(Err(error)) => { + warn!("Noise relay harness key validation task failed: {error}"); + let stream_ids = pending_handshakes.keys().cloned().collect::>(); + pending_handshakes.clear(); + for stream_id in stream_ids { + send_reset(&physical_outgoing_tx, stream_id); + } + } + None => {} } + continue; + } + 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) => { + warn!("dropping malformed Noise relay frame from harness: {error}"); + continue; + } + }, Some(Ok(Message::Close(_))) | None => break, Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue, Some(Ok(Message::Text(_))) => { - warn!("dropping non-binary relay message frame from harness"); + warn!("dropping non-binary Noise relay frame from harness"); continue; } - Some(Err(err)) => { - debug!("multiplexed environment websocket read failed: {err}"); + Some(Err(error)) => { + debug!("Noise multiplexed environment websocket read failed: {error}"); break; } } @@ -452,128 +640,207 @@ pub(crate) async fn run_multiplexed_environment( let kind = match frame.validate() { Ok(kind) => kind, - Err(err) => { - warn!("dropping invalid relay message frame: {err}"); + Err(error) => { + warn!("dropping invalid Noise relay frame: {error}"); continue; } }; - + let stream_id = frame.stream_id.clone(); match kind { - RelayFrameBodyKind::Data => { - let stream_id = frame.stream_id.clone(); - let message = match frame.into_jsonrpc_message() { - Ok(message) => message, - Err(err) => { - warn!("dropping malformed relay data message frame: {err}"); + 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) { + send_reset(&physical_outgoing_tx, stream_id); + continue; + } + // Removing pending state makes the in-flight validation result stale. + if pending_handshakes.remove(&stream_id).is_some() { + send_reset(&physical_outgoing_tx, stream_id); + if failed_handshake_budget_exhausted(&mut failed_handshakes) { + warn!("closing Noise relay after repeated handshake failures"); + break; + } + continue; + } + if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS { + warn!("Noise relay has too many active streams"); + 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); + continue; + } + let prologue = + noise_channel_prologue(&environment_id, &executor_registration_id, &stream_id); + let request = match frame.into_handshake_payload() { + Ok(request) => request, + Err(error) => { + warn!("failed to read Noise relay handshake frame: {error}"); + send_reset(&physical_outgoing_tx, stream_id); continue; } }; + let mut pending = + match PendingResponderHandshake::read_request(&identity, &prologue, &request) { + Ok(pending) => pending, + Err(error) => { + warn!("failed to read Noise relay handshake request: {error}"); + send_reset(&physical_outgoing_tx, stream_id); + if failed_handshake_budget_exhausted(&mut failed_handshakes) { + warn!("closing Noise relay after repeated handshake failures"); + break; + } + continue; + } + }; - let stream = streams.entry(stream_id.clone()).or_insert_with(|| { - spawn_virtual_stream( - stream_id.clone(), - processor.clone(), - physical_outgoing_tx.clone(), - ) - }); - let incoming_tx = stream.incoming_tx.clone(); - match send_event_with_keepalive( - &mut websocket, - &mut keepalive, - &incoming_tx, - JsonRpcConnectionEvent::Message(message), - ) - .await - { - Ok(()) => {} - Err(RelayEventSendError::IncomingClosed) => { - streams.remove(&stream_id); + // The authorization and authenticated harness key come from the + // same encrypted IK message and are validated together. + let authorization = match String::from_utf8(std::mem::take(&mut pending.payload)) { + Ok(authorization) + if authorization.len() <= MAX_HARNESS_KEY_AUTHORIZATION_BYTES => + { + Some(authorization) } - Err(RelayEventSendError::WebSocketClosed) => break 'websocket, + Ok(_) => { + warn!("Noise relay handshake authorization is too long"); + None + } + Err(_) => { + warn!("Noise relay handshake authorization is not UTF-8"); + None + } + }; + let Some(authorization) = authorization else { + send_reset(&physical_outgoing_tx, stream_id); + if failed_handshake_budget_exhausted(&mut failed_handshakes) { + warn!("closing Noise relay after repeated handshake failures"); + break; + } + continue; + }; + let harness_public_key = pending.initiator_public_key.clone(); + let validation_id = next_validation_id; + next_validation_id += 1; + pending_handshakes.insert( + stream_id.clone(), + PendingHandshake { + validation_id, + handshake: pending, + }, + ); + let validator = validator.clone(); + + // Failed validation leaves no transport state and sends only a + // generic reset. + validation_tasks.spawn(async move { + let result = match timeout( + HARNESS_KEY_VALIDATION_TIMEOUT, + validator.validate_harness_key(&harness_public_key, &authorization), + ) + .await + { + Ok(result) => result, + Err(_) => Err(ExecServerError::Protocol( + "timed out validating Noise relay harness key".to_string(), + )), + }; + HarnessKeyValidationResult { + stream_id, + validation_id, + result, + } + }); + } + RelayFrameBodyKind::Data => { + // Removing pending state also makes any in-flight validation stale. + let Some(stream) = streams.get_mut(&stream_id) else { + let canceled_pending_handshake = + pending_handshakes.remove(&stream_id).is_some(); + send_reset(&physical_outgoing_tx, stream_id); + if canceled_pending_handshake + && failed_handshake_budget_exhausted(&mut failed_handshakes) + { + warn!("closing Noise relay after repeated handshake failures"); + break; + } + continue; + }; + let data = match frame.into_data() { + Ok(data) => data, + Err(error) => { + warn!("dropping malformed Noise relay data frame: {error}"); + streams.remove(&stream_id); + send_reset(&physical_outgoing_tx, stream_id); + continue; + } + }; + 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); } } RelayFrameBodyKind::Reset => { - if let Some(stream) = streams.remove(&frame.stream_id) { - stream.disconnect(frame.into_reset_reason()).await; + pending_handshakes.remove(&stream_id); + if let Some(stream) = streams.remove(&stream_id) { + // The reset reason is unauthenticated, so do not log it. + stream.disconnect(/*reason*/ None); } } RelayFrameBodyKind::Ack | RelayFrameBodyKind::Resume - | RelayFrameBodyKind::Heartbeat - | RelayFrameBodyKind::Handshake => {} + | RelayFrameBodyKind::Heartbeat => {} } } for (_stream_id, stream) in streams { - stream.disconnect(/*reason*/ None).await; + stream.disconnect(/*reason*/ None); } - drop(physical_outgoing_tx); -} - -struct VirtualStream { - incoming_tx: mpsc::Sender, - disconnected_tx: watch::Sender, -} - -impl VirtualStream { - async fn disconnect(self, reason: Option) { - let _ = self.disconnected_tx.send(true); - let _ = self - .incoming_tx - .send(JsonRpcConnectionEvent::Disconnected { reason }) - .await; + // Dropping the JoinSet aborts any registry checks still running. + if !physical_writer_task.is_finished() { + physical_writer_task.abort(); + let _ = physical_writer_task.await; } } -fn spawn_virtual_stream( +/// Charge one failed authenticated-channel attempt to this physical relay. +/// +/// Closing after a small fixed budget prevents a peer that has not been +/// authorized from triggering unbounded hybrid handshakes or registry checks. +fn failed_handshake_budget_exhausted(failed_handshakes: &mut usize) -> bool { + *failed_handshakes += 1; + *failed_handshakes >= MAX_FAILED_NOISE_HANDSHAKES +} + +/// Responder state held while registry authorization is pending. +struct PendingHandshake { + validation_id: u64, + handshake: PendingResponderHandshake, +} + +/// `validation_id` prevents an old check from completing a reused `stream_id`. +struct HarnessKeyValidationResult { stream_id: String, - processor: ConnectionProcessor, - physical_outgoing_tx: mpsc::Sender>, -) -> VirtualStream { - 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 writer_stream_id = stream_id; - let writer_task = tokio::spawn(async move { - let mut next_seq = 0u32; - while let Some(message) = json_outgoing_rx.recv().await { - let payload = match jsonrpc_payload(&message) { - Ok(payload) => payload, - Err(err) => { - warn!("failed to serialize virtual stream JSON-RPC payload: {err}"); - break; - } - }; - let frame = RelayMessageFrame::data(writer_stream_id.clone(), next_seq, payload); - next_seq = next_seq.wrapping_add(1); - if physical_outgoing_tx - .send(encode_relay_message_frame(&frame)) - .await - .is_err() - { - break; - } - } - }); - - let connection = JsonRpcConnection { - outgoing_tx: json_outgoing_tx, - incoming_rx, - disconnected_rx, - task_handles: vec![writer_task], - transport: JsonRpcTransport::Plain, - }; - tokio::spawn(async move { - processor.run_connection(connection).await; - }); - - VirtualStream { - incoming_tx, - disconnected_tx, - } + validation_id: u64, + result: Result<(), ExecServerError>, } +/// Queue a best-effort reset without blocking the shared websocket loop. +/// Reset reasons are relay control data and are not treated as trusted text. +fn send_reset(physical_outgoing_tx: &mpsc::Sender>, stream_id: String) { + let reset = RelayMessageFrame::reset(stream_id, NOISE_RELAY_RESET_REASON.to_string()); + let _ = physical_outgoing_tx.try_send(encode_relay_message_frame(&reset)); +} + +#[cfg(test)] +#[path = "relay_noise_tests.rs"] +mod noise_tests; + #[cfg(test)] mod tests { use std::pin::Pin; @@ -584,7 +851,6 @@ mod tests { use std::task::Poll; use std::time::Duration; - use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RequestId; use futures::Sink; @@ -633,8 +899,7 @@ mod tests { } #[tokio::test] - async fn multiplexed_environment_sends_keepalive_and_processes_relay_data() -> anyhow::Result<()> - { + async fn multiplexed_environment_sends_keepalive() -> anyhow::Result<()> { let (client_websocket, mut server_websocket) = websocket_pair().await?; let runtime_paths = crate::ExecServerRuntimePaths::new( std::env::current_exe()?, @@ -644,67 +909,32 @@ mod tests { let environment_task = tokio::spawn(run_multiplexed_environment( client_websocket, ConnectionProcessor::new(runtime_paths), + "test-environment".to_string(), + "test-registration".to_string(), + NoiseChannelIdentity::generate()?, + AllowHarnessKeyValidator, )); read_keepalive_ping(&mut server_websocket).await?; - server_websocket - .send(Message::Pong(b"keepalive".to_vec().into())) - .await?; - let stream_id = "test-stream".to_string(); - let request = test_jsonrpc_message(); - server_websocket - .send(Message::Binary( - encode_relay_message_frame(&RelayMessageFrame::data( - stream_id.clone(), - /*seq*/ 0, - jsonrpc_payload(&request)?, - )) - .into(), - )) - .await?; - - let response_frame = timeout(Duration::from_secs(1), async { - loop { - let message = server_websocket - .next() - .await - .expect("websocket should stay open")?; - match message { - Message::Binary(payload) => { - return decode_relay_message_frame(payload.as_ref()) - .map_err(anyhow::Error::from); - } - Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => {} - Message::Text(text) => { - anyhow::bail!("expected binary relay frame, got text: {text}"); - } - Message::Close(_) => { - anyhow::bail!("websocket closed before relay response"); - } - } - } - }) - .await??; - let expected_message = JSONRPCMessage::Error(JSONRPCError { - id: RequestId::Integer(1), - error: crate::rpc::method_not_found( - "exec-server stub does not implement `test` yet".to_string(), - ), - }); - assert_eq!( - response_frame, - RelayMessageFrame::data( - stream_id, - /*seq*/ 0, - jsonrpc_payload(&expected_message)?, - ) - ); environment_task.abort(); let _ = environment_task.await; Ok(()) } + #[derive(Clone)] + struct AllowHarnessKeyValidator; + + impl HarnessKeyValidator for AllowHarnessKeyValidator { + async fn validate_harness_key( + &self, + _harness_public_key: &NoiseChannelPublicKey, + _authorization: &str, + ) -> Result<(), ExecServerError> { + Ok(()) + } + } + #[tokio::test] async fn send_event_with_keepalive_pings_while_incoming_queue_is_full() -> anyhow::Result<()> { let (mut websocket, _control, mut outbound_rx) = diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs new file mode 100644 index 000000000..0020efe34 --- /dev/null +++ b/codex-rs/exec-server/src/relay_noise_tests.rs @@ -0,0 +1,358 @@ +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Result; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use tokio::net::TcpListener; +use tokio::sync::Notify; +use tokio::time::timeout; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::tungstenite::Message; + +use super::HarnessKeyValidator; +use super::MAX_FAILED_NOISE_HANDSHAKES; +use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES; +use super::run_multiplexed_environment; +use crate::ExecServerError; +use crate::ExecServerRuntimePaths; +use crate::noise_channel::InitiatorHandshake; +use crate::noise_channel::NoiseChannelIdentity; +use crate::noise_channel::NoiseChannelPublicKey; +use crate::noise_channel::noise_channel_prologue; +use crate::relay::RelayFrameBodyKind; +use crate::relay::decode_relay_message_frame; +use crate::relay::encode_relay_message_frame; +use crate::relay_proto::RelayMessageFrame; +use crate::server::ConnectionProcessor; + +const ENVIRONMENT_ID: &str = "environment-1"; +const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; + +#[derive(Clone)] +struct BlockingValidator { + calls: Arc, + release: Arc, +} + +impl HarnessKeyValidator for BlockingValidator { + fn validate_harness_key( + &self, + _harness_public_key: &NoiseChannelPublicKey, + _authorization: &str, + ) -> impl std::future::Future> + Send { + let calls = Arc::clone(&self.calls); + let release = Arc::clone(&self.release); + async move { + calls.fetch_add(1, Ordering::SeqCst); + release.notified().await; + Ok(()) + } + } +} + +#[tokio::test] +async fn pending_harness_key_validation_does_not_block_new_handshakes() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (mut harness_websocket, _response) = harness_connection.await??; + + let environment_identity = NoiseChannelIdentity::generate()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let calls = Arc::new(AtomicUsize::new(0)); + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + environment_identity.clone(), + BlockingValidator { + calls: Arc::clone(&calls), + release: Arc::new(Notify::new()), + }, + )); + + for stream_id in ["stream-1", "stream-2"] { + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id); + let (_handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + b"authorization", + )?; + let frame = RelayMessageFrame::handshake(stream_id.to_string(), request); + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + } + + timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) != 2 { + tokio::task::yield_now().await; + } + }) + .await?; + + harness_websocket.close(None).await?; + timeout(Duration::from_secs(1), environment_task).await??; + Ok(()) +} + +#[tokio::test] +async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (mut harness_websocket, _response) = harness_connection.await??; + + let environment_identity = NoiseChannelIdentity::generate()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let calls = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(Notify::new()); + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + environment_identity.clone(), + BlockingValidator { + calls: Arc::clone(&calls), + release: Arc::clone(&release), + }, + )); + + let stream_id = "stream-1"; + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id); + let (_handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + b"authorization", + )?; + let frame = RelayMessageFrame::handshake(stream_id.to_string(), request); + let encoded = encode_relay_message_frame(&frame); + harness_websocket + .send(Message::Binary(encoded.clone().into())) + .await?; + timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) != 1 { + tokio::task::yield_now().await; + } + }) + .await?; + + for attempt in 1..MAX_FAILED_NOISE_HANDSHAKES { + if attempt > 1 { + harness_websocket + .send(Message::Binary(encoded.clone().into())) + .await?; + timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) != attempt { + tokio::task::yield_now().await; + } + }) + .await?; + } + harness_websocket + .send(Message::Binary(encoded.clone().into())) + .await?; + let payload = timeout(Duration::from_secs(1), async { + loop { + match harness_websocket.next().await { + Some(Ok(Message::Binary(payload))) => break Ok(payload), + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {} + Some(Ok(message)) => anyhow::bail!("expected reset frame, got {message:?}"), + Some(Err(error)) => break Err(error.into()), + None => anyhow::bail!("environment closed before sending reset"), + } + } + }) + .await??; + let reset = decode_relay_message_frame(payload.as_ref())?; + assert_eq!(reset.stream_id, stream_id); + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); + } + + harness_websocket + .send(Message::Binary(encoded.clone().into())) + .await?; + timeout(Duration::from_secs(1), async { + while calls.load(Ordering::SeqCst) != MAX_FAILED_NOISE_HANDSHAKES { + tokio::task::yield_now().await; + } + }) + .await?; + harness_websocket + .send(Message::Binary(encoded.into())) + .await?; + timeout(Duration::from_secs(1), environment_task).await??; + release.notify_waiters(); + Ok(()) +} + +#[tokio::test] +async fn oversized_harness_authorization_is_rejected_before_validation() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (mut harness_websocket, _response) = harness_connection.await??; + + let environment_identity = NoiseChannelIdentity::generate()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let calls = Arc::new(AtomicUsize::new(0)); + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + environment_identity.clone(), + BlockingValidator { + calls: Arc::clone(&calls), + release: Arc::new(Notify::new()), + }, + )); + + let stream_id = "stream-1"; + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, stream_id); + let oversized_authorization = vec![b'a'; MAX_HARNESS_KEY_AUTHORIZATION_BYTES + 1]; + let (_handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + &oversized_authorization, + )?; + let frame = RelayMessageFrame::handshake(stream_id.to_string(), request); + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + + let Message::Binary(payload) = timeout(Duration::from_secs(1), harness_websocket.next()) + .await? + .ok_or_else(|| anyhow::anyhow!("environment closed before sending reset"))?? + else { + anyhow::bail!("expected binary reset frame"); + }; + let reset = decode_relay_message_frame(payload.as_ref())?; + assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset); + assert_eq!(calls.load(Ordering::SeqCst), 0); + + harness_websocket.close(None).await?; + timeout(Duration::from_secs(1), environment_task).await??; + Ok(()) +} + +#[tokio::test] +async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (mut harness_websocket, _response) = harness_connection.await??; + + let environment_identity = NoiseChannelIdentity::generate()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + environment_identity.clone(), + BlockingValidator { + calls: Arc::new(AtomicUsize::new(0)), + release: Arc::new(Notify::new()), + }, + )); + + for attempt in 0..MAX_FAILED_NOISE_HANDSHAKES { + let stream_id = format!("malformed-{attempt}"); + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, &stream_id); + let (_handshake, mut request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + b"authorization", + )?; + let last_byte = request.last_mut().expect("handshake request is not empty"); + *last_byte ^= 1; + let frame = RelayMessageFrame::handshake(stream_id, request); + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + } + + timeout(Duration::from_secs(1), environment_task).await??; + Ok(()) +} + +#[tokio::test] +async fn repeated_early_data_during_validation_closes_the_physical_relay() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let harness_connection = tokio::spawn(connect_async(websocket_url)); + let (socket, _peer_addr) = listener.accept().await?; + let environment_websocket = accept_async(socket).await?; + let (mut harness_websocket, _response) = harness_connection.await??; + + let environment_identity = NoiseChannelIdentity::generate()?; + let harness_identity = NoiseChannelIdentity::generate()?; + let environment_task = tokio::spawn(run_multiplexed_environment( + environment_websocket, + ConnectionProcessor::new(ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?), + ENVIRONMENT_ID.to_string(), + EXECUTOR_REGISTRATION_ID.to_string(), + environment_identity.clone(), + BlockingValidator { + calls: Arc::new(AtomicUsize::new(0)), + release: Arc::new(Notify::new()), + }, + )); + + for attempt in 0..MAX_FAILED_NOISE_HANDSHAKES { + let stream_id = format!("early-data-{attempt}"); + let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, &stream_id); + let (_handshake, request) = InitiatorHandshake::start( + &harness_identity, + &environment_identity.public_key(), + &prologue, + b"authorization", + )?; + for frame in [ + RelayMessageFrame::handshake(stream_id.clone(), request), + RelayMessageFrame::data(stream_id, /*seq*/ 0, vec![0]), + ] { + harness_websocket + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await?; + } + } + + timeout(Duration::from_secs(1), environment_task).await??; + Ok(()) +} diff --git a/codex-rs/exec-server/src/relay_proto.rs b/codex-rs/exec-server/src/relay_proto.rs index 4da5be5e6..da7cb5296 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 ae7242377..7a470f3d9 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -3,18 +3,26 @@ use std::time::Duration; use codex_api::SharedAuthProvider; use reqwest::StatusCode; use serde::Deserialize; +use serde::Serialize; use tokio::time::sleep; -use tokio_tungstenite::connect_async; +use tokio_tungstenite::connect_async_with_config; +use tracing::debug; +use tracing::info; 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::noise_relay_websocket_config; +use crate::relay::HarnessKeyValidator; use crate::relay::run_multiplexed_environment; use crate::server::ConnectionProcessor; const ERROR_BODY_PREVIEW_BYTES: usize = 4096; +const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1"; #[derive(Clone)] struct EnvironmentRegistryClient { @@ -44,9 +52,12 @@ impl EnvironmentRegistryClient { }) } + /// Register the executor public key and obtain the rendezvous allocation. + /// The returned registration ID is included in each stream's Noise prologue. async fn register_environment( &self, environment_id: &str, + executor_public_key: &NoiseChannelPublicKey, ) -> Result { let response = self .http @@ -55,9 +66,37 @@ impl EnvironmentRegistryClient { &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 + let response: EnvironmentRegistryRegistrationResponse = + self.parse_json_response(response).await?; + if response.environment_id != 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 + ))); + } + info!( + noise_event = "registration", + noise_outcome = "ok", + security_profile = NOISE_RELAY_SECURITY_PROFILE, + "Noise executor registration completed" + ); + debug!( + environment_id = response.environment_id, + executor_registration_id = response.executor_registration_id, + "Noise executor registration details" + ); + Ok(response) } async fn parse_json_response( @@ -81,10 +120,89 @@ impl EnvironmentRegistryClient { } } +#[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, + 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 { + /// Authorize the harness key recovered from the first IK message. + /// Noise proves key possession; the registry decides whether that key may use + /// this executor. The authorization token and public key are checked together. + async fn validate_harness_key( + &self, + harness_public_key: &NoiseChannelPublicKey, + authorization: &str, + ) -> Result<(), ExecServerError> { + let environment_id = &self.environment_id; + let response = self + .client + .http + .post(endpoint_url( + &self.client.base_url, + &format!("/cloud/environment/{environment_id}/validate"), + )) + .headers(self.client.auth_provider.to_auth_headers()) + .json(&EnvironmentRegistryHarnessKeyValidationRequest { + executor_registration_id: &self.executor_registration_id, + harness_public_key, + harness_key_authorization: authorization, + }) + .send() + .await?; + let status = response.status(); + if !status.is_success() { + // The request contains the short-lived authorization. Do not include + // a response body that might echo it in logs or 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(()) + } } /// Configuration for registering an exec-server for remote use. @@ -123,8 +241,12 @@ impl RemoteEnvironmentConfig { } } -/// Register an exec-server for remote use and serve requests over the returned -/// rendezvous websocket. +/// Register an exec-server for remote use and serve requests over Noise. +/// +/// The executor identity is generated once per process and reused across +/// reconnects. The registration and rendezvous URL are also reused until +/// rendezvous rejects the URL, at which point the next attempt registers again. +/// The websocket carries cleartext routing metadata and encrypted payloads. pub async fn run_remote_environment( config: RemoteEnvironmentConfig, runtime_paths: ExecServerRuntimePaths, @@ -133,22 +255,62 @@ 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 identity = NoiseChannelIdentity::generate().map_err(|error| { + ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}")) + })?; let mut backoff = Duration::from_secs(1); + let mut response = client + .register_environment(&config.environment_id, &identity.public_key()) + .await?; loop { - let response = client.register_environment(&config.environment_id).await?; - eprintln!( - "codex exec-server remote environment registered with environment_id {}", - response.environment_id - ); - - match connect_async(response.url.as_str()).await { + match connect_async_with_config( + response.url.as_str(), + Some(noise_relay_websocket_config()), + /*disable_nagle*/ false, + ) + .await + { Ok((websocket, _)) => { backoff = Duration::from_secs(1); - run_multiplexed_environment(websocket, processor.clone()).await; + let executor_registration_id = response.executor_registration_id.clone(); + info!( + noise_event = "rendezvous_connection", + noise_outcome = "ok", + "Noise executor connected to rendezvous" + ); + run_multiplexed_environment( + websocket, + processor.clone(), + response.environment_id.clone(), + executor_registration_id.clone(), + identity.clone(), + RegistryHarnessKeyValidator { + client: client.clone(), + environment_id: config.environment_id.clone(), + executor_registration_id, + }, + ) + .await; } - Err(err) => { - warn!("failed to connect remote exec-server websocket: {err}"); + Err(error) => { + let registration_rejected = matches!( + &error, + tokio_tungstenite::tungstenite::Error::Http(response) + if response.status().is_client_error() + ); + warn!( + noise_event = "rendezvous_connection", + noise_outcome = "error", + noise_reason = "websocket_error", + "Noise executor failed to connect to rendezvous" + ); + debug!(error = %error, "Noise executor rendezvous connection error"); + if registration_rejected { + response = client + .register_environment(&config.environment_id, &identity.public_key()) + .await?; + } } } @@ -252,6 +414,7 @@ mod tests { use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; + use wiremock::matchers::body_partial_json; use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; @@ -281,19 +444,22 @@ mod tests { #[tokio::test] async fn register_environment_posts_with_auth_provider_headers() { let server = MockServer::start().await; - let config = RemoteEnvironmentConfig::new( - server.uri(), - "environment-requested".to_string(), - static_registry_auth_provider(), - ) - .expect("config"); + let executor_public_key = NoiseChannelIdentity::generate() + .expect("identity") + .public_key(); Mock::given(method("POST")) .and(path("/cloud/environment/environment-requested/register")) .and(header("authorization", "Bearer registry-token")) .and(header("chatgpt-account-id", "workspace-123")) + .and(body_partial_json(serde_json::json!({ + "security_profile": NOISE_RELAY_SECURITY_PROFILE, + "executor_public_key": executor_public_key.clone(), + }))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "environment_id": "env-1", - "url": "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc" + "environment_id": "environment-requested", + "url": "wss://rendezvous.test/cloud-agent/default/ws/environment/environment-requested?role=environment&sig=abc", + "security_profile": NOISE_RELAY_SECURITY_PROFILE, + "executor_registration_id": "registration-1", }))) .mount(&server) .await; @@ -301,15 +467,17 @@ mod tests { .expect("client"); let response = client - .register_environment(&config.environment_id) + .register_environment("environment-requested", &executor_public_key) .await .expect("register environment"); assert_eq!( response, EnvironmentRegistryRegistrationResponse { - environment_id: "env-1".to_string(), - url: "wss://rendezvous.test/cloud-agent/default/ws/environment/env-1?role=environment&sig=abc".to_string(), + environment_id: "environment-requested".to_string(), + url: "wss://rendezvous.test/cloud-agent/default/ws/environment/environment-requested?role=environment&sig=abc".to_string(), + security_profile: NOISE_RELAY_SECURITY_PROFILE.to_string(), + executor_registration_id: "registration-1".to_string(), } ); } @@ -317,6 +485,9 @@ mod tests { #[tokio::test] async fn register_environment_does_not_follow_redirects_with_auth_headers() { let server = MockServer::start().await; + let executor_public_key = NoiseChannelIdentity::generate() + .expect("identity") + .public_key(); Mock::given(method("POST")) .and(path("/cloud/environment/environment-requested/register")) .and(header("authorization", "Bearer registry-token")) @@ -336,7 +507,7 @@ mod tests { .expect("client"); let error = client - .register_environment("environment-requested") + .register_environment("environment-requested", &executor_public_key) .await .expect_err("redirect response should not be followed"); @@ -364,3 +535,7 @@ mod tests { assert!(!debug.contains("workspace-123")); } } + +#[cfg(test)] +#[path = "remote/noise_tests.rs"] +mod noise_tests; diff --git a/codex-rs/exec-server/src/remote/noise_tests.rs b/codex-rs/exec-server/src/remote/noise_tests.rs new file mode 100644 index 000000000..1c4525e63 --- /dev/null +++ b/codex-rs/exec-server/src/remote/noise_tests.rs @@ -0,0 +1,163 @@ +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; +use codex_api::AuthProvider; +use codex_api::SharedAuthProvider; +use http::HeaderMap; +use http::HeaderValue; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpListener; +use tokio::time::timeout; +use tokio_tungstenite::accept_async; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +use super::*; + +const HARNESS_KEY_AUTHORIZATION: &str = "authorization-that-must-not-leak"; + +#[derive(Debug)] +struct StaticRegistryAuthProvider; + +impl AuthProvider for StaticRegistryAuthProvider { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + let _ = headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer registry-token"), + ); + } +} + +fn static_registry_auth_provider() -> SharedAuthProvider { + Arc::new(StaticRegistryAuthProvider) +} + +#[tokio::test] +async fn reconnect_reuses_registration_until_url_is_rejected() -> Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let rendezvous_url = format!("ws://{}", listener.local_addr()?); + let registry = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/cloud/environment/environment-requested/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "environment_id": "environment-requested", + "url": rendezvous_url, + "security_profile": NOISE_RELAY_SECURITY_PROFILE, + "executor_registration_id": "registration-1", + }))) + .expect(2) + .mount(®istry) + .await; + let config = RemoteEnvironmentConfig::new( + registry.uri(), + "environment-requested".to_string(), + static_registry_auth_provider(), + )?; + let environment_task = tokio::spawn(run_remote_environment( + config, + ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + )?, + )); + + let (first_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??; + let mut first_websocket = accept_async(first_socket).await?; + first_websocket.close(None).await?; + + // An ordinary disconnect retries the same URL without registering again. + let (mut rejected_socket, _peer_addr) = + timeout(Duration::from_secs(5), listener.accept()).await??; + let mut request = [0u8; 4096]; + let _ = rejected_socket.read(&mut request).await?; + rejected_socket + .write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n") + .await?; + rejected_socket.shutdown().await?; + + // The 4xx response discards the old registration before this attempt. + let (third_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??; + let _third_websocket = accept_async(third_socket).await?; + registry.verify().await; + + environment_task.abort(); + let _ = environment_task.await; + Ok(()) +} + +#[tokio::test] +async fn validate_harness_key_requires_explicit_valid_response() { + let server = MockServer::start().await; + let harness_public_key = NoiseChannelIdentity::generate() + .expect("identity") + .public_key(); + Mock::given(method("POST")) + .and(path("/cloud/environment/environment-requested/validate")) + .and(header("authorization", "Bearer registry-token")) + .and(body_partial_json(serde_json::json!({ + "executor_registration_id": "registration-1", + "harness_public_key": harness_public_key.clone(), + "harness_key_authorization": HARNESS_KEY_AUTHORIZATION, + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "valid": false, + }))) + .mount(&server) + .await; + let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) + .expect("client"); + + let error = RegistryHarnessKeyValidator { + client, + environment_id: "environment-requested".to_string(), + executor_registration_id: "registration-1".to_string(), + } + .validate_harness_key(&harness_public_key, HARNESS_KEY_AUTHORIZATION) + .await + .expect_err("a false validation response must fail closed"); + + assert!(matches!( + error, + ExecServerError::Protocol(message) + if message == "environment registry rejected Noise relay harness key" + )); +} + +#[tokio::test] +async fn validate_harness_key_does_not_expose_error_body() { + let server = MockServer::start().await; + let harness_public_key = NoiseChannelIdentity::generate() + .expect("identity") + .public_key(); + Mock::given(method("POST")) + .and(path("/cloud/environment/environment-requested/validate")) + .respond_with(ResponseTemplate::new(500).set_body_string(HARNESS_KEY_AUTHORIZATION)) + .mount(&server) + .await; + let client = EnvironmentRegistryClient::new(server.uri(), static_registry_auth_provider()) + .expect("client"); + + let error = RegistryHarnessKeyValidator { + client, + environment_id: "environment-requested".to_string(), + executor_registration_id: "registration-1".to_string(), + } + .validate_harness_key(&harness_public_key, HARNESS_KEY_AUTHORIZATION) + .await + .expect_err("validation HTTP error should fail closed"); + + let display = error.to_string(); + assert!(!display.contains(HARNESS_KEY_AUTHORIZATION)); + assert!(matches!( + error, + ExecServerError::EnvironmentRegistryHttp { message, .. } + if message == "environment registry harness key validation failed" + )); +} diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index ff41dd525..114fe122a 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/codex-rs/exec-server/tests/relay.rs @@ -4,40 +4,49 @@ mod common; mod relay_proto; use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Context; use anyhow::Result; -use anyhow::anyhow; -use anyhow::bail; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD; use codex_api::AuthProvider; -use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCRequest; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecParams; +use codex_exec_server::ExecResponse; +use codex_exec_server::ExecServerClient; +use codex_exec_server::ExecServerError; use codex_exec_server::ExecServerRuntimePaths; -use codex_exec_server::InitializeParams; -use codex_exec_server::InitializeResponse; +use codex_exec_server::FsReadFileParams; +use codex_exec_server::NoiseChannelIdentity; +use codex_exec_server::NoiseChannelPublicKey; +use codex_exec_server::NoiseRendezvousConnectArgs; +use codex_exec_server::NoiseRendezvousConnectBundle; +use codex_exec_server::NoiseRendezvousConnectProvider; +use codex_exec_server::ProcessId; use codex_exec_server::RemoteEnvironmentConfig; +use codex_utils_path_uri::PathUri; +use futures::FutureExt; use futures::SinkExt; use futures::StreamExt; +use futures::future::BoxFuture; use http::HeaderMap; use http::HeaderValue; use pretty_assertions::assert_eq; use prost::Message as ProstMessage; -use relay_proto::RelayData; use relay_proto::RelayMessageFrame; -use relay_proto::RelayReset; use relay_proto::relay_message_frame; -use std::sync::Arc; +use tempfile::TempDir; use tokio::net::TcpListener; +use tokio::net::TcpStream; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; use tokio_tungstenite::tungstenite::Message; -use uuid::Uuid; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; @@ -45,10 +54,11 @@ use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; -const ENVIRONMENT_ID: &str = "env-mux-test"; +const ENVIRONMENT_ID: &str = "env-noise-relay-test"; +const EXECUTOR_REGISTRATION_ID: &str = "registration-1"; +const HARNESS_KEY_AUTHORIZATION: &str = "harness-key-authorization"; const REGISTRY_TOKEN: &str = "registry-token"; -const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; -const TEST_TIMEOUT: Duration = Duration::from_secs(5); +const TEST_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Debug)] struct StaticRegistryAuthProvider; @@ -62,12 +72,70 @@ impl AuthProvider for StaticRegistryAuthProvider { } } +struct FailingNoiseConnectProvider { + attempts: Arc, +} + +impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider { + fn connect_bundle( + &self, + _: NoiseChannelPublicKey, + ) -> BoxFuture<'_, Result> { + self.attempts.fetch_add(1, Ordering::SeqCst); + async { + Err(ExecServerError::Protocol( + "test registry connect failure".to_string(), + )) + } + .boxed() + } +} + fn static_registry_auth_provider() -> codex_api::SharedAuthProvider { Arc::new(StaticRegistryAuthProvider) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn multiplexed_remote_environment_routes_independent_virtual_streams() -> Result<()> { +#[tokio::test] +async fn noise_environment_refreshes_bundle_for_each_connection_attempt() -> Result<()> { + let attempts = Arc::new(AtomicUsize::new(0)); + let manager = EnvironmentManager::without_environments(); + manager.upsert_noise_environment( + ENVIRONMENT_ID.to_string(), + Arc::new(FailingNoiseConnectProvider { + attempts: Arc::clone(&attempts), + }), + )?; + let backend = manager + .get_environment(ENVIRONMENT_ID) + .context("Noise environment should be materialized")? + .get_exec_backend(); + + for attempt in 1..=2 { + let result = backend + .start(ExecParams { + process_id: ProcessId::new(format!("proc-{attempt}")), + argv: vec!["true".to_string()], + cwd: PathUri::from_path(std::env::current_dir()?)?, + env_policy: None, + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + }) + .await; + assert!(matches!( + result, + Err(ExecServerError::Protocol(ref message)) + if message == "test registry connect failure" + )); + } + + assert_eq!(attempts.load(Ordering::SeqCst), 2); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; let rendezvous_url = format!("ws://{}", listener.local_addr()?); let registry = MockServer::start().await; @@ -78,7 +146,19 @@ async fn multiplexed_remote_environment_routes_independent_virtual_streams() -> .and(header("authorization", format!("Bearer {REGISTRY_TOKEN}"))) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ "environment_id": ENVIRONMENT_ID, - "url": rendezvous_url, + "url": format!("{rendezvous_url}/relay?role=environment"), + "security_profile": "noise_hybrid_ik_v1", + "executor_registration_id": EXECUTOR_REGISTRATION_ID, + }))) + .mount(®istry) + .await; + Mock::given(method("POST")) + .and(path(format!( + "/cloud/environment/{ENVIRONMENT_ID}/validate" + ))) + .and(header("authorization", format!("Bearer {REGISTRY_TOKEN}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "valid": true, }))) .mount(®istry) .await; @@ -95,281 +175,163 @@ async fn multiplexed_remote_environment_routes_independent_virtual_streams() -> runtime_paths, )); - let (socket, _peer_addr) = timeout(TEST_TIMEOUT, listener.accept()) + let environment_websocket = accept_websocket(&listener, "environment").await?; + let executor_public_key = registered_executor_public_key(®istry).await?; + let harness_identity = NoiseChannelIdentity::generate()?; + let client_args = NoiseRendezvousConnectArgs { + bundle: NoiseRendezvousConnectBundle { + websocket_url: format!("{rendezvous_url}/relay?role=harness"), + environment_id: ENVIRONMENT_ID.to_string(), + executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(), + executor_public_key, + harness_key_authorization: HARNESS_KEY_AUTHORIZATION.to_string(), + }, + harness_identity, + client_name: "noise-relay-test".to_string(), + connect_timeout: TEST_TIMEOUT, + initialize_timeout: TEST_TIMEOUT, + resume_session_id: None, + }; + let client_task = + 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 relay_task = tokio::spawn(proxy_relay_frames( + environment_websocket, + harness_websocket, + Arc::clone(&captured_frames), + )); + let client = timeout(TEST_TIMEOUT, client_task) .await - .context("remote environment should connect to fake rendezvous")??; - let mut websocket = timeout(TEST_TIMEOUT, accept_async(socket)) - .await - .context("fake rendezvous should accept environment websocket")??; + .context("Noise harness client should connect")???; - let stream_a = "stream-a"; - let stream_b = "stream-b"; - send_relay_message( - &mut websocket, - stream_a, - /*seq*/ 0, - initialize_request(/*id*/ 1, "relay-test-a")?, - ) - .await?; - send_relay_message( - &mut websocket, - stream_b, - /*seq*/ 0, - initialize_request(/*id*/ 1, "relay-test-b")?, - ) - .await?; + let response = client + .exec(ExecParams { + process_id: ProcessId::from("proc-1"), + argv: vec!["true".to_string()], + cwd: PathUri::from_path(std::env::current_dir()?)?, + env_policy: None, + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + }) + .await?; + assert_eq!( + response, + ExecResponse { + process_id: ProcessId::from("proc-1") + } + ); - let initialize_responses = read_relay_messages_by_stream(&mut websocket, /*count*/ 2).await?; - let session_a = - assert_initialize_response(initialize_responses.get(stream_a), stream_a, /*id*/ 1)?; - let session_b = - assert_initialize_response(initialize_responses.get(stream_b), stream_b, /*id*/ 1)?; - assert_ne!(session_a, session_b); + let temp_dir = TempDir::new()?; + let large_file_path = temp_dir.path().join("large-response.bin"); + let large_file_contents = vec![0x5a; 128 * 1024]; + std::fs::write(&large_file_path, &large_file_contents)?; + let read_response = client + .fs_read_file(FsReadFileParams { + path: PathUri::from_path(large_file_path)?, + sandbox: None, + }) + .await?; + assert_eq!( + STANDARD.decode(read_response.data_base64)?, + large_file_contents + ); - send_relay_message( - &mut websocket, - stream_a, - /*seq*/ 1, - notification("initialized", serde_json::json!({})), - ) - .await?; - send_relay_message( - &mut websocket, - stream_b, - /*seq*/ 1, - notification("initialized", serde_json::json!({})), - ) - .await?; + assert_relay_data_is_encrypted(&captured_frames)?; - send_relay_message( - &mut websocket, - stream_a, - /*seq*/ 2, - request(/*id*/ 2, "test/unknown-a", serde_json::json!({})), - ) - .await?; - send_relay_message( - &mut websocket, - stream_b, - /*seq*/ 2, - request(/*id*/ 2, "test/unknown-b", serde_json::json!({})), - ) - .await?; - - let unknown_method_responses = - read_relay_messages_by_stream(&mut websocket, /*count*/ 2).await?; - assert_error_response( - unknown_method_responses.get(stream_a), - stream_a, - /*id*/ 2, - "test/unknown-a", - )?; - assert_error_response( - unknown_method_responses.get(stream_b), - stream_b, - /*id*/ 2, - "test/unknown-b", - )?; - - send_relay_reset(&mut websocket, stream_a, "test_reset").await?; - send_relay_message( - &mut websocket, - stream_b, - /*seq*/ 3, - request( - /*id*/ 3, - "test/unknown-b-after-reset", - serde_json::json!({}), - ), - ) - .await?; - - let (stream_id, message) = read_relay_message(&mut websocket).await?; - assert_eq!(stream_id, stream_b); - assert_error_response( - Some(&message), - stream_b, - /*id*/ 3, - "test/unknown-b-after-reset", - )?; - - websocket.close(None).await?; + drop(client); + relay_task.abort(); remote_environment.abort(); + let _ = relay_task.await; let _ = remote_environment.await; Ok(()) } -async fn send_relay_message( - websocket: &mut WebSocketStream, - stream_id: &str, - seq: u32, - message: JSONRPCMessage, -) -> Result<()> -where - S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, -{ - let payload = serde_json::to_vec(&message)?; - let frame = RelayMessageFrame { - version: RELAY_MESSAGE_FRAME_VERSION, - stream_id: stream_id.to_string(), - ack: 0, - ack_bits: 0, - body: Some(relay_message_frame::Body::Data(RelayData { - seq, - segment_index: 0, - segment_count: 1, - payload, - })), - }; - send_relay_frame(websocket, frame).await +async fn accept_websocket( + listener: &TcpListener, + role: &str, +) -> Result> { + let (socket, _peer_addr) = timeout(TEST_TIMEOUT, listener.accept()) + .await + .with_context(|| format!("remote {role} should connect to fake rendezvous"))??; + timeout(TEST_TIMEOUT, accept_async(socket)) + .await + .with_context(|| format!("fake rendezvous should accept {role} websocket"))? + .map_err(Into::into) } -async fn send_relay_reset( - websocket: &mut WebSocketStream, - stream_id: &str, - reason: &str, -) -> Result<()> -where - S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, -{ - send_relay_frame( - websocket, - RelayMessageFrame { - version: RELAY_MESSAGE_FRAME_VERSION, - stream_id: stream_id.to_string(), - ack: 0, - ack_bits: 0, - body: Some(relay_message_frame::Body::Reset(RelayReset { - reason: reason.to_string(), - })), - }, - ) - .await +async fn registered_executor_public_key(registry: &MockServer) -> Result { + let requests = registry + .received_requests() + .await + .context("wiremock should retain requests")?; + let request = requests + .iter() + .find(|request| request.url.path().ends_with("/register")) + .context("exec-server should register before connecting")?; + let body: serde_json::Value = serde_json::from_slice(&request.body)?; + let key = serde_json::from_value(body["executor_public_key"].clone())?; + Ok(key) } -async fn send_relay_frame( - websocket: &mut WebSocketStream, - frame: RelayMessageFrame, -) -> Result<()> -where - S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, -{ - websocket - .send(Message::Binary(frame.encode_to_vec().into())) - .await?; +async fn proxy_relay_frames( + mut environment: WebSocketStream, + mut harness: WebSocketStream, + captured_frames: Arc>>>, +) -> Result<()> { + loop { + tokio::select! { + message = environment.next() => { + let Some(message) = message else { + break; + }; + let message = message?; + capture_binary_frame(&captured_frames, &message); + harness.send(message).await?; + } + message = harness.next() => { + let Some(message) = message else { + break; + }; + let message = message?; + capture_binary_frame(&captured_frames, &message); + environment.send(message).await?; + } + } + } Ok(()) } -async fn read_relay_messages_by_stream( - websocket: &mut WebSocketStream, - count: usize, -) -> Result> -where - S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, -{ - let mut messages = HashMap::new(); - for _ in 0..count { - let (stream_id, message) = read_relay_message(websocket).await?; - if messages.insert(stream_id.clone(), message).is_some() { - bail!("received duplicate response for stream {stream_id}"); - } - } - Ok(messages) -} - -async fn read_relay_message( - websocket: &mut WebSocketStream, -) -> Result<(String, JSONRPCMessage)> -where - S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, -{ - loop { - let frame = timeout(TEST_TIMEOUT, websocket.next()) - .await - .context("timed out waiting for relay frame")? - .ok_or_else(|| anyhow!("environment websocket closed"))??; - match frame { - Message::Binary(bytes) => { - let frame = RelayMessageFrame::decode(bytes.as_ref())?; - let stream_id = frame.stream_id; - let Some(relay_message_frame::Body::Data(data)) = frame.body else { - continue; - }; - let message = serde_json::from_slice(&data.payload)?; - return Ok((stream_id, message)); - } - Message::Ping(_) | Message::Pong(_) => {} - Message::Close(_) => bail!("environment websocket closed"), - Message::Text(_) => bail!("environment sent text frame on relay websocket"), - Message::Frame(_) => {} - } +fn capture_binary_frame(captured_frames: &Mutex>>, message: &Message) { + if let Message::Binary(bytes) = message { + captured_frames + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(bytes.to_vec()); } } -fn initialize_request(id: i64, client_name: &str) -> Result { - Ok(request( - id, - "initialize", - serde_json::to_value(InitializeParams { - client_name: client_name.to_string(), - resume_session_id: None, - })?, - )) -} - -fn request(id: i64, method: &str, params: serde_json::Value) -> JSONRPCMessage { - JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(id), - method: method.to_string(), - params: Some(params), - trace: None, - }) -} - -fn notification(method: &str, params: serde_json::Value) -> JSONRPCMessage { - JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - }) -} - -fn assert_initialize_response( - message: Option<&JSONRPCMessage>, - stream_id: &str, - id: i64, -) -> Result { - let message = message.ok_or_else(|| anyhow!("missing initialize response for {stream_id}"))?; - let JSONRPCMessage::Response(JSONRPCResponse { - id: response_id, - result, - }) = message - else { - bail!("expected initialize response for {stream_id}, got {message:?}"); - }; - assert_eq!(response_id, &RequestId::Integer(id)); - let response: InitializeResponse = serde_json::from_value(result.clone())?; - Ok(Uuid::parse_str(&response.session_id)?) -} - -fn assert_error_response( - message: Option<&JSONRPCMessage>, - stream_id: &str, - id: i64, - expected_method: &str, -) -> Result<()> { - let message = message.ok_or_else(|| anyhow!("missing error response for {stream_id}"))?; - let JSONRPCMessage::Error(JSONRPCError { - id: response_id, - error, - }) = message - else { - bail!("expected error response for {stream_id}, got {message:?}"); - }; - assert_eq!(response_id, &RequestId::Integer(id)); +fn assert_relay_data_is_encrypted(captured_frames: &Mutex>>) -> Result<()> { + let captured_frames = captured_frames + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut data_frames = 0; + for encoded in captured_frames.iter() { + let frame = RelayMessageFrame::decode(encoded.as_slice())?; + let Some(relay_message_frame::Body::Data(data)) = frame.body else { + continue; + }; + data_frames += 1; + let payload = String::from_utf8_lossy(&data.payload); + assert!(!payload.contains("initialize")); + assert!(!payload.contains("process/start")); + assert!(!payload.contains("noise-relay-test")); + } assert!( - error.message.contains(expected_method), - "expected error for {stream_id} to mention {expected_method}, got {}", - error.message + data_frames >= 4, + "expected encrypted request and response frames" ); Ok(()) }