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

## 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
@@ -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,
}
}
@@ -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.
@@ -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;
@@ -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;
@@ -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(())
}
+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() {
@@ -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?;