From 3ebabf7eef5e02b40ac729f3f7960f5a0d282f66 Mon Sep 17 00:00:00 2001 From: viyatb-oai Date: Wed, 3 Jun 2026 12:23:01 -0700 Subject: [PATCH] feat: add Noise relay handshake frames Co-authored-by: Codex noreply@openai.com --- codex-rs/exec-server/src/connection.rs | 11 +- .../proto/codex.exec_server.relay.v1.proto | 5 + .../src/proto/codex.exec_server.relay.v1.rs | 9 +- codex-rs/exec-server/src/relay.rs | 147 +++++++++++++++--- codex-rs/exec-server/src/relay_proto.rs | 2 + 5 files changed, 149 insertions(+), 25 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index b211c504ac..1cace9df4e 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -43,8 +43,11 @@ pub(crate) enum JsonRpcConnectionEvent { } #[derive(Clone)] +/// Describes resources owned by a JSON-RPC connection, not the protection of +/// its byte stream. `External` covers websocket, relay, and Noise connections +/// whose transport task owns no child process for Codex to terminate. pub(crate) enum JsonRpcTransport { - Plain, + External, Stdio { transport: StdioTransport }, } @@ -57,7 +60,7 @@ impl JsonRpcTransport { pub(crate) fn terminate(&self) { match self { - Self::Plain => {} + Self::External => {} Self::Stdio { transport } => transport.terminate(), } } @@ -315,7 +318,7 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles: vec![reader_task, writer_task], - transport: JsonRpcTransport::Plain, + transport: JsonRpcTransport::External, } } @@ -452,7 +455,7 @@ impl JsonRpcConnection { incoming_rx, disconnected_rx, task_handles: vec![websocket_task], - transport: JsonRpcTransport::Plain, + transport: JsonRpcTransport::External, } } diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto index 46527d80cc..3de1bfbe99 100644 --- a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto +++ b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.proto @@ -14,6 +14,7 @@ message RelayMessageFrame { RelayResume resume = 7; RelayReset reset = 8; RelayHeartbeat heartbeat = 9; + RelayHandshake handshake = 10; } } @@ -35,3 +36,7 @@ message RelayReset { } message RelayHeartbeat {} + +message RelayHandshake { + bytes payload = 1; +} diff --git a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs index 072a003dac..f65e1329f5 100644 --- a/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs +++ b/codex-rs/exec-server/src/proto/codex.exec_server.relay.v1.rs @@ -9,7 +9,7 @@ pub struct RelayMessageFrame { pub ack: u32, #[prost(uint32, tag = "4")] pub ack_bits: u32, - #[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9")] + #[prost(oneof = "relay_message_frame::Body", tags = "5, 6, 7, 8, 9, 10")] pub body: ::core::option::Option, } pub mod relay_message_frame { @@ -25,6 +25,8 @@ pub mod relay_message_frame { Reset(super::RelayReset), #[prost(message, tag = "9")] Heartbeat(super::RelayHeartbeat), + #[prost(message, tag = "10")] + Handshake(super::RelayHandshake), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -52,3 +54,8 @@ pub struct RelayReset { } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct RelayHeartbeat {} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct RelayHandshake { + #[prost(bytes = "vec", tag = "1")] + pub payload: ::prost::alloc::vec::Vec, +} diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index ae07d7652c..04c5e5061e 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -22,24 +22,29 @@ use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; 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_RELAY_RESET_REASON_BYTES: usize = 256; +const MAX_RELAY_STREAM_ID_BYTES: usize = 128; #[derive(Debug, Clone, Copy, Eq, PartialEq)] -enum RelayFrameBodyKind { +pub(crate) enum RelayFrameBodyKind { Data, Ack, Resume, Reset, Heartbeat, + Handshake, } impl RelayMessageFrame { - fn data(stream_id: String, seq: u32, payload: Vec) -> Self { + pub(crate) fn data(stream_id: String, seq: u32, payload: Vec) -> Self { Self { version: RELAY_MESSAGE_FRAME_VERSION, stream_id, @@ -54,7 +59,7 @@ impl RelayMessageFrame { } } - fn resume(stream_id: String) -> Self { + pub(crate) fn resume(stream_id: String) -> Self { Self { version: RELAY_MESSAGE_FRAME_VERSION, stream_id, @@ -66,18 +71,55 @@ impl RelayMessageFrame { } } - fn validate(&self) -> Result { + pub(crate) fn handshake(stream_id: String, payload: Vec) -> Self { + Self { + version: RELAY_MESSAGE_FRAME_VERSION, + stream_id, + ack: 0, + ack_bits: 0, + body: Some(relay_message_frame::Body::Handshake(RelayHandshake { + payload, + })), + } + } + + 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 })), + } + } + + /// Validate cleartext routing metadata before a frame reaches either the + /// direct JSON-RPC parser or the Noise relay state machine. + /// + /// The encrypted path intentionally leaves this metadata visible to + /// rendezvous, so it must be canonical and tightly bounded everywhere. + pub(crate) fn validate(&self) -> Result { if self.version != RELAY_MESSAGE_FRAME_VERSION { return Err(ExecServerError::Protocol(format!( "unsupported relay message frame version {}", self.version ))); } - if self.stream_id.trim().is_empty() { + if self.stream_id.is_empty() { return Err(ExecServerError::Protocol( "relay message frame is missing stream_id".to_string(), )); } + if self.stream_id.trim() != self.stream_id { + return Err(ExecServerError::Protocol( + "relay message frame stream_id is not canonical".to_string(), + )); + } + if self.stream_id.len() > MAX_RELAY_STREAM_ID_BYTES { + return Err(ExecServerError::Protocol( + "relay message frame stream_id is too long".to_string(), + )); + } match self.body.as_ref() { Some(relay_message_frame::Body::Data(data)) => { if data.segment_index != 0 || data.segment_count != 1 || data.payload.is_empty() { @@ -90,35 +132,64 @@ impl RelayMessageFrame { Some(relay_message_frame::Body::AckFrame(_)) => Ok(RelayFrameBodyKind::Ack), Some(relay_message_frame::Body::Resume(_)) => Ok(RelayFrameBodyKind::Resume), Some(relay_message_frame::Body::Reset(reset)) => { - if reset.reason.is_empty() { + if reset.reason.is_empty() || reset.reason.len() > MAX_RELAY_RESET_REASON_BYTES { return Err(ExecServerError::Protocol( - "relay reset message frame is missing reason".to_string(), + "relay reset message frame has invalid reason".to_string(), )); } Ok(RelayFrameBodyKind::Reset) } Some(relay_message_frame::Body::Heartbeat(_)) => Ok(RelayFrameBodyKind::Heartbeat), + Some(relay_message_frame::Body::Handshake(handshake)) => { + if handshake.payload.is_empty() { + return Err(ExecServerError::Protocol( + "relay handshake message frame is missing payload".to_string(), + )); + } + Ok(RelayFrameBodyKind::Handshake) + } None => Err(ExecServerError::Protocol( "relay message frame is missing body".to_string(), )), } } - fn into_jsonrpc_message(self) -> Result { + pub(crate) fn into_data(self) -> Result { let kind = self.validate()?; if kind != RelayFrameBodyKind::Data { return Err(ExecServerError::Protocol( "expected relay data message frame".to_string(), )); } - let payload = match self.body { - Some(relay_message_frame::Body::Data(data)) => data.payload, - _ => Vec::new(), - }; + match self.body { + Some(relay_message_frame::Body::Data(data)) => Ok(data), + _ => Err(ExecServerError::Protocol( + "expected relay data message frame".to_string(), + )), + } + } + + fn into_jsonrpc_message(self) -> Result { + let payload = self.into_data()?.payload; serde_json::from_slice(&payload).map_err(ExecServerError::Json) } - fn into_reset_reason(self) -> Option { + pub(crate) fn into_handshake_payload(self) -> Result, ExecServerError> { + let kind = self.validate()?; + if kind != RelayFrameBodyKind::Handshake { + return Err(ExecServerError::Protocol( + "expected relay handshake message frame".to_string(), + )); + } + match self.body { + Some(relay_message_frame::Body::Handshake(handshake)) => Ok(handshake.payload), + _ => Err(ExecServerError::Protocol( + "expected relay handshake message frame".to_string(), + )), + } + } + + pub(crate) fn into_reset_reason(self) -> Option { match self.body { Some(relay_message_frame::Body::Reset(reset)) if !reset.reason.is_empty() => { Some(reset.reason) @@ -128,16 +199,18 @@ impl RelayMessageFrame { } } -fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec { +pub(crate) fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec { frame.encode_to_vec() } -fn decode_relay_message_frame(payload: &[u8]) -> Result { +pub(crate) fn decode_relay_message_frame( + payload: &[u8], +) -> Result { RelayMessageFrame::decode(payload) .map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}"))) } -fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { +pub(crate) fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> { serde_json::to_vec(message).map_err(ExecServerError::Json) } @@ -253,7 +326,8 @@ where } RelayFrameBodyKind::Ack | RelayFrameBodyKind::Resume - | RelayFrameBodyKind::Heartbeat => {} + | RelayFrameBodyKind::Heartbeat + | RelayFrameBodyKind::Handshake => {} } } Some(Ok(Message::Close(_))) | None => { @@ -294,10 +368,22 @@ where incoming_rx, disconnected_rx, task_handles: vec![websocket_task], - transport: JsonRpcTransport::Plain, + transport: JsonRpcTransport::External, } } +/// Runs the backwards-compatible, cleartext multiplexed relay protocol. +/// +/// The physical websocket carries protobuf [`RelayMessageFrame`] values. Each +/// frame's `stream_id` selects an independent logical JSON-RPC connection, so a +/// single registered exec-server can serve multiple orchestrator sessions. +/// This runner intentionally preserves the protocol used before Noise relay +/// support was introduced. Callers must select the separate secure runner when +/// Noise protection was explicitly negotiated during registration. +/// +/// Relay framing is not an authentication boundary. Every frame is validated +/// before its routing metadata or payload is used, and malformed frames are +/// dropped without affecting the other virtual streams on the connection. pub(crate) async fn run_multiplexed_environment( stream: WebSocketStream, processor: ConnectionProcessor, @@ -310,6 +396,9 @@ pub(crate) async fn run_multiplexed_environment( let mut streams: HashMap = HashMap::new(); loop { + // Serialize all logical-stream writes through this task so only one + // owner writes to the physical websocket. Reads remain in the same + // select loop, which also keeps disconnect handling deterministic. let frame = tokio::select! { maybe_encoded = physical_outgoing_rx.recv() => { let Some(encoded) = maybe_encoded else { @@ -361,6 +450,11 @@ pub(crate) async fn run_multiplexed_environment( continue; } }; + + // A logical connection is created lazily on its first data + // frame. The connection processor owns the JSON-RPC lifecycle; + // this task only translates between relay and connection + // events. let stream = streams.entry(stream_id.clone()).or_insert_with(|| { spawn_virtual_stream( stream_id.clone(), @@ -382,24 +476,34 @@ pub(crate) async fn run_multiplexed_environment( stream.disconnect(frame.into_reset_reason()).await; } } + // These control frames are meaningful to other relay protocol + // variants. The legacy protocol has no resume or handshake state, + // so ignoring them preserves its existing wire behavior. RelayFrameBodyKind::Ack | RelayFrameBodyKind::Resume - | RelayFrameBodyKind::Heartbeat => {} + | RelayFrameBodyKind::Heartbeat + | RelayFrameBodyKind::Handshake => {} } } + // A physical disconnect ends every virtual connection. Notify each + // processor explicitly so requests do not remain live after the relay + // websocket has disappeared. for (_stream_id, stream) in streams { stream.disconnect(/*reason*/ None).await; } drop(physical_outgoing_tx); } +/// The exec-server-facing half of one logical connection on the shared relay. struct VirtualStream { incoming_tx: mpsc::Sender, disconnected_tx: watch::Sender, } impl VirtualStream { + /// Marks the logical connection disconnected and supplies the peer's reset + /// reason, when one was provided. async fn disconnect(self, reason: Option) { let _ = self.disconnected_tx.send(true); let _ = self @@ -409,6 +513,8 @@ impl VirtualStream { } } +/// Creates a logical JSON-RPC connection and forwards its outgoing messages +/// back to the physical relay writer as legacy cleartext data frames. fn spawn_virtual_stream( stream_id: String, processor: ConnectionProcessor, @@ -446,7 +552,7 @@ fn spawn_virtual_stream( incoming_rx, disconnected_rx, task_handles: vec![writer_task], - transport: JsonRpcTransport::Plain, + transport: JsonRpcTransport::External, }; tokio::spawn(async move { processor.run_connection(connection).await; @@ -476,6 +582,7 @@ mod tests { use futures::task::AtomicWaker; use tokio::net::TcpListener; use tokio::time::timeout; + use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message; diff --git a/codex-rs/exec-server/src/relay_proto.rs b/codex-rs/exec-server/src/relay_proto.rs index b8a938b8c7..da7cb5296d 100644 --- a/codex-rs/exec-server/src/relay_proto.rs +++ b/codex-rs/exec-server/src/relay_proto.rs @@ -2,6 +2,8 @@ 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;