openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-07 15:40:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: harden Noise executor relay

Browse Source 
Co-authored-by: Codex noreply@openai.com
This commit is contained in:
viyatb-oai
2026-06-03 17:18:00 -07:00
parent e7554548e0
commit 8ae85753fc
7 changed files with 574 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

265
codex-rs/exec-server/src/noise_relay/environment.rs
View File

@@ -1,6 +1,4 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use futures::SinkExt;
@@ -8,7 +6,6 @@ use futures::StreamExt;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
@@ -18,29 +15,20 @@ use tracing::warn;
use crate::ExecServerError;
use crate::connection::CHANNEL_CAPACITY;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_channel::NoiseChannelPublicKey;
use crate::noise_channel::NoiseTransport;
use crate::noise_channel::PendingResponderHandshake;
use crate::noise_channel::noise_channel_prologue;
use crate::noise_relay::message_framing::JsonRpcMessageDecoder;
use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN;
use crate::noise_relay::message_framing::frame_jsonrpc_message;
use crate::noise_relay::ordered_ciphertext::OrderedCiphertextFrames;
use crate::noise_relay::take_next_sequence;
use crate::noise_relay::NOISE_RELAY_RESET_REASON;
use crate::noise_relay::executor_stream::ClosedNoiseVirtualStream;
use crate::noise_relay::executor_stream::NoiseVirtualStream;
use crate::noise_relay::executor_stream::spawn_noise_virtual_stream;
use crate::relay::RelayFrameBodyKind;
use crate::relay::decode_relay_message_frame;
use crate::relay::encode_relay_message_frame;
use crate::relay_proto::RelayData;
use crate::relay_proto::RelayMessageFrame;
use crate::server::ConnectionProcessor;
// This value is already part of the relay wire contract. Keep it stable even
// though the source module now uses the more precise Noise terminology.
const NOISE_RELAY_RESET_REASON: &str = "secure_relay_protocol_error";
const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128;
const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096;
const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32;
@@ -76,9 +64,23 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
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::<Vec<u8>>(CHANNEL_CAPACITY);
let (closed_stream_tx, mut closed_stream_rx) =
mpsc::channel::<ClosedNoiseVirtualStream>(MAX_ACTIVE_NOISE_RELAY_STREAMS);
// A separate writer task is required because this state machine also
// produces resets and handshake responses. If the same task both sent into
// and drained the bounded outgoing channel, backpressure could make it wait
// on itself and stop servicing the physical websocket.
let mut physical_writer_task = tokio::spawn(async move {
while let Some(encoded) = physical_outgoing_rx.recv().await {
if let Err(error) = websocket_sink.send(Message::Binary(encoded.into())).await {
debug!("Noise multiplexed environment websocket write failed: {error}");
break;
}
}
});
let mut streams: HashMap<String, NoiseVirtualStream> = HashMap::new();
let mut pending_handshakes: HashMap<String, PendingHandshake> = HashMap::new();
let mut validation_tasks: JoinSet<HarnessKeyValidationResult> = JoinSet::new();
@@ -89,12 +91,21 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
// malicious authorization request must not block existing streams or
// prevent other handshakes from being received and bounded.
let frame = tokio::select! {
maybe_encoded = physical_outgoing_rx.recv() => {
let Some(encoded) = maybe_encoded else {
break;
};
if websocket.send(Message::Binary(encoded.into())).await.is_err() {
break;
writer_result = &mut physical_writer_task => {
if let Err(error) = writer_result {
warn!("Noise multiplexed environment websocket writer failed: {error}");
}
break;
}
Some(closed_stream) = closed_stream_rx.recv() => {
// A writer can finish after its peer resets and reuses the same
// routing ID. Remove only the exact authenticated stream
// instance that produced this close notification.
let is_current = streams
.get(&closed_stream.stream_id)
.is_some_and(|stream| stream.is_instance(closed_stream.instance_id));
if is_current {
streams.remove(&closed_stream.stream_id);
}
continue;
}
@@ -117,14 +128,18 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
else {
continue;
};
if let Err(error) = validation_result.result {
warn!("Noise relay harness key validation failed: {error}");
send_reset(&physical_outgoing_tx, validation_result.stream_id).await;
if validation_result.result.is_err() {
// Validators receive the short-lived authorization.
// Keep their error text out of logs even though the
// registry implementation below also sanitizes
// response bodies.
warn!("Noise relay harness key validation failed");
send_reset(&physical_outgoing_tx, validation_result.stream_id);
continue;
}
if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS {
warn!("Noise relay has too many active streams");
send_reset(&physical_outgoing_tx, validation_result.stream_id).await;
send_reset(&physical_outgoing_tx, validation_result.stream_id);
continue;
}
@@ -135,7 +150,7 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
Ok(completed) => completed,
Err(error) => {
warn!("failed to complete Noise relay handshake: {error}");
send_reset(&physical_outgoing_tx, validation_result.stream_id).await;
send_reset(&physical_outgoing_tx, validation_result.stream_id);
continue;
}
};
@@ -143,9 +158,13 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
validation_result.stream_id.clone(),
response,
);
// The shared state machine must never wait behind an
// overloaded writer queue. If a successful handshake
// response cannot be queued immediately, close this
// physical connection rather than expose a half-open
// virtual stream.
if physical_outgoing_tx
.send(encode_relay_message_frame(&response))
.await
.try_send(encode_relay_message_frame(&response))
.is_err()
{
break;
@@ -154,8 +173,10 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
validation_result.stream_id.clone(),
spawn_noise_virtual_stream(
validation_result.stream_id,
validation_result.validation_id,
processor.clone(),
physical_outgoing_tx.clone(),
closed_stream_tx.clone(),
transport,
),
);
@@ -165,14 +186,14 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
let stream_ids = pending_handshakes.keys().cloned().collect::<Vec<_>>();
pending_handshakes.clear();
for stream_id in stream_ids {
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
}
}
None => {}
}
continue;
}
incoming_message = websocket.next() => match incoming_message {
incoming_message = websocket_stream.next() => match incoming_message {
Some(Ok(Message::Binary(payload))) => match decode_relay_message_frame(payload.as_ref()) {
Ok(frame) => frame,
Err(error) => {
@@ -206,17 +227,17 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
// Bound all pre-authentication state before doing expensive
// hybrid cryptography or starting an external validation.
if streams.contains_key(&stream_id) || pending_handshakes.contains_key(&stream_id) {
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS {
warn!("Noise relay has too many active streams");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
if validation_tasks.len() >= MAX_PENDING_HANDSHAKE_VALIDATIONS {
warn!("Noise relay has too many pending harness key validations");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
let prologue = match noise_channel_prologue(
@@ -227,7 +248,7 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
Ok(prologue) => prologue,
Err(error) => {
warn!("failed to build Noise relay prologue: {error}");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
};
@@ -235,7 +256,7 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
Ok(request) => request,
Err(error) => {
warn!("failed to read Noise relay handshake frame: {error}");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
};
@@ -244,7 +265,7 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
Ok(pending) => pending,
Err(error) => {
warn!("failed to read Noise relay handshake request: {error}");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
};
@@ -260,12 +281,12 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
}
Ok(_) => {
warn!("Noise relay handshake authorization is too long");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
Err(_) => {
warn!("Noise relay handshake authorization is not UTF-8");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
};
@@ -273,7 +294,7 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
let validation_id = next_validation_id;
let Some(next_id) = next_validation_id.checked_add(1) else {
warn!("Noise relay harness key validation id exhausted");
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
};
next_validation_id = next_id;
@@ -310,11 +331,11 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
}
RelayFrameBodyKind::Data => {
// Data before handshake completion is always invalid. Removing
// a pending handshake also ensures a peer cannot keep its
// authorization task alive while sending application records.
// pending state makes the time-bounded validation result stale,
// so it can never complete a stream after this protocol error.
let Some(stream) = streams.get_mut(&stream_id) else {
pending_handshakes.remove(&stream_id);
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
};
let data = match frame.into_data() {
@@ -322,20 +343,23 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
Err(error) => {
warn!("dropping malformed Noise relay data frame: {error}");
streams.remove(&stream_id);
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
};
if let Err(error) = stream.receive_data(data).await {
if let Err(error) = stream.receive_data(data) {
warn!("failed to process Noise relay payload: {error}");
streams.remove(&stream_id);
send_reset(&physical_outgoing_tx, stream_id).await;
send_reset(&physical_outgoing_tx, stream_id);
}
}
RelayFrameBodyKind::Reset => {
pending_handshakes.remove(&stream_id);
if let Some(stream) = streams.remove(&stream_id) {
stream.disconnect(frame.into_reset_reason()).await;
// Reset is cleartext relay control and is not authenticated
// by Noise. Honor its availability effect, but never
// forward attacker-controlled reason text into logs.
stream.disconnect(/*reason*/ None);
}
}
RelayFrameBodyKind::Ack
@@ -345,7 +369,14 @@ pub(crate) async fn run_noise_multiplexed_environment<S, V>(
}
for (_stream_id, stream) in streams {
stream.disconnect(/*reason*/ None).await;
stream.disconnect(/*reason*/ None);
}
// Dropping the JoinSet below aborts any still-running registry validations.
// Await an abort only when the select loop did not already consume the
// writer result.
if !physical_writer_task.is_finished() {
physical_writer_task.abort();
let _ = physical_writer_task.await;
}
}
@@ -360,138 +391,10 @@ struct HarnessKeyValidationResult {
result: Result<(), ExecServerError>,
}
struct NoiseVirtualStream {
incoming_tx: mpsc::Sender<JsonRpcConnectionEvent>,
disconnected_tx: watch::Sender<bool>,
transport: Arc<Mutex<NoiseTransport>>,
inbound_ciphertexts: OrderedCiphertextFrames,
inbound_decoder: JsonRpcMessageDecoder,
}
impl NoiseVirtualStream {
async fn disconnect(self, reason: Option<String>) {
let _ = self.disconnected_tx.send(true);
let _ = self
.incoming_tx
.send(JsonRpcConnectionEvent::Disconnected { reason })
.await;
}
async fn receive_data(&mut self, data: RelayData) -> Result<(), ExecServerError> {
// Relay sequence ordering is enforced before taking the transport lock
// and decrypting. Each virtual stream owns one ordered Noise nonce
// space shared by its reader and writer transport halves.
for ciphertext in self.inbound_ciphertexts.push(data.seq, data.payload)? {
let plaintext = {
let mut transport = self
.transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.decrypt(&ciphertext).map_err(|error| {
ExecServerError::Protocol(format!("Noise relay decryption failed: {error}"))
})?
};
for message in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
.send(JsonRpcConnectionEvent::Message(message))
.await
.map_err(|_| ExecServerError::Closed)?;
}
}
Ok(())
}
}
fn spawn_noise_virtual_stream(
stream_id: String,
processor: ConnectionProcessor,
physical_outgoing_tx: mpsc::Sender<Vec<u8>>,
transport: NoiseTransport,
) -> NoiseVirtualStream {
let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (disconnected_tx, disconnected_rx) = watch::channel(false);
let transport = Arc::new(Mutex::new(transport));
let writer_transport = Arc::clone(&transport);
let writer_stream_id = stream_id;
let writer_task = tokio::spawn(async move {
let mut next_seq = 0u32;
'writer: while let Some(message) = json_outgoing_rx.recv().await {
// Frame first, then split into bounded Noise records. Each record
// receives one checked relay sequence and is encrypted exactly
// once, preserving the implicit Noise sending nonce.
let framed = match frame_jsonrpc_message(&message) {
Ok(framed) => framed,
Err(error) => {
warn!("failed to frame Noise virtual stream JSON-RPC payload: {error}");
break;
}
};
for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) {
let seq = match take_next_sequence(&mut next_seq) {
Ok(seq) => seq,
Err(error) => {
warn!("Noise virtual stream sequence exhausted: {error}");
break 'writer;
}
};
let ciphertext = {
let mut transport = writer_transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.encrypt(plaintext_record)
};
let ciphertext = match ciphertext {
Ok(ciphertext) => ciphertext,
Err(error) => {
warn!("failed to encrypt Noise virtual stream payload: {error}");
break 'writer;
}
};
let frame = RelayMessageFrame::data(writer_stream_id.clone(), seq, ciphertext);
if physical_outgoing_tx
.send(encode_relay_message_frame(&frame))
.await
.is_err()
{
break 'writer;
}
}
}
// Tell the harness to discard this virtual stream whenever its writer
// exits, including processor shutdown or a cryptographic/send failure.
// Otherwise the peer could wait indefinitely on a dead stream.
let reset =
RelayMessageFrame::reset(writer_stream_id, NOISE_RELAY_RESET_REASON.to_string());
let _ = physical_outgoing_tx
.send(encode_relay_message_frame(&reset))
.await;
});
let connection = JsonRpcConnection {
outgoing_tx: json_outgoing_tx,
incoming_rx,
disconnected_rx,
task_handles: vec![writer_task],
transport: JsonRpcTransport::External,
};
tokio::spawn(async move {
processor.run_connection(connection).await;
});
NoiseVirtualStream {
incoming_tx,
disconnected_tx,
transport,
inbound_ciphertexts: OrderedCiphertextFrames::default(),
inbound_decoder: JsonRpcMessageDecoder::default(),
}
}
async fn send_reset(physical_outgoing_tx: &mpsc::Sender<Vec<u8>>, stream_id: String) {
fn send_reset(physical_outgoing_tx: &mpsc::Sender<Vec<u8>>, stream_id: String) {
let reset = RelayMessageFrame::reset(stream_id, NOISE_RELAY_RESET_REASON.to_string());
// Resets are best effort. Untrusted relay input must never block the shared
// state machine behind an overloaded physical writer queue.
let _ = physical_outgoing_tx
.send(encode_relay_message_frame(&reset))
.await;
.try_send(encode_relay_message_frame(&reset));
}

181
codex-rs/exec-server/src/noise_relay/executor_stream.rs Normal file
View File

@@ -0,0 +1,181 @@
use std::sync::Arc;
use std::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tracing::warn;
use crate::ExecServerError;
use crate::connection::CHANNEL_CAPACITY;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
use crate::noise_channel::NoiseTransport;
use crate::noise_relay::NOISE_RELAY_RESET_REASON;
use crate::noise_relay::message_framing::JsonRpcMessageDecoder;
use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN;
use crate::noise_relay::message_framing::frame_jsonrpc_message;
use crate::noise_relay::ordered_ciphertext::OrderedCiphertextFrames;
use crate::noise_relay::take_next_sequence;
use crate::relay::encode_relay_message_frame;
use crate::relay_proto::RelayData;
use crate::relay_proto::RelayMessageFrame;
use crate::server::ConnectionProcessor;
/// Identifies one completed virtual-stream instance.
///
/// Stream IDs are supplied by the untrusted relay peer and may be reused. The
/// instance ID prevents a delayed writer notification from removing a newer
/// stream that happens to use the same routing ID.
pub(super) struct ClosedNoiseVirtualStream {
pub(super) stream_id: String,
pub(super) instance_id: u64,
}
/// One authenticated JSON-RPC stream carried by the executor's physical relay.
///
/// Inbound delivery is intentionally nonblocking. An overloaded or abandoned
/// stream fails independently instead of stalling every stream multiplexed over
/// the same physical websocket.
pub(super) struct NoiseVirtualStream {
incoming_tx: mpsc::Sender<JsonRpcConnectionEvent>,
disconnected_tx: watch::Sender<bool>,
transport: Arc<Mutex<NoiseTransport>>,
inbound_ciphertexts: OrderedCiphertextFrames,
inbound_decoder: JsonRpcMessageDecoder,
instance_id: u64,
}
impl NoiseVirtualStream {
pub(super) fn disconnect(self, reason: Option<String>) {
let _ = self.disconnected_tx.send(true);
let _ = self
.incoming_tx
.try_send(JsonRpcConnectionEvent::Disconnected { reason });
}
pub(super) fn is_instance(&self, instance_id: u64) -> bool {
self.instance_id == instance_id
}
pub(super) fn receive_data(&mut self, data: RelayData) -> Result<(), ExecServerError> {
// Relay sequence ordering is enforced before taking the transport lock
// and decrypting. Each virtual stream owns one ordered Noise nonce
// space shared by its reader and writer transport halves.
for ciphertext in self.inbound_ciphertexts.push(data.seq, data.payload)? {
let plaintext = {
let mut transport = self
.transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.decrypt(&ciphertext).map_err(|error| {
ExecServerError::Protocol(format!("Noise relay decryption failed: {error}"))
})?
};
for message in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
.try_send(JsonRpcConnectionEvent::Message(message))
.map_err(|_| {
ExecServerError::Protocol(
"Noise virtual stream inbound queue is full or closed".to_string(),
)
})?;
}
}
Ok(())
}
}
pub(super) fn spawn_noise_virtual_stream(
stream_id: String,
instance_id: u64,
processor: ConnectionProcessor,
physical_outgoing_tx: mpsc::Sender<Vec<u8>>,
closed_stream_tx: mpsc::Sender<ClosedNoiseVirtualStream>,
transport: NoiseTransport,
) -> NoiseVirtualStream {
let (json_outgoing_tx, mut json_outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (disconnected_tx, disconnected_rx) = watch::channel(false);
let transport = Arc::new(Mutex::new(transport));
let writer_transport = Arc::clone(&transport);
let writer_stream_id = stream_id;
let writer_task = tokio::spawn(async move {
let mut next_seq = 0u32;
'writer: while let Some(message) = json_outgoing_rx.recv().await {
// Frame first, then split into bounded Noise records. Each record
// receives one checked relay sequence and is encrypted exactly
// once, preserving the implicit Noise sending nonce.
let framed = match frame_jsonrpc_message(&message) {
Ok(framed) => framed,
Err(error) => {
warn!("failed to frame Noise virtual stream JSON-RPC payload: {error}");
break;
}
};
for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) {
let seq = match take_next_sequence(&mut next_seq) {
Ok(seq) => seq,
Err(error) => {
warn!("Noise virtual stream sequence exhausted: {error}");
break 'writer;
}
};
let ciphertext = {
let mut transport = writer_transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.encrypt(plaintext_record)
};
let ciphertext = match ciphertext {
Ok(ciphertext) => ciphertext,
Err(error) => {
warn!("failed to encrypt Noise virtual stream payload: {error}");
break 'writer;
}
};
let frame = RelayMessageFrame::data(writer_stream_id.clone(), seq, ciphertext);
if physical_outgoing_tx
.send(encode_relay_message_frame(&frame))
.await
.is_err()
{
break 'writer;
}
}
}
// Reset is best effort because an overloaded physical writer must not
// keep this dead stream alive. The reliable local close notification
// below lets the shared state machine reap the exact stream instance.
let closed_stream = ClosedNoiseVirtualStream {
stream_id: writer_stream_id.clone(),
instance_id,
};
let reset =
RelayMessageFrame::reset(writer_stream_id, NOISE_RELAY_RESET_REASON.to_string());
let _ = physical_outgoing_tx
.try_send(encode_relay_message_frame(&reset));
let _ = closed_stream_tx.send(closed_stream).await;
});
let connection = JsonRpcConnection {
outgoing_tx: json_outgoing_tx,
incoming_rx,
disconnected_rx,
task_handles: vec![writer_task],
transport: JsonRpcTransport::External,
};
tokio::spawn(async move {
processor.run_connection(connection).await;
});
NoiseVirtualStream {
incoming_tx,
disconnected_tx,
transport,
inbound_ciphertexts: OrderedCiphertextFrames::default(),
inbound_decoder: JsonRpcMessageDecoder::default(),
instance_id,
}
}

16
codex-rs/exec-server/src/noise_relay/harness.rs
View File

@@ -30,6 +30,11 @@ use crate::relay::encode_relay_message_frame;
use crate::relay_proto::RelayData;
use crate::relay_proto::RelayMessageFrame;
// Reset frames are cleartext relay control and are not authenticated by Noise.
// Preserve the availability signal while replacing attacker-controlled reason
// text before it reaches disconnect diagnostics.
const NOISE_RELAY_RESET_DISCONNECT_REASON: &str = "Noise relay stream reset";
/// Adapt one harness rendezvous websocket into an authenticated JSON-RPC connection.
///
/// The returned connection is not usable until the background task completes
@@ -199,9 +204,7 @@ where
send_disconnected(
&incoming_tx,
&disconnected_tx,
frame
.into_reset_reason()
.unwrap_or_else(|| "Noise relay reset during handshake".to_string()),
NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(),
)
.await;
return;
@@ -306,9 +309,12 @@ where
}
}
Ok(RelayFrameBodyKind::Reset) => {
let reason = frame.into_reset_reason();
let _ = incoming_tx
.send(JsonRpcConnectionEvent::Disconnected { reason })
.send(JsonRpcConnectionEvent::Disconnected {
reason: Some(
NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(),
),
})
.await;
break;
}

5
codex-rs/exec-server/src/noise_relay/mod.rs
View File

@@ -1,3 +1,4 @@
mod executor_stream;
mod environment;
mod harness;
mod message_framing;
@@ -11,6 +12,10 @@ pub(crate) use environment::HarnessKeyValidator;
pub(crate) use environment::run_noise_multiplexed_environment;
pub(crate) use harness::noise_harness_connection_from_websocket;
// This value is already part of the relay wire contract. Keep it stable even
// though the source module now uses the more precise Noise terminology.
const NOISE_RELAY_RESET_REASON: &str = "secure_relay_protocol_error";
// This bounds allocation in tungstenite before protobuf and Noise record
// validation run. It comfortably fits one maximum Noise record plus metadata.
const MAX_NOISE_RELAY_WEBSOCKET_MESSAGE_SIZE: usize = 256 * 1024;

1
codex-rs/exec-server/src/relay_proto.rs
View File

@@ -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;

275
codex-rs/exec-server/src/remote.rs
View File

@@ -3,29 +3,20 @@ use std::time::Duration;
use codex_api::SharedAuthProvider;
use reqwest::StatusCode;
use serde::Deserialize;
use serde::Serialize;
use tokio::time::sleep;
use tokio::time::timeout;
use tokio_tungstenite::connect_async_with_config;
use tracing::info;
use tokio_tungstenite::connect_async;
use tracing::warn;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::NoiseChannelIdentity;
use crate::NoiseChannelPublicKey;
use crate::noise_relay::HarnessKeyValidator;
use crate::noise_relay::noise_relay_websocket_config;
use crate::noise_relay::run_noise_multiplexed_environment;
use crate::relay::run_multiplexed_environment;
use crate::server::ConnectionProcessor;
mod noise;
const ERROR_BODY_PREVIEW_BYTES: usize = 4096;
const ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_REMOTE_ENVIRONMENT_ID_LEN: usize = 256;
const REMOTE_RENDEZVOUS_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
#[derive(Clone)]
struct EnvironmentRegistryClient {
@@ -51,16 +42,16 @@ impl EnvironmentRegistryClient {
auth_provider,
http: reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT)
.build()?,
})
}
/// Register using the original body-less registry contract.
///
/// This remains the default so a new exec-server can still connect through
/// registries and harnesses that have not rolled out Noise relay support.
async fn register_legacy_environment(
/// This method intentionally preserves the legacy request shape and timeout
/// behavior. Noise support is opt-in and must not silently alter existing
/// remote exec-server registrations.
async fn register_environment(
&self,
environment_id: &str,
) -> Result<EnvironmentRegistryRegistrationResponse, ExecServerError> {
@@ -76,64 +67,6 @@ impl EnvironmentRegistryClient {
self.parse_json_response(response).await
}
/// Register the exec-server's static Noise identity with a Noise-aware registry.
///
/// Supplying this request body is the protocol-level opt in. Registries can
/// therefore distinguish Noise registrations from the legacy body-less
/// contract without guessing based on binary version or rollout state.
async fn register_noise_environment(
&self,
environment_id: &str,
executor_public_key: &NoiseChannelPublicKey,
) -> Result<EnvironmentRegistryNoiseRegistrationResponse, ExecServerError> {
let response = self
.http
.post(endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/register"),
))
.headers(self.auth_provider.to_auth_headers())
.json(&EnvironmentRegistryRegistrationRequest {
security_profile: NOISE_RELAY_SECURITY_PROFILE,
executor_public_key,
})
.send()
.await?;
self.parse_json_response(response).await
}
async fn validate_harness_key(
&self,
environment_id: &str,
executor_registration_id: &str,
harness_public_key: &NoiseChannelPublicKey,
harness_key_authorization: &str,
) -> Result<(), ExecServerError> {
let response = self
.http
.post(endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/validate"),
))
.headers(self.auth_provider.to_auth_headers())
.json(&EnvironmentRegistryHarnessKeyValidationRequest {
executor_registration_id,
harness_public_key,
harness_key_authorization,
})
.send()
.await?;
let response = self
.parse_json_response::<EnvironmentRegistryHarnessKeyValidationResponse>(response)
.await?;
if !response.valid {
return Err(ExecServerError::Protocol(
"environment registry rejected Noise relay harness key".to_string(),
));
}
Ok(())
}
async fn parse_json_response<R>(
&self,
response: reqwest::Response,
@@ -155,64 +88,12 @@ impl EnvironmentRegistryClient {
}
}
const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1";
#[derive(Serialize)]
struct EnvironmentRegistryRegistrationRequest<'a> {
security_profile: &'static str,
executor_public_key: &'a NoiseChannelPublicKey,
}
#[derive(Debug, Clone, Eq, PartialEq, Deserialize)]
struct EnvironmentRegistryRegistrationResponse {
environment_id: String,
url: String,
}
#[derive(Debug, Clone, Eq, PartialEq, Deserialize)]
struct EnvironmentRegistryNoiseRegistrationResponse {
environment_id: String,
url: String,
security_profile: String,
executor_registration_id: String,
}
#[derive(Serialize)]
struct EnvironmentRegistryHarnessKeyValidationRequest<'a> {
executor_registration_id: &'a str,
harness_public_key: &'a NoiseChannelPublicKey,
harness_key_authorization: &'a str,
}
#[derive(Deserialize)]
struct EnvironmentRegistryHarnessKeyValidationResponse {
valid: bool,
}
#[derive(Clone)]
struct RegistryHarnessKeyValidator {
client: EnvironmentRegistryClient,
environment_id: String,
executor_registration_id: String,
}
impl HarnessKeyValidator for RegistryHarnessKeyValidator {
async fn validate_harness_key(
&self,
harness_public_key: &NoiseChannelPublicKey,
authorization: &str,
) -> Result<(), ExecServerError> {
self.client
.validate_harness_key(
&self.environment_id,
&self.executor_registration_id,
harness_public_key,
authorization,
)
.await
}
}
/// Protocol used for an exec-server's registered remote relay.
///
/// Legacy is intentionally the default during rollout. Noise must be selected
@@ -220,40 +101,13 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator {
/// contract until both endpoints are ready for Noise.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum RemoteRelayProtocol {
/// Original cleartext JSON-RPC relay protocol.
#[default]
Legacy,
/// Authenticated, end-to-end encrypted Noise relay protocol.
Noise,
}
enum RemoteRelayProtocolState {
Legacy,
Noise(NoiseChannelIdentity),
}
enum RegisteredRemoteEnvironment {
Legacy(EnvironmentRegistryRegistrationResponse),
Noise {
response: EnvironmentRegistryNoiseRegistrationResponse,
identity: NoiseChannelIdentity,
},
}
impl RegisteredRemoteEnvironment {
fn environment_id(&self) -> &str {
match self {
Self::Legacy(response) => &response.environment_id,
Self::Noise { response, .. } => &response.environment_id,
}
}
fn websocket_url(&self) -> &str {
match self {
Self::Legacy(response) => &response.url,
Self::Noise { response, .. } => &response.url,
}
}
}
/// Configuration for registering an exec-server for remote use.
#[derive(Clone)]
pub struct RemoteEnvironmentConfig {
@@ -303,92 +157,27 @@ pub async fn run_remote_environment(
let client =
EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?;
let processor = ConnectionProcessor::new(runtime_paths);
let protocol_state = match config.relay_protocol {
RemoteRelayProtocol::Legacy => RemoteRelayProtocolState::Legacy,
RemoteRelayProtocol::Noise => RemoteRelayProtocolState::Noise(
NoiseChannelIdentity::generate().map_err(|error| {
ExecServerError::Protocol(format!(
"failed to generate Noise relay identity: {error}"
))
})?,
),
};
if config.relay_protocol == RemoteRelayProtocol::Noise {
return noise::run_remote_environment(&config, &client, processor).await;
}
let mut backoff = Duration::from_secs(1);
loop {
let registration = match &protocol_state {
RemoteRelayProtocolState::Legacy => RegisteredRemoteEnvironment::Legacy(
client
.register_legacy_environment(&config.environment_id)
.await?,
),
RemoteRelayProtocolState::Noise(identity) => {
let response = client
.register_noise_environment(&config.environment_id, &identity.public_key())
.await?;
if response.security_profile != NOISE_RELAY_SECURITY_PROFILE {
return Err(ExecServerError::Protocol(format!(
"environment registry returned unsupported security profile `{}`",
response.security_profile
)));
}
RegisteredRemoteEnvironment::Noise {
response,
identity: identity.clone(),
}
}
};
if registration.environment_id() != config.environment_id {
return Err(ExecServerError::Protocol(
"environment registry returned a different environment id".to_string(),
));
}
let environment_id = registration.environment_id();
info!(
"codex exec-server remote environment registered with environment_id {environment_id}"
let response = client.register_environment(&config.environment_id).await?;
eprintln!(
"codex exec-server remote environment registered with environment_id {}",
response.environment_id
);
let websocket_config = match &registration {
RegisteredRemoteEnvironment::Legacy(_) => None,
RegisteredRemoteEnvironment::Noise { .. } => Some(noise_relay_websocket_config()),
};
match timeout(
REMOTE_RENDEZVOUS_CONNECT_TIMEOUT,
connect_async_with_config(
registration.websocket_url(),
websocket_config,
/*disable_nagle*/ false,
),
)
.await
{
Ok(Ok((websocket, _))) => {
match connect_async(response.url.as_str()).await {
Ok((websocket, _)) => {
backoff = Duration::from_secs(1);
match registration {
RegisteredRemoteEnvironment::Legacy(_) => {
run_multiplexed_environment(websocket, processor.clone()).await;
}
RegisteredRemoteEnvironment::Noise { response, identity } => {
run_noise_multiplexed_environment(
websocket,
processor.clone(),
response.environment_id,
response.executor_registration_id.clone(),
identity,
RegistryHarnessKeyValidator {
client: client.clone(),
environment_id: config.environment_id.clone(),
executor_registration_id: response.executor_registration_id,
},
)
.await;
}
}
run_multiplexed_environment(websocket, processor.clone()).await;
}
Ok(Err(err)) => {
Err(err) => {
warn!("failed to connect remote exec-server websocket: {err}");
}
Err(_) => warn!("timed out connecting remote exec-server websocket"),
}
sleep(backoff).await;
@@ -397,32 +186,12 @@ pub async fn run_remote_environment(
}
fn normalize_environment_id(environment_id: String) -> Result<String, ExecServerError> {
let environment_id = environment_id.trim().to_string();
if environment_id.is_empty() {
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id is required for remote exec-server registration".to_string(),
));
}
if environment_id.trim() != environment_id {
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id must not contain surrounding whitespace".to_string(),
));
}
if environment_id.len() > MAX_REMOTE_ENVIRONMENT_ID_LEN {
return Err(ExecServerError::EnvironmentRegistryConfig(format!(
"environment id cannot be longer than {MAX_REMOTE_ENVIRONMENT_ID_LEN} characters"
)));
}
// The ID is interpolated into authenticated registry request paths below.
// Keep it to one URL path component so a caller cannot use a delimiter to
// redirect the exec-server's registration credential to another endpoint.
if !environment_id
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_')
{
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id must contain only ASCII letters, numbers, '-' or '_'".to_string(),
));
}
Ok(environment_id)
}

270
codex-rs/exec-server/src/remote/noise.rs Normal file
View File

@@ -0,0 +1,270 @@
//! Explicitly selected Noise registration and encrypted relay runtime.
//!
//! The legacy remote-exec path stays in the parent module. Keeping the Noise
//! path here makes the opt-in boundary visible and keeps its stricter registry,
//! identifier, timeout, and websocket requirements from changing legacy
//! behavior accidentally.
use std::time::Duration;
use reqwest::StatusCode;
use serde::Deserialize;
use serde::Serialize;
use tokio::time::sleep;
use tokio::time::timeout;
use tokio_tungstenite::connect_async_with_config;
use tracing::info;
use tracing::warn;
use super::EnvironmentRegistryClient;
use super::RemoteEnvironmentConfig;
use super::endpoint_url;
use crate::ExecServerError;
use crate::NoiseChannelIdentity;
use crate::NoiseChannelPublicKey;
use crate::noise_relay::HarnessKeyValidator;
use crate::noise_relay::noise_relay_websocket_config;
use crate::noise_relay::run_noise_multiplexed_environment;
use crate::server::ConnectionProcessor;
const ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_EXECUTOR_REGISTRATION_ID_LEN: usize = 256;
const MAX_REMOTE_ENVIRONMENT_ID_LEN: usize = 256;
const NOISE_RELAY_SECURITY_PROFILE: &str = "noise_hybrid_ik_v1";
const REMOTE_RENDEZVOUS_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
impl EnvironmentRegistryClient {
/// Register the exec-server's static Noise identity with a Noise-aware registry.
///
/// Supplying this request body is the protocol-level opt in. Registries can
/// therefore distinguish Noise registrations from the legacy body-less
/// contract without guessing based on binary version or rollout state.
async fn register_noise_environment(
&self,
environment_id: &str,
executor_public_key: &NoiseChannelPublicKey,
) -> Result<EnvironmentRegistryNoiseRegistrationResponse, ExecServerError> {
let response = self
.http
.post(endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/register"),
))
.headers(self.auth_provider.to_auth_headers())
.timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT)
.json(&EnvironmentRegistryRegistrationRequest {
security_profile: NOISE_RELAY_SECURITY_PROFILE,
executor_public_key,
})
.send()
.await?;
self.parse_json_response(response).await
}
/// Validate the authenticated harness key without exposing its authorization.
async fn validate_harness_key(
&self,
environment_id: &str,
executor_registration_id: &str,
harness_public_key: &NoiseChannelPublicKey,
harness_key_authorization: &str,
) -> Result<(), ExecServerError> {
let response = self
.http
.post(endpoint_url(
&self.base_url,
&format!("/cloud/environment/{environment_id}/validate"),
))
.headers(self.auth_provider.to_auth_headers())
.timeout(ENVIRONMENT_REGISTRY_REQUEST_TIMEOUT)
.json(&EnvironmentRegistryHarnessKeyValidationRequest {
executor_registration_id,
harness_public_key,
harness_key_authorization,
})
.send()
.await?;
let status = response.status();
if !status.is_success() {
// This request contains the short-lived harness authorization.
// Never propagate a response body that might echo it into logs or
// user-visible error chains.
if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) {
return Err(ExecServerError::EnvironmentRegistryAuth(format!(
"environment registry harness key validation authentication failed ({status})"
)));
}
return Err(ExecServerError::EnvironmentRegistryHttp {
status,
code: None,
message: "environment registry harness key validation failed".to_string(),
});
}
let response = response
.json::<EnvironmentRegistryHarnessKeyValidationResponse>()
.await?;
if !response.valid {
return Err(ExecServerError::Protocol(
"environment registry rejected Noise relay harness key".to_string(),
));
}
Ok(())
}
}
#[derive(Serialize)]
struct EnvironmentRegistryRegistrationRequest<'a> {
security_profile: &'static str,
executor_public_key: &'a NoiseChannelPublicKey,
}
#[derive(Deserialize)]
struct EnvironmentRegistryNoiseRegistrationResponse {
environment_id: String,
url: String,
security_profile: String,
executor_registration_id: String,
}
#[derive(Serialize)]
struct EnvironmentRegistryHarnessKeyValidationRequest<'a> {
executor_registration_id: &'a str,
harness_public_key: &'a NoiseChannelPublicKey,
harness_key_authorization: &'a str,
}
#[derive(Deserialize)]
struct EnvironmentRegistryHarnessKeyValidationResponse {
valid: bool,
}
#[derive(Clone)]
struct RegistryHarnessKeyValidator {
client: EnvironmentRegistryClient,
environment_id: String,
executor_registration_id: String,
}
impl HarnessKeyValidator for RegistryHarnessKeyValidator {
async fn validate_harness_key(
&self,
harness_public_key: &NoiseChannelPublicKey,
authorization: &str,
) -> Result<(), ExecServerError> {
self.client
.validate_harness_key(
&self.environment_id,
&self.executor_registration_id,
harness_public_key,
authorization,
)
.await
}
}
/// Run the Noise registration and encrypted relay loop.
///
/// A new executor identity is generated once per process invocation and reused
/// across physical reconnects. Every registry registration still receives a new
/// registration ID, which is bound into each virtual stream's Noise prologue.
pub(super) async fn run_remote_environment(
config: &RemoteEnvironmentConfig,
client: &EnvironmentRegistryClient,
processor: ConnectionProcessor,
) -> Result<(), ExecServerError> {
validate_environment_id(&config.environment_id)?;
let identity = NoiseChannelIdentity::generate().map_err(|error| {
ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
})?;
let mut backoff = Duration::from_secs(1);
loop {
let response = client
.register_noise_environment(&config.environment_id, &identity.public_key())
.await?;
if response.environment_id != config.environment_id {
return Err(ExecServerError::Protocol(
"environment registry returned a different environment id".to_string(),
));
}
if response.security_profile != NOISE_RELAY_SECURITY_PROFILE {
return Err(ExecServerError::Protocol(format!(
"environment registry returned unsupported security profile `{}`",
response.security_profile
)));
}
validate_executor_registration_id(&response.executor_registration_id)?;
let environment_id = &response.environment_id;
info!(
"codex exec-server Noise environment registered with environment_id {environment_id}"
);
match timeout(
REMOTE_RENDEZVOUS_CONNECT_TIMEOUT,
connect_async_with_config(
response.url.as_str(),
Some(noise_relay_websocket_config()),
/*disable_nagle*/ false,
),
)
.await
{
Ok(Ok((websocket, _))) => {
backoff = Duration::from_secs(1);
let executor_registration_id = response.executor_registration_id;
run_noise_multiplexed_environment(
websocket,
processor.clone(),
response.environment_id,
executor_registration_id.clone(),
identity.clone(),
RegistryHarnessKeyValidator {
client: client.clone(),
environment_id: config.environment_id.clone(),
executor_registration_id,
},
)
.await;
}
Ok(Err(err)) => warn!("failed to connect Noise remote exec-server websocket: {err}"),
Err(_) => warn!("timed out connecting Noise remote exec-server websocket"),
}
sleep(backoff).await;
backoff = (backoff * 2).min(Duration::from_secs(30));
}
}
fn validate_environment_id(environment_id: &str) -> Result<(), ExecServerError> {
if environment_id.len() > MAX_REMOTE_ENVIRONMENT_ID_LEN {
return Err(ExecServerError::EnvironmentRegistryConfig(format!(
"environment id cannot be longer than {MAX_REMOTE_ENVIRONMENT_ID_LEN} characters"
)));
}
// The ID is interpolated into authenticated registry request paths below.
// Keep it to one URL path component so a caller cannot use a delimiter to
// redirect the exec-server's registration credential to another endpoint.
if !environment_id
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_')
{
return Err(ExecServerError::EnvironmentRegistryConfig(
"environment id must contain only ASCII letters, numbers, '-' or '_'".to_string(),
));
}
Ok(())
}
fn validate_executor_registration_id(
executor_registration_id: &str,
) -> Result<(), ExecServerError> {
if executor_registration_id.is_empty()
|| executor_registration_id.trim() != executor_registration_id
|| executor_registration_id.len() > MAX_EXECUTOR_REGISTRATION_ID_LEN
{
return Err(ExecServerError::Protocol(
"environment registry returned an invalid executor registration id".to_string(),
));
}
Ok(())
}