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

Decouple Noise relay streams from JSON-RPC processing (#39235)

Browse Source 
## What changed

- Add a `NoiseStreamHandler` abstraction that owns payload encoding, decoding,
  and connection processing while the virtual stream handles framing,
  encryption, and multiplexing.
- Reassemble authenticated payloads as opaque bytes, with the existing
  JSON-RPC processor supplied as a handler implementation.
- Have the physical relay send a reset after the current stream instance
  closes, including when its processor exits before the writer task.

## Testing

- Verify local JSON-RPC decoding creates the queued request span before queue
  admission.
- Verify a processor exit resets the corresponding Noise harness stream.

GitOrigin-RevId: 270befcc8bfbf27e46c84d6a2c6e349cb98ea3f8
This commit is contained in:
pakrym-oai
2026-08-18 17:18:50 +00:00
committed by copyberry
parent 9b9b614b02
commit 3df5087f75
10 changed files with 582 additions and 302 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/noise_relay/executor_stream.rs
+34 -54
View File
@@ -7,28 +7,22 @@
use std::sync::Arc;
use std::sync::Mutex;
use codex_exec_server_protocol::JSONRPCMessage;
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::MessageDecoder;
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::stream_handler::NoiseStreamConnection;
use crate::noise_relay::stream_handler::NoiseStreamHandler;
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;
use crate::telemetry::ConnectionTransport;
/// Identifies one completed virtual-stream instance.
///
@@ -40,29 +34,26 @@ pub(crate) struct ClosedNoiseVirtualStream {
pub(crate) instance_id: u64,
}
/// One authenticated JSON-RPC stream carried by the executor's physical relay.
/// One authenticated application 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<JsonRpcConnectionEvent>,
pub(crate) struct NoiseVirtualStream<H: NoiseStreamHandler> {
incoming_tx: mpsc::Sender<H::Incoming>,
disconnected_tx: watch::Sender<bool>,
transport: Arc<Mutex<NoiseTransport>>,
inbound_ciphertexts: OrderedCiphertextFrames,
inbound_decoder: JsonRpcMessageDecoder,
inbound_decoder: MessageDecoder,
pub(crate) instance_id: u64,
}
impl NoiseVirtualStream {
pub(crate) fn disconnect(self, reason: Option<String>) {
impl<H: NoiseStreamHandler> NoiseVirtualStream<H> {
pub(crate) fn disconnect(self) {
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.
/// Reorder and decrypt one record, then deliver complete payloads to the handler.
/// 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)? {
@@ -77,7 +68,7 @@ impl NoiseVirtualStream {
};
for message in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
.try_send(JsonRpcConnectionEvent::message(message))
.try_send(H::decode(message)?)
.map_err(|_| {
ExecServerError::Protocol(
"Noise virtual stream inbound queue is full or closed".to_string(),
@@ -89,44 +80,39 @@ impl NoiseVirtualStream {
}
}
/// Start JSON-RPC processing for a completed handshake.
/// Hand a completed handshake to its execution or forwarding owner.
///
/// 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(
pub(crate) fn spawn_noise_virtual_stream<H: NoiseStreamHandler>(
stream_id: String,
instance_id: u64,
processor: ConnectionProcessor,
handler: H,
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);
) -> NoiseVirtualStream<H> {
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 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 owner_stream_id = stream_id.clone();
let owner_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,
'writer: while let Some(message) = outgoing_rx.recv().await {
let message = match H::encode(message) {
Ok(message) => message,
Err(error) => {
warn!("failed to frame Noise virtual stream JSON-RPC payload: {error}");
warn!("failed to encode Noise virtual stream payload: {error}");
break;
}
};
let mut trace = match message {
JSONRPCMessage::Request(request) => request.trace,
JSONRPCMessage::Notification(_)
| JSONRPCMessage::Response(_)
| JSONRPCMessage::Error(_) => None,
};
for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) {
// Each chunk becomes one Noise record and consumes one nonce.
let mut trace = message.trace;
for plaintext_record in message.framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) {
let seq = match take_next_sequence(&mut next_seq) {
Ok(seq) => seq,
Err(error) => {
@@ -163,31 +149,25 @@ pub(crate) fn spawn_noise_virtual_stream(
}
}
// The reset is best effort; the local close notification is not.
// The physical relay owns reset delivery and rejects stale instance IDs.
let closed_stream = ClosedNoiseVirtualStream {
stream_id: writer_stream_id.clone(),
stream_id: writer_stream_id,
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,
let connection = NoiseStreamConnection {
outgoing_tx,
incoming_rx,
disconnected_rx,
task_handles: vec![writer_task],
transport: JsonRpcTransport::Plain,
writer_task,
};
tokio::spawn(async move {
processor
.run_connection(connection, ConnectionTransport::Relay)
.await;
let _ = processor_closed_stream_tx
handler.run_connection(connection).await;
let _ = owner_closed_stream_tx
.send(ClosedNoiseVirtualStream {
stream_id: processor_stream_id,
stream_id: owner_stream_id,
instance_id,
})
.await;
@@ -198,7 +178,7 @@ pub(crate) fn spawn_noise_virtual_stream(
disconnected_tx,
transport,
inbound_ciphertexts: OrderedCiphertextFrames::default(),
inbound_decoder: JsonRpcMessageDecoder::default(),
inbound_decoder: MessageDecoder::default(),
instance_id,
}
}
codex-rs/exec-server/src/noise_relay/message_framing.rs
+29 -10
View File
@@ -1,4 +1,5 @@
use bytes::Buf;
use bytes::Bytes;
use bytes::BytesMut;
use codex_exec_server_protocol::JSONRPCMessage;
@@ -17,14 +18,18 @@ pub(crate) const NOISE_RECORD_PLAINTEXT_LEN: usize = 60 * 1024;
pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> {
let mut framed = vec![0; LENGTH_PREFIX_BYTES];
serde_json::to_writer(&mut framed, message)?;
let message_len = framed.len() - LENGTH_PREFIX_BYTES;
if message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN {
let prefix = message_length_prefix(framed.len() - LENGTH_PREFIX_BYTES)?;
framed[..LENGTH_PREFIX_BYTES].copy_from_slice(&prefix);
Ok(framed)
}
fn message_length_prefix(message_len: usize) -> Result<[u8; LENGTH_PREFIX_BYTES], ExecServerError> {
if message_len == 0 || message_len > MAX_NOISE_JSONRPC_MESSAGE_LEN {
return Err(ExecServerError::Protocol(
"Noise relay JSON-RPC message exceeds maximum length".to_string(),
));
}
framed[..LENGTH_PREFIX_BYTES].copy_from_slice(&(message_len as u32).to_be_bytes());
Ok(framed)
Ok((message_len as u32).to_be_bytes())
}
/// Incrementally reconstructs authenticated JSON-RPC messages from Noise records.
@@ -33,15 +38,31 @@ pub(crate) fn frame_jsonrpc_message(message: &JSONRPCMessage) -> Result<Vec<u8>,
/// here so a bad authenticated peer cannot grow the reassembly buffer forever.
#[derive(Default)]
pub(crate) struct JsonRpcMessageDecoder {
buffered: BytesMut,
decoder: MessageDecoder,
}
impl JsonRpcMessageDecoder {
/// Append one decrypted record and return all complete framed messages.
pub(crate) fn push(
&mut self,
plaintext_record: &[u8],
) -> Result<Vec<JSONRPCMessage>, ExecServerError> {
self.decoder
.push(plaintext_record)?
.into_iter()
.map(|message| serde_json::from_slice(&message).map_err(Into::into))
.collect()
}
}
/// Reassembles opaque application payloads without interpreting their schema.
#[derive(Default)]
pub(crate) struct MessageDecoder {
buffered: BytesMut,
}
impl MessageDecoder {
/// Append one decrypted record and return all complete framed messages.
pub(crate) fn push(&mut self, plaintext_record: &[u8]) -> Result<Vec<Bytes>, ExecServerError> {
if plaintext_record.len() > NOISE_RECORD_PLAINTEXT_LEN {
return Err(ExecServerError::Protocol(
"Noise relay plaintext record exceeds maximum length".to_string(),
@@ -66,10 +87,8 @@ impl JsonRpcMessageDecoder {
if self.buffered.len() < framed_len {
break;
}
messages.push(serde_json::from_slice(
&self.buffered[LENGTH_PREFIX_BYTES..framed_len],
)?);
self.buffered.advance(framed_len);
self.buffered.advance(LENGTH_PREFIX_BYTES);
messages.push(self.buffered.split_to(message_len).freeze());
}
// Even before a message is complete, keep reassembly memory bounded.
codex-rs/exec-server/src/noise_relay/mod.rs
+1
View File
@@ -2,6 +2,7 @@ pub(crate) mod executor_stream;
mod harness;
mod message_framing;
mod ordered_ciphertext;
pub(crate) mod stream_handler;
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
codex-rs/exec-server/src/noise_relay/stream_handler.rs
+87
View File
@@ -0,0 +1,87 @@
use std::future::Future;
use bytes::Bytes;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_protocol::protocol::W3cTraceContext;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinHandle;
use crate::ExecServerError;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
use crate::noise_relay::message_framing::frame_jsonrpc_message;
use crate::server::ConnectionProcessor;
use crate::telemetry::ConnectionTransport;
pub(crate) struct NoiseStreamConnection<I, O> {
pub(crate) outgoing_tx: mpsc::Sender<O>,
pub(crate) incoming_rx: mpsc::Receiver<I>,
pub(crate) disconnected_rx: watch::Receiver<bool>,
pub(crate) writer_task: JoinHandle<()>,
}
pub(crate) struct NoiseOutboundMessage {
/// Payload with the authenticated length prefix supplied by message_framing.
pub(crate) framed: Vec<u8>,
pub(crate) trace: Option<W3cTraceContext>,
}
/// Adapts complete Noise payloads to one owner without adding an inbound queue.
/// Decoding is synchronous so execution spans begin before queue admission.
pub(crate) trait NoiseStreamHandler: Clone + Send + 'static {
type Incoming: Send + 'static;
type Outgoing: Send + 'static;
fn decode(payload: Bytes) -> Result<Self::Incoming, ExecServerError>;
fn encode(message: Self::Outgoing) -> Result<NoiseOutboundMessage, ExecServerError>;
fn run_connection(
self,
connection: NoiseStreamConnection<Self::Incoming, Self::Outgoing>,
) -> impl Future<Output = ()> + Send;
}
impl NoiseStreamHandler for ConnectionProcessor {
type Incoming = JsonRpcConnectionEvent;
type Outgoing = JSONRPCMessage;
fn decode(payload: Bytes) -> Result<Self::Incoming, ExecServerError> {
Ok(JsonRpcConnectionEvent::message(serde_json::from_slice(
&payload,
)?))
}
fn encode(message: Self::Outgoing) -> Result<NoiseOutboundMessage, ExecServerError> {
let framed = frame_jsonrpc_message(&message)?;
let trace = match message {
JSONRPCMessage::Request(request) => request.trace,
JSONRPCMessage::Notification(_)
| JSONRPCMessage::Response(_)
| JSONRPCMessage::Error(_) => None,
};
Ok(NoiseOutboundMessage { framed, trace })
}
async fn run_connection(
self,
connection: NoiseStreamConnection<Self::Incoming, Self::Outgoing>,
) {
ConnectionProcessor::run_connection(
&self,
JsonRpcConnection {
outgoing_tx: connection.outgoing_tx,
incoming_rx: connection.incoming_rx,
disconnected_rx: connection.disconnected_rx,
task_handles: vec![connection.writer_task],
transport: JsonRpcTransport::Plain,
},
ConnectionTransport::Relay,
)
.await;
}
}
#[cfg(test)]
#[path = "stream_handler_tests.rs"]
mod tests;
codex-rs/exec-server/src/noise_relay/stream_handler_tests.rs
+41
View File
@@ -0,0 +1,41 @@
use std::time::Instant;
use anyhow::Result;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::RequestId;
use pretty_assertions::assert_eq;
use super::NoiseStreamHandler;
use crate::connection::JsonRpcConnectionEvent;
use crate::server::ConnectionProcessor;
#[test]
fn local_decode_starts_request_span_before_queueing() -> Result<()> {
let _subscriber = tracing::subscriber::set_default(tracing_subscriber::registry());
tracing::callsite::rebuild_interest_cache();
let request = JSONRPCRequest {
id: RequestId::Integer(1),
method: "test/queued".to_string(),
params: None,
trace: None,
};
let before_decode = Instant::now();
let event = <ConnectionProcessor as NoiseStreamHandler>::decode(
serde_json::to_vec(&JSONRPCMessage::Request(request.clone()))?.into(),
)?;
let after_decode = Instant::now();
let JsonRpcConnectionEvent::QueuedRequest {
request: actual_request,
request_span,
queued_at,
} = event
else {
panic!("local ingress must create the queued request synchronously");
};
assert_eq!(actual_request, request);
assert!(!request_span.is_disabled());
assert!((before_decode..=after_decode).contains(&queued_at));
Ok(())
}
codex-rs/exec-server/src/relay.rs
+10 -6
View File
@@ -32,12 +32,14 @@ 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::noise_relay::stream_handler::NoiseStreamHandler;
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;
#[cfg(test)]
use crate::server::ConnectionProcessor;
use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT;
use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT_REASON;
@@ -475,9 +477,9 @@ pub(crate) trait HarnessKeyValidator: Send + Sync {
/// 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"))]
pub(crate) async fn run_multiplexed_environment<T, E, V>(
pub(crate) async fn run_multiplexed_environment<T, E, V, H>(
stream: T,
processor: ConnectionProcessor,
handler: H,
environment_id: String,
executor_registration_id: String,
identity: NoiseChannelIdentity,
@@ -487,6 +489,7 @@ where
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
E: std::fmt::Display + Send + 'static,
V: HarnessKeyValidator + Clone + 'static,
H: NoiseStreamHandler,
{
debug!(
environment_id,
@@ -563,7 +566,7 @@ where
}
}
});
let mut streams: HashMap<String, NoiseVirtualStream> = HashMap::new();
let mut streams: HashMap<String, NoiseVirtualStream<H>> = HashMap::new();
let mut pending_handshakes: HashMap<String, PendingHandshake> = HashMap::new();
let mut validation_tasks: JoinSet<HarnessKeyValidationResult> = JoinSet::new();
let mut failed_handshakes = 0usize;
@@ -591,6 +594,7 @@ where
.is_some_and(|stream| stream.instance_id == closed_stream.instance_id);
if is_current {
streams.remove(&closed_stream.stream_id);
send_reset(&physical_outgoing_tx, closed_stream.stream_id);
}
continue;
}
@@ -678,7 +682,7 @@ where
spawn_noise_virtual_stream(
validation_result.stream_id,
validation_result.validation_id,
processor.clone(),
handler.clone(),
physical_outgoing_tx.clone(),
closed_stream_tx.clone(),
transport,
@@ -877,7 +881,7 @@ where
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);
stream.disconnect();
}
}
RelayFrameBodyKind::Ack
@@ -887,7 +891,7 @@ where
}
for (_stream_id, stream) in streams {
stream.disconnect(/*reason*/ None);
stream.disconnect();
}
// Dropping the JoinSet aborts any registry checks still running.
if !physical_writer_task.is_finished() {
codex-rs/exec-server/src/relay_noise_tests.rs
+65
View File
@@ -4,6 +4,9 @@ use std::sync::atomic::Ordering;
use std::time::Duration;
use anyhow::Result;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_exec_server_protocol::RequestId;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
@@ -13,6 +16,7 @@ use tokio::time::timeout;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::task::AbortOnDropHandle;
use super::HarnessKeyValidator;
use super::MAX_FAILED_NOISE_HANDSHAKES;
@@ -21,10 +25,13 @@ use super::RendezvousDisconnectReason;
use super::run_multiplexed_environment;
use crate::ExecServerError;
use crate::ExecServerRuntimePaths;
use crate::connection::JsonRpcConnectionEvent;
use crate::noise_channel::InitiatorHandshake;
use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_channel::NoiseChannelPublicKey;
use crate::noise_channel::noise_channel_prologue;
use crate::noise_relay::NoiseHarnessConnectionArgs;
use crate::noise_relay::noise_harness_connection_from_websocket;
use crate::relay::RelayFrameBodyKind;
use crate::relay::decode_relay_message_frame;
use crate::relay::encode_relay_message_frame;
@@ -136,6 +143,64 @@ impl HarnessKeyValidator for BlockingValidator {
}
}
#[tokio::test]
async fn processor_exit_resets_noise_harness_stream() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let connecting = tokio::spawn(connect_async(format!("ws://{}", listener.local_addr()?)));
let (socket, _) = listener.accept().await?;
let environment_websocket = accept_async(socket).await?;
let (harness_websocket, _) = connecting.await??;
let identity = NoiseChannelIdentity::generate()?;
let release = Arc::new(Notify::new());
release.notify_one();
let environment_task = AbortOnDropHandle::new(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(),
identity.clone(),
BlockingValidator {
calls: Arc::new(AtomicUsize::new(0)),
release,
},
)));
let mut connection = noise_harness_connection_from_websocket(
harness_websocket,
NoiseHarnessConnectionArgs {
connection_label: "processor exit test".to_string(),
environment_id: ENVIRONMENT_ID.to_string(),
executor_registration_id: EXECUTOR_REGISTRATION_ID.to_string(),
identity: NoiseChannelIdentity::generate()?,
responder_public_key: identity.public_key(),
harness_key_authorization: "authorization".to_string(),
},
);
// Valid JSON reaches the processor; an unsolicited response closes it and
// aborts its writer. The physical relay must still deliver the reset.
connection
.outgoing_tx
.send(JSONRPCMessage::Response(JSONRPCResponse {
id: RequestId::Integer(1),
result: serde_json::Value::Null,
}))
.await?;
assert!(matches!(
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
Some(JsonRpcConnectionEvent::Disconnected { reason: Some(reason) })
if reason == "Noise relay stream reset"
));
for task in connection.task_handles {
task.abort();
let _ = task.await;
}
environment_task.abort();
let _ = environment_task.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?;
codex-rs/exec-server/tests/common/relay.rs
+254
View File
@@ -0,0 +1,254 @@
#[path = "../../src/proto/codex.exec_server.relay.v1.rs"]
mod relay_proto;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_api::AuthProvider;
use codex_exec_server::ExecServerClient;
use codex_exec_server::NoiseChannelIdentity;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectArgs;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::RemoteEnvironmentConfig;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use futures::SinkExt;
use futures::StreamExt;
use http::HeaderMap;
use http::HeaderValue;
use prost::Message as ProstMessage;
use relay_proto::RelayMessageFrame;
use relay_proto::relay_message_frame;
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 tokio_util::task::AbortOnDropHandle;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
pub(crate) const ENVIRONMENT_ID: &str = "env-noise-relay-test";
pub(crate) const EXECUTOR_REGISTRATION_ID: &str = "registration-1";
pub(crate) const HARNESS_KEY_AUTHORIZATION: &str = "harness-key-authorization";
pub(crate) const REGISTRY_TOKEN: &str = "registry-token";
pub(crate) const TEST_TIMEOUT: Duration = Duration::from_secs(30);
#[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"),
);
}
}
pub(crate) fn static_registry_auth_provider() -> codex_api::SharedAuthProvider {
Arc::new(StaticRegistryAuthProvider)
}
pub(crate) struct RelayTest {
registry: MockServer,
listener: TcpListener,
}
pub(crate) struct RelayConnection {
pub(crate) client: ExecServerClient,
captured_frames: Arc<Mutex<Vec<Vec<u8>>>>,
relay_task: AbortOnDropHandle<Result<()>>,
}
impl RelayTest {
pub(crate) async fn new() -> Result<Self> {
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(format!(
"/cloud/environment/{ENVIRONMENT_ID}/register"
)))
.and(header("authorization", format!("Bearer {REGISTRY_TOKEN}")))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": ENVIRONMENT_ID,
"url": format!("{rendezvous_url}/relay?role=environment"),
"security_profile": "noise_hybrid_ik_v1",
"executor_registration_id": EXECUTOR_REGISTRATION_ID,
})))
.mount(&registry)
.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(&registry)
.await;
Ok(Self { registry, listener })
}
pub(crate) fn config(&self) -> Result<RemoteEnvironmentConfig> {
Ok(RemoteEnvironmentConfig::new(
self.registry.uri(),
ENVIRONMENT_ID.to_string(),
static_registry_auth_provider(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)?)
}
pub(crate) async fn connect(&self) -> Result<RelayConnection> {
let rendezvous_url = format!("ws://{}", self.listener.local_addr()?);
let environment_websocket = accept_websocket(&self.listener, "environment").await?;
let executor_public_key = registered_executor_public_key(&self.registry).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,
http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
};
let client_task =
tokio::spawn(
async move { ExecServerClient::connect_noise_rendezvous(client_args).await },
);
let harness_websocket = accept_websocket(&self.listener, "harness").await?;
let captured_frames = Arc::new(Mutex::new(Vec::new()));
let relay_task = AbortOnDropHandle::new(tokio::spawn(proxy_relay_frames(
environment_websocket,
harness_websocket,
Arc::clone(&captured_frames),
)));
let client = timeout(TEST_TIMEOUT, client_task)
.await
.context("Noise harness client should connect")???;
Ok(RelayConnection {
client,
captured_frames,
relay_task,
})
}
}
impl RelayConnection {
pub(crate) fn assert_encrypted(&self) -> Result<()> {
assert_relay_data_is_encrypted(&self.captured_frames)
}
pub(crate) async fn close(self) {
drop(self.client);
self.relay_task.abort();
let _ = self.relay_task.await;
}
}
pub(crate) async fn accept_websocket(
listener: &TcpListener,
role: &str,
) -> Result<WebSocketStream<TcpStream>> {
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)
}
pub(crate) async fn registered_executor_public_key(
registry: &MockServer,
) -> Result<NoiseChannelPublicKey> {
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)
}
pub(crate) async fn proxy_relay_frames(
mut environment: WebSocketStream<TcpStream>,
mut harness: WebSocketStream<TcpStream>,
captured_frames: Arc<Mutex<Vec<Vec<u8>>>>,
) -> 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(())
}
fn capture_binary_frame(captured_frames: &Mutex<Vec<Vec<u8>>>, message: &Message) {
if let Message::Binary(bytes) = message {
captured_frames
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(bytes.to_vec());
}
}
pub(crate) fn assert_relay_data_is_encrypted(captured_frames: &Mutex<Vec<Vec<u8>>>) -> 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!(
data_frames >= 4,
"expected encrypted request and response frames"
);
Ok(())
}
codex-rs/exec-server/tests/relay.rs
+53 -224
View File
@@ -4,33 +4,28 @@ mod common;
#[path = "relay/version_skew.rs"]
mod version_skew;
#[path = "../src/proto/codex.exec_server.relay.v1.rs"]
mod relay_proto;
#[path = "common/relay.rs"]
mod relay_support;
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 base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_api::AuthProvider;
use codex_exec_server::EnvironmentConnectionState;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::EnvironmentReadyInfo;
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::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;
@@ -42,15 +37,17 @@ use codex_http_client::cache_system_proxy_route_for_test;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
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::RelayMessageFrame;
use relay_proto::relay_message_frame;
use relay_support::ENVIRONMENT_ID;
use relay_support::EXECUTOR_REGISTRATION_ID;
use relay_support::HARNESS_KEY_AUTHORIZATION;
use relay_support::RelayTest;
use relay_support::TEST_TIMEOUT;
use relay_support::accept_websocket;
use relay_support::proxy_relay_frames;
use relay_support::registered_executor_public_key;
use relay_support::static_registry_auth_provider;
use tempfile::TempDir;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
@@ -60,39 +57,13 @@ use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use tokio_util::task::AbortOnDropHandle;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
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 TEST_TIMEOUT: Duration = Duration::from_secs(30);
#[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() -> codex_api::SharedAuthProvider {
Arc::new(StaticRegistryAuthProvider)
}
struct FreshBundleNoiseConnectProvider {
websocket_url: String,
executor_public_key: NoiseChannelPublicKey,
@@ -123,6 +94,7 @@ impl NoiseRendezvousConnectProvider for FreshBundleNoiseConnectProvider {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
#[serial_test::serial]
async fn pending_noise_environment_connects_and_reconnects_after_ready_report() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let rendezvous_address = listener.local_addr()?;
@@ -199,9 +171,15 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report()
static_registry_auth_provider(),
http_client_factory.clone(),
)?;
let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment(
config,
runtime_paths,
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let remote_environment = AbortOnDropHandle::new(tokio::spawn(
codex_exec_server::run_remote_environment_until_shutdown(
config,
runtime_paths,
async move {
let _ = shutdown_rx.await;
},
),
));
let environment_websocket = accept_websocket(&listener, "environment").await?;
@@ -316,9 +294,9 @@ async fn pending_noise_environment_connects_and_reconnects_after_ready_report()
registry.verify().await;
second_relay.abort();
remote_environment.abort();
let _ = second_relay.await;
let _ = remote_environment.await;
let _ = shutdown_tx.send(());
timeout(TEST_TIMEOUT, remote_environment).await???;
Ok(())
}
@@ -331,95 +309,37 @@ async fn next_connection_state(
#[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;
Mock::given(method("POST"))
.and(path(format!(
"/cloud/environment/{ENVIRONMENT_ID}/register"
)))
.and(header("authorization", format!("Bearer {REGISTRY_TOKEN}")))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"environment_id": ENVIRONMENT_ID,
"url": format!("{rendezvous_url}/relay?role=environment"),
"security_profile": "noise_hybrid_ik_v1",
"executor_registration_id": EXECUTOR_REGISTRATION_ID,
})))
.mount(&registry)
.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(&registry)
.await;
let relay = RelayTest::new().await?;
let (codex_exe, codex_linux_sandbox_exe) = common::current_test_binary_helper_paths()?;
let runtime_paths = ExecServerRuntimePaths::new(codex_exe, codex_linux_sandbox_exe)?;
let config = RemoteEnvironmentConfig::new(
registry.uri(),
ENVIRONMENT_ID.to_string(),
static_registry_auth_provider(),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)?;
let remote_environment = tokio::spawn(codex_exec_server::run_remote_environment(
config,
runtime_paths,
));
let environment_websocket = accept_websocket(&listener, "environment").await?;
let executor_public_key = registered_executor_public_key(&registry).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,
http_client_factory: codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let remote_environment = AbortOnDropHandle::new(tokio::spawn(
codex_exec_server::run_remote_environment_until_shutdown(
relay.config()?,
runtime_paths,
async move {
let _ = shutdown_rx.await;
},
),
};
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("Noise harness client should connect")???;
let connection = relay.connect().await?;
let client = &connection.client;
let response = client
.exec(ExecParams {
process_id: ProcessId::from("proc-1"),
argv: vec!["true".to_string()],
cwd: PathUri::from_host_native_path(std::env::current_dir()?)?,
env_policy: None,
env: HashMap::new(),
tty: false,
pipe_stdin: false,
arg0: None,
sandbox: None,
enforce_managed_network: false,
managed_network: None,
network_proxy: None,
})
.await?;
let exec_params = ExecParams {
process_id: ProcessId::from("proc-1"),
argv: vec!["true".to_string()],
cwd: PathUri::from_host_native_path(std::env::current_dir()?)?,
env_policy: None,
env: HashMap::new(),
tty: false,
pipe_stdin: false,
arg0: None,
sandbox: None,
enforce_managed_network: false,
managed_network: None,
network_proxy: None,
};
let response = client.exec(exec_params).await?;
assert_eq!(
response,
ExecResponse {
@@ -442,100 +362,9 @@ async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> {
STANDARD.decode(read_response.data_base64)?,
large_file_contents
);
assert_relay_data_is_encrypted(&captured_frames)?;
drop(client);
relay_task.abort();
remote_environment.abort();
let _ = relay_task.await;
let _ = remote_environment.await;
Ok(())
}
async fn accept_websocket(
listener: &TcpListener,
role: &str,
) -> Result<WebSocketStream<TcpStream>> {
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 registered_executor_public_key(registry: &MockServer) -> Result<NoiseChannelPublicKey> {
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 proxy_relay_frames(
mut environment: WebSocketStream<TcpStream>,
mut harness: WebSocketStream<TcpStream>,
captured_frames: Arc<Mutex<Vec<Vec<u8>>>>,
) -> 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(())
}
fn capture_binary_frame(captured_frames: &Mutex<Vec<Vec<u8>>>, message: &Message) {
if let Message::Binary(bytes) = message {
captured_frames
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.push(bytes.to_vec());
}
}
fn assert_relay_data_is_encrypted(captured_frames: &Mutex<Vec<Vec<u8>>>) -> 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!(
data_frames >= 4,
"expected encrypted request and response frames"
);
connection.assert_encrypted()?;
connection.close().await;
let _ = shutdown_tx.send(());
timeout(TEST_TIMEOUT, remote_environment).await???;
Ok(())
}
codex-rs/exec-server/tests/relay/version_skew.rs
+8 -8
View File
@@ -31,14 +31,14 @@ use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
use super::ENVIRONMENT_ID;
use super::EXECUTOR_REGISTRATION_ID;
use super::HARNESS_KEY_AUTHORIZATION;
use super::REGISTRY_TOKEN;
use super::accept_websocket;
use super::assert_relay_data_is_encrypted;
use super::proxy_relay_frames;
use super::registered_executor_public_key;
use crate::relay_support::ENVIRONMENT_ID;
use crate::relay_support::EXECUTOR_REGISTRATION_ID;
use crate::relay_support::HARNESS_KEY_AUTHORIZATION;
use crate::relay_support::REGISTRY_TOKEN;
use crate::relay_support::accept_websocket;
use crate::relay_support::assert_relay_data_is_encrypted;
use crate::relay_support::proxy_relay_frames;
use crate::relay_support::registered_executor_public_key;
const RELEASED_CODEX_ENV_VAR: &str = "CODEX_TEST_RELEASED_CODEX";
const CURRENT_CODEX_ENV_VAR: &str = "CODEX_TEST_CURRENT_CODEX";