mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
viyatb-oai
6e50b22e55
## Why The transport in [openai/codex#26242](https://github.com/openai/codex/pull/26242) needs to be used by every remote orchestrator-to-executor connection before JSON-RPC traffic starts. ## Changes - Generates one executor Noise identity when remote exec-server starts and registers its public key. - Creates a harness identity for each physical remote environment connection. - Fetches a fresh registry bundle before connecting and validates the authenticated harness key before completing the executor handshake. - Multiplexes encrypted logical streams over the existing executor WebSocket. - Adds bounded stream, handshake-failure, and reassembly state. - Adds safe lifecycle diagnostics without logging keys, authorizations, plaintext, or ciphertext. - Covers reconnects, replay rejection, validation failure, framing limits, and encrypted JSON-RPC tool traffic. ## Stack 1. [openai/codex#26242](https://github.com/openai/codex/pull/26242): Noise channel and relay transport 2. **[openai/codex#26245](https://github.com/openai/codex/pull/26245)**: remote registration and runtime activation ## Verification - `just test -p codex-exec-server` - `just fix -p codex-exec-server` - `just bazel-lock-check` - `cargo shear` --------- Co-authored-by: Codex <noreply@openai.com>
1241 lines
50 KiB
Rust
1241 lines
50 KiB
Rust
use std::collections::HashMap;
|
|
use std::time::Duration;
|
|
|
|
use codex_app_server_protocol::JSONRPCMessage;
|
|
use futures::Sink;
|
|
use futures::SinkExt;
|
|
use futures::Stream;
|
|
use futures::StreamExt;
|
|
use prost::Message as ProstMessage;
|
|
use tokio::io::AsyncRead;
|
|
use tokio::io::AsyncWrite;
|
|
use tokio::sync::mpsc;
|
|
use tokio::sync::watch;
|
|
use tokio::task::JoinSet;
|
|
use tokio::time::timeout;
|
|
use tokio_tungstenite::WebSocketStream;
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
use tracing::debug;
|
|
use tracing::info;
|
|
use tracing::warn;
|
|
use uuid::Uuid;
|
|
|
|
use crate::ExecServerError;
|
|
use crate::connection::CHANNEL_CAPACITY;
|
|
use crate::connection::JsonRpcConnection;
|
|
use crate::connection::JsonRpcConnectionEvent;
|
|
use crate::connection::JsonRpcTransport;
|
|
use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL;
|
|
use crate::noise_channel::NoiseChannelIdentity;
|
|
use crate::noise_channel::NoiseChannelPublicKey;
|
|
use crate::noise_channel::PendingResponderHandshake;
|
|
use crate::noise_channel::noise_channel_prologue;
|
|
use crate::noise_relay::NOISE_RELAY_RESET_REASON;
|
|
use crate::noise_relay::executor_stream::ClosedNoiseVirtualStream;
|
|
use crate::noise_relay::executor_stream::NoiseVirtualStream;
|
|
use crate::noise_relay::executor_stream::spawn_noise_virtual_stream;
|
|
use crate::relay_proto::RelayData;
|
|
use crate::relay_proto::RelayHandshake;
|
|
use crate::relay_proto::RelayMessageFrame;
|
|
use crate::relay_proto::RelayReset;
|
|
use crate::relay_proto::RelayResume;
|
|
use crate::relay_proto::relay_message_frame;
|
|
use crate::server::ConnectionProcessor;
|
|
|
|
const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
|
|
const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128;
|
|
const MAX_FAILED_NOISE_HANDSHAKES: usize = 8;
|
|
const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096;
|
|
const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32;
|
|
const HARNESS_KEY_VALIDATION_TIMEOUT: Duration = Duration::from_secs(10);
|
|
|
|
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
|
|
pub(crate) enum RelayFrameBodyKind {
|
|
Data,
|
|
Ack,
|
|
Resume,
|
|
Reset,
|
|
Heartbeat,
|
|
Handshake,
|
|
}
|
|
|
|
impl RelayMessageFrame {
|
|
pub(crate) fn data(stream_id: String, seq: u32, payload: Vec<u8>) -> Self {
|
|
Self {
|
|
version: RELAY_MESSAGE_FRAME_VERSION,
|
|
stream_id,
|
|
ack: 0,
|
|
ack_bits: 0,
|
|
body: Some(relay_message_frame::Body::Data(RelayData {
|
|
seq,
|
|
segment_index: 0,
|
|
segment_count: 1,
|
|
payload,
|
|
})),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn resume(stream_id: String) -> Self {
|
|
Self {
|
|
version: RELAY_MESSAGE_FRAME_VERSION,
|
|
stream_id,
|
|
ack: 0,
|
|
ack_bits: 0,
|
|
body: Some(relay_message_frame::Body::Resume(RelayResume {
|
|
next_seq: 0,
|
|
})),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn handshake(stream_id: String, payload: Vec<u8>) -> Self {
|
|
Self {
|
|
version: RELAY_MESSAGE_FRAME_VERSION,
|
|
stream_id,
|
|
ack: 0,
|
|
ack_bits: 0,
|
|
body: Some(relay_message_frame::Body::Handshake(RelayHandshake {
|
|
payload,
|
|
})),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn reset(stream_id: String, reason: String) -> Self {
|
|
Self {
|
|
version: RELAY_MESSAGE_FRAME_VERSION,
|
|
stream_id,
|
|
ack: 0,
|
|
ack_bits: 0,
|
|
body: Some(relay_message_frame::Body::Reset(RelayReset { reason })),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn validate(&self) -> Result<RelayFrameBodyKind, ExecServerError> {
|
|
if self.version != RELAY_MESSAGE_FRAME_VERSION {
|
|
return Err(ExecServerError::Protocol(format!(
|
|
"unsupported relay message frame version {}",
|
|
self.version
|
|
)));
|
|
}
|
|
if self.stream_id.trim().is_empty() {
|
|
return Err(ExecServerError::Protocol(
|
|
"relay message frame is missing stream_id".to_string(),
|
|
));
|
|
}
|
|
match self.body.as_ref() {
|
|
Some(relay_message_frame::Body::Data(data)) => {
|
|
if data.segment_index != 0 || data.segment_count != 1 || data.payload.is_empty() {
|
|
return Err(ExecServerError::Protocol(
|
|
"relay data message frame is missing required fields".to_string(),
|
|
));
|
|
}
|
|
Ok(RelayFrameBodyKind::Data)
|
|
}
|
|
Some(relay_message_frame::Body::AckFrame(_)) => Ok(RelayFrameBodyKind::Ack),
|
|
Some(relay_message_frame::Body::Resume(_)) => Ok(RelayFrameBodyKind::Resume),
|
|
Some(relay_message_frame::Body::Reset(reset)) => {
|
|
if reset.reason.is_empty() {
|
|
return Err(ExecServerError::Protocol(
|
|
"relay reset message frame is missing reason".to_string(),
|
|
));
|
|
}
|
|
Ok(RelayFrameBodyKind::Reset)
|
|
}
|
|
Some(relay_message_frame::Body::Heartbeat(_)) => Ok(RelayFrameBodyKind::Heartbeat),
|
|
Some(relay_message_frame::Body::Handshake(handshake)) => {
|
|
if handshake.payload.is_empty() {
|
|
return Err(ExecServerError::Protocol(
|
|
"relay handshake message frame is missing payload".to_string(),
|
|
));
|
|
}
|
|
Ok(RelayFrameBodyKind::Handshake)
|
|
}
|
|
None => Err(ExecServerError::Protocol(
|
|
"relay message frame is missing body".to_string(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn into_data(self) -> Result<RelayData, ExecServerError> {
|
|
let kind = self.validate()?;
|
|
if kind != RelayFrameBodyKind::Data {
|
|
return Err(ExecServerError::Protocol(
|
|
"expected relay data message frame".to_string(),
|
|
));
|
|
}
|
|
match self.body {
|
|
Some(relay_message_frame::Body::Data(data)) => Ok(data),
|
|
_ => Err(ExecServerError::Protocol(
|
|
"expected relay data message frame".to_string(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn into_jsonrpc_message(self) -> Result<JSONRPCMessage, ExecServerError> {
|
|
let payload = self.into_data()?.payload;
|
|
serde_json::from_slice(&payload).map_err(ExecServerError::Json)
|
|
}
|
|
|
|
pub(crate) fn into_handshake_payload(self) -> Result<Vec<u8>, ExecServerError> {
|
|
let kind = self.validate()?;
|
|
if kind != RelayFrameBodyKind::Handshake {
|
|
return Err(ExecServerError::Protocol(
|
|
"expected relay handshake message frame".to_string(),
|
|
));
|
|
}
|
|
match self.body {
|
|
Some(relay_message_frame::Body::Handshake(handshake)) => Ok(handshake.payload),
|
|
_ => Err(ExecServerError::Protocol(
|
|
"expected relay handshake message frame".to_string(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(crate) fn into_reset_reason(self) -> Option<String> {
|
|
match self.body {
|
|
Some(relay_message_frame::Body::Reset(reset)) if !reset.reason.is_empty() => {
|
|
Some(reset.reason)
|
|
}
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn encode_relay_message_frame(frame: &RelayMessageFrame) -> Vec<u8> {
|
|
frame.encode_to_vec()
|
|
}
|
|
|
|
pub(crate) fn decode_relay_message_frame(
|
|
payload: &[u8],
|
|
) -> Result<RelayMessageFrame, ExecServerError> {
|
|
RelayMessageFrame::decode(payload)
|
|
.map_err(|err| ExecServerError::Protocol(format!("invalid relay message frame: {err}")))
|
|
}
|
|
|
|
pub(crate) fn jsonrpc_payload(message: &JSONRPCMessage) -> Result<Vec<u8>, ExecServerError> {
|
|
serde_json::to_vec(message).map_err(ExecServerError::Json)
|
|
}
|
|
|
|
enum RelayEventSendError {
|
|
IncomingClosed,
|
|
WebSocketClosed,
|
|
}
|
|
|
|
async fn send_event_with_keepalive<T, E>(
|
|
websocket: &mut T,
|
|
keepalive: &mut tokio::time::Interval,
|
|
incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
|
|
event: JsonRpcConnectionEvent,
|
|
) -> Result<(), RelayEventSendError>
|
|
where
|
|
T: Sink<Message, Error = E> + Unpin,
|
|
{
|
|
let send = incoming_tx.send(event);
|
|
tokio::pin!(send);
|
|
loop {
|
|
tokio::select! {
|
|
result = &mut send => {
|
|
return result.map_err(|_| RelayEventSendError::IncomingClosed);
|
|
}
|
|
_ = keepalive.tick() => {
|
|
websocket
|
|
.send(Message::Ping(Vec::new().into()))
|
|
.await
|
|
.map_err(|_| RelayEventSendError::WebSocketClosed)?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn harness_connection_from_websocket<T, E>(
|
|
stream: T,
|
|
connection_label: String,
|
|
) -> JsonRpcConnection
|
|
where
|
|
T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,
|
|
E: std::fmt::Display + Send + 'static,
|
|
{
|
|
let stream_id = Uuid::new_v4().to_string();
|
|
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 websocket_task = tokio::spawn(async move {
|
|
let mut websocket = stream;
|
|
let reader_label = connection_label;
|
|
let reader_stream_id = stream_id.clone();
|
|
let resume = RelayMessageFrame::resume(stream_id.clone());
|
|
if websocket
|
|
.send(Message::Binary(encode_relay_message_frame(&resume).into()))
|
|
.await
|
|
.is_err()
|
|
{
|
|
let _ = disconnected_tx.send(true);
|
|
return;
|
|
}
|
|
|
|
let mut keepalive = tokio::time::interval_at(
|
|
tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL,
|
|
WEBSOCKET_KEEPALIVE_INTERVAL,
|
|
);
|
|
keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
let mut next_seq = 0u32;
|
|
loop {
|
|
tokio::select! {
|
|
maybe_message = outgoing_rx.recv() => {
|
|
let Some(message) = maybe_message else {
|
|
break;
|
|
};
|
|
let payload = match jsonrpc_payload(&message) {
|
|
Ok(payload) => payload,
|
|
Err(err) => {
|
|
warn!("failed to serialize JSON-RPC payload for relay transport: {err}");
|
|
break;
|
|
}
|
|
};
|
|
let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload);
|
|
next_seq = next_seq.wrapping_add(1);
|
|
if websocket
|
|
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
|
|
.await
|
|
.is_err()
|
|
{
|
|
let _ = disconnected_tx.send(true);
|
|
break;
|
|
}
|
|
}
|
|
_ = keepalive.tick() => {
|
|
if websocket.send(Message::Ping(Vec::new().into())).await.is_err() {
|
|
let _ = disconnected_tx.send(true);
|
|
break;
|
|
}
|
|
}
|
|
incoming_message = websocket.next() => {
|
|
match incoming_message {
|
|
Some(Ok(Message::Binary(payload))) => {
|
|
let frame = match decode_relay_message_frame(payload.as_ref()) {
|
|
Ok(frame) => frame,
|
|
Err(err) => {
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::MalformedMessage {
|
|
reason: format!(
|
|
"failed to parse relay message frame from {reader_label}: {err}"
|
|
),
|
|
})
|
|
.await;
|
|
continue;
|
|
}
|
|
};
|
|
if frame.stream_id != reader_stream_id {
|
|
continue;
|
|
}
|
|
let kind = match frame.validate() {
|
|
Ok(kind) => kind,
|
|
Err(err) => {
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::MalformedMessage {
|
|
reason: err.to_string(),
|
|
})
|
|
.await;
|
|
continue;
|
|
}
|
|
};
|
|
match kind {
|
|
RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() {
|
|
Ok(message) => {
|
|
match send_event_with_keepalive(
|
|
&mut websocket,
|
|
&mut keepalive,
|
|
&incoming_tx,
|
|
JsonRpcConnectionEvent::Message(message),
|
|
)
|
|
.await
|
|
{
|
|
Ok(()) => {}
|
|
Err(RelayEventSendError::IncomingClosed) => break,
|
|
Err(RelayEventSendError::WebSocketClosed) => {
|
|
let _ = disconnected_tx.send(true);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
Err(err) => {
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::MalformedMessage {
|
|
reason: err.to_string(),
|
|
})
|
|
.await;
|
|
}
|
|
},
|
|
RelayFrameBodyKind::Reset => {
|
|
let _ = disconnected_tx.send(true);
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::Disconnected {
|
|
reason: frame.into_reset_reason(),
|
|
})
|
|
.await;
|
|
break;
|
|
}
|
|
RelayFrameBodyKind::Ack
|
|
| RelayFrameBodyKind::Resume
|
|
| RelayFrameBodyKind::Heartbeat
|
|
| RelayFrameBodyKind::Handshake => {}
|
|
}
|
|
}
|
|
Some(Ok(Message::Close(_))) | None => {
|
|
let _ = disconnected_tx.send(true);
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::Disconnected { reason: None })
|
|
.await;
|
|
break;
|
|
}
|
|
Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {}
|
|
Some(Ok(Message::Text(_))) => {
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::MalformedMessage {
|
|
reason: "relay exec-server transport expects binary protobuf frames"
|
|
.to_string(),
|
|
})
|
|
.await;
|
|
}
|
|
Some(Err(err)) => {
|
|
let _ = disconnected_tx.send(true);
|
|
let _ = incoming_tx
|
|
.send(JsonRpcConnectionEvent::Disconnected {
|
|
reason: Some(format!(
|
|
"failed to read relay websocket frame from {reader_label}: {err}"
|
|
)),
|
|
})
|
|
.await;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
JsonRpcConnection {
|
|
outgoing_tx,
|
|
incoming_rx,
|
|
disconnected_rx,
|
|
task_handles: vec![websocket_task],
|
|
transport: JsonRpcTransport::Plain,
|
|
}
|
|
}
|
|
|
|
/// Validates that a Noise-authenticated harness public key is authorized.
|
|
///
|
|
/// Implementations must consult an authority independent of rendezvous. The
|
|
/// exec-server invokes this after parsing the first IK message and before
|
|
/// completing the responder handshake.
|
|
pub(crate) trait HarnessKeyValidator: Send + Sync {
|
|
fn validate_harness_key(
|
|
&self,
|
|
harness_public_key: &NoiseChannelPublicKey,
|
|
authorization: &str,
|
|
) -> impl std::future::Future<Output = Result<(), ExecServerError>> + Send;
|
|
}
|
|
|
|
/// Serve authenticated virtual JSON-RPC streams over one executor websocket.
|
|
///
|
|
/// Parsing the first Noise message authenticates the harness key. Only a
|
|
/// successful registry check turns that pending handshake into a virtual stream.
|
|
#[tracing::instrument(
|
|
level = "debug",
|
|
skip_all,
|
|
fields(
|
|
noise_side = "executor",
|
|
environment_id = %environment_id,
|
|
executor_registration_id = %executor_registration_id,
|
|
)
|
|
)]
|
|
pub(crate) async fn run_multiplexed_environment<S, V>(
|
|
stream: WebSocketStream<S>,
|
|
processor: ConnectionProcessor,
|
|
environment_id: String,
|
|
executor_registration_id: String,
|
|
identity: NoiseChannelIdentity,
|
|
validator: V,
|
|
) where
|
|
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
|
V: HarnessKeyValidator + Clone + 'static,
|
|
{
|
|
let (mut websocket_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);
|
|
// Use a separate writer so this loop never waits on the channel it drains.
|
|
let mut physical_writer_task = tokio::spawn(async move {
|
|
let mut keepalive = tokio::time::interval_at(
|
|
tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL,
|
|
WEBSOCKET_KEEPALIVE_INTERVAL,
|
|
);
|
|
keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
loop {
|
|
let message = tokio::select! {
|
|
encoded = physical_outgoing_rx.recv() => {
|
|
let Some(encoded) = encoded else {
|
|
break;
|
|
};
|
|
Message::Binary(encoded.into())
|
|
}
|
|
_ = keepalive.tick() => Message::Ping(Vec::new().into()),
|
|
};
|
|
if let Err(error) = websocket_sink.send(message).await {
|
|
warn!("Noise multiplexed environment websocket write failed: {error}");
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
let mut streams: HashMap<String, NoiseVirtualStream> = HashMap::new();
|
|
let mut pending_handshakes: HashMap<String, PendingHandshake> = HashMap::new();
|
|
let mut validation_tasks: JoinSet<HarnessKeyValidationResult> = JoinSet::new();
|
|
let mut failed_handshakes = 0usize;
|
|
let mut next_validation_id = 0u64;
|
|
|
|
loop {
|
|
// Registry calls run separately so a slow check does not block the relay.
|
|
let frame = tokio::select! {
|
|
writer_result = &mut physical_writer_task => {
|
|
if let Err(error) = writer_result {
|
|
warn!("Noise multiplexed environment websocket writer failed: {error}");
|
|
}
|
|
break;
|
|
}
|
|
Some(closed_stream) = closed_stream_rx.recv() => {
|
|
// A stream ID may have been reused before this writer exits.
|
|
// Remove only the instance that sent the notification.
|
|
let is_current = streams
|
|
.get(&closed_stream.stream_id)
|
|
.is_some_and(|stream| stream.instance_id == closed_stream.instance_id);
|
|
if is_current {
|
|
streams.remove(&closed_stream.stream_id);
|
|
}
|
|
continue;
|
|
}
|
|
validation_result = validation_tasks.join_next(), if !validation_tasks.is_empty() => {
|
|
match validation_result {
|
|
Some(Ok(validation_result)) => {
|
|
// The stream ID may have been reused while validation ran.
|
|
let is_current = pending_handshakes
|
|
.get(&validation_result.stream_id)
|
|
.is_some_and(|pending| {
|
|
pending.validation_id == validation_result.validation_id
|
|
});
|
|
if !is_current {
|
|
continue;
|
|
}
|
|
let Some(pending) =
|
|
pending_handshakes.remove(&validation_result.stream_id)
|
|
else {
|
|
continue;
|
|
};
|
|
if validation_result.result.is_err() {
|
|
// Validator errors may contain authorization details.
|
|
warn!(
|
|
noise_event = "authorization",
|
|
noise_outcome = "error",
|
|
noise_reason = "authorization_failed",
|
|
"Noise harness authorization failed"
|
|
);
|
|
debug!(
|
|
stream_id = validation_result.stream_id,
|
|
"Noise harness authorization failure details"
|
|
);
|
|
send_reset(&physical_outgoing_tx, validation_result.stream_id);
|
|
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
|
|
warn!("closing Noise relay after repeated handshake failures");
|
|
break;
|
|
}
|
|
continue;
|
|
}
|
|
if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS {
|
|
warn!("Noise relay has too many active streams");
|
|
send_reset(&physical_outgoing_tx, validation_result.stream_id);
|
|
continue;
|
|
}
|
|
|
|
// This is the only point where the responder completes
|
|
// IK and exposes a JSON-RPC stream: Noise authenticated
|
|
// the harness key and the registry authorized it.
|
|
let (transport, response) = match pending.handshake.complete() {
|
|
Ok(completed) => completed,
|
|
Err(error) => {
|
|
warn!("failed to complete Noise relay handshake: {error}");
|
|
send_reset(&physical_outgoing_tx, validation_result.stream_id);
|
|
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
|
|
warn!("closing Noise relay after repeated handshake failures");
|
|
break;
|
|
}
|
|
continue;
|
|
}
|
|
};
|
|
let response = RelayMessageFrame::handshake(
|
|
validation_result.stream_id.clone(),
|
|
response,
|
|
);
|
|
// Do not leave a half-open stream if the handshake reply
|
|
// cannot be queued immediately.
|
|
if physical_outgoing_tx
|
|
.try_send(encode_relay_message_frame(&response))
|
|
.is_err()
|
|
{
|
|
break;
|
|
}
|
|
info!(
|
|
noise_event = "handshake",
|
|
noise_outcome = "ok",
|
|
"Noise executor handshake completed"
|
|
);
|
|
debug!(
|
|
stream_id = validation_result.stream_id,
|
|
active_streams = streams.len() + 1,
|
|
"Noise executor stream activated"
|
|
);
|
|
streams.insert(
|
|
validation_result.stream_id.clone(),
|
|
spawn_noise_virtual_stream(
|
|
validation_result.stream_id,
|
|
validation_result.validation_id,
|
|
processor.clone(),
|
|
physical_outgoing_tx.clone(),
|
|
closed_stream_tx.clone(),
|
|
transport,
|
|
),
|
|
);
|
|
}
|
|
Some(Err(error)) => {
|
|
warn!("Noise relay harness key validation task failed: {error}");
|
|
let stream_ids = pending_handshakes.keys().cloned().collect::<Vec<_>>();
|
|
pending_handshakes.clear();
|
|
for stream_id in stream_ids {
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
}
|
|
}
|
|
None => {}
|
|
}
|
|
continue;
|
|
}
|
|
incoming_message = websocket_stream.next() => match incoming_message {
|
|
Some(Ok(Message::Binary(payload))) => match decode_relay_message_frame(payload.as_ref()) {
|
|
Ok(frame) => frame,
|
|
Err(error) => {
|
|
warn!("dropping malformed Noise relay frame from harness: {error}");
|
|
continue;
|
|
}
|
|
},
|
|
Some(Ok(Message::Close(_))) | None => break,
|
|
Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue,
|
|
Some(Ok(Message::Text(_))) => {
|
|
warn!("dropping non-binary Noise relay frame from harness");
|
|
continue;
|
|
}
|
|
Some(Err(error)) => {
|
|
debug!("Noise multiplexed environment websocket read failed: {error}");
|
|
break;
|
|
}
|
|
}
|
|
};
|
|
|
|
let kind = match frame.validate() {
|
|
Ok(kind) => kind,
|
|
Err(error) => {
|
|
warn!("dropping invalid Noise relay frame: {error}");
|
|
continue;
|
|
}
|
|
};
|
|
let stream_id = frame.stream_id.clone();
|
|
match kind {
|
|
RelayFrameBodyKind::Handshake => {
|
|
// Reject duplicate or busy streams before paying for a hybrid
|
|
// handshake. Malformed attempts that reach cryptography are
|
|
// covered by the connection-wide failure budget below.
|
|
if streams.contains_key(&stream_id) {
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
continue;
|
|
}
|
|
// Removing pending state makes the in-flight validation result stale.
|
|
if pending_handshakes.remove(&stream_id).is_some() {
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
|
|
warn!("closing Noise relay after repeated handshake failures");
|
|
break;
|
|
}
|
|
continue;
|
|
}
|
|
if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS {
|
|
warn!("Noise relay has too many active streams");
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
continue;
|
|
}
|
|
if validation_tasks.len() >= MAX_PENDING_HANDSHAKE_VALIDATIONS {
|
|
warn!("Noise relay has too many pending harness key validations");
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
continue;
|
|
}
|
|
let prologue =
|
|
noise_channel_prologue(&environment_id, &executor_registration_id, &stream_id);
|
|
let request = match frame.into_handshake_payload() {
|
|
Ok(request) => request,
|
|
Err(error) => {
|
|
warn!("failed to read Noise relay handshake frame: {error}");
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
continue;
|
|
}
|
|
};
|
|
let mut pending =
|
|
match PendingResponderHandshake::read_request(&identity, &prologue, &request) {
|
|
Ok(pending) => pending,
|
|
Err(error) => {
|
|
warn!("failed to read Noise relay handshake request: {error}");
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
|
|
warn!("closing Noise relay after repeated handshake failures");
|
|
break;
|
|
}
|
|
continue;
|
|
}
|
|
};
|
|
|
|
// The authorization and authenticated harness key come from the
|
|
// same encrypted IK message and are validated together.
|
|
let authorization = match String::from_utf8(std::mem::take(&mut pending.payload)) {
|
|
Ok(authorization)
|
|
if authorization.len() <= MAX_HARNESS_KEY_AUTHORIZATION_BYTES =>
|
|
{
|
|
Some(authorization)
|
|
}
|
|
Ok(_) => {
|
|
warn!("Noise relay handshake authorization is too long");
|
|
None
|
|
}
|
|
Err(_) => {
|
|
warn!("Noise relay handshake authorization is not UTF-8");
|
|
None
|
|
}
|
|
};
|
|
let Some(authorization) = authorization else {
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
|
|
warn!("closing Noise relay after repeated handshake failures");
|
|
break;
|
|
}
|
|
continue;
|
|
};
|
|
let harness_public_key = pending.initiator_public_key.clone();
|
|
let validation_id = next_validation_id;
|
|
next_validation_id += 1;
|
|
pending_handshakes.insert(
|
|
stream_id.clone(),
|
|
PendingHandshake {
|
|
validation_id,
|
|
handshake: pending,
|
|
},
|
|
);
|
|
let validator = validator.clone();
|
|
|
|
// Failed validation leaves no transport state and sends only a
|
|
// generic reset.
|
|
validation_tasks.spawn(async move {
|
|
let result = match timeout(
|
|
HARNESS_KEY_VALIDATION_TIMEOUT,
|
|
validator.validate_harness_key(&harness_public_key, &authorization),
|
|
)
|
|
.await
|
|
{
|
|
Ok(result) => result,
|
|
Err(_) => Err(ExecServerError::Protocol(
|
|
"timed out validating Noise relay harness key".to_string(),
|
|
)),
|
|
};
|
|
HarnessKeyValidationResult {
|
|
stream_id,
|
|
validation_id,
|
|
result,
|
|
}
|
|
});
|
|
}
|
|
RelayFrameBodyKind::Data => {
|
|
// Removing pending state also makes any in-flight validation stale.
|
|
let Some(stream) = streams.get_mut(&stream_id) else {
|
|
let canceled_pending_handshake =
|
|
pending_handshakes.remove(&stream_id).is_some();
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
if canceled_pending_handshake
|
|
&& failed_handshake_budget_exhausted(&mut failed_handshakes)
|
|
{
|
|
warn!("closing Noise relay after repeated handshake failures");
|
|
break;
|
|
}
|
|
continue;
|
|
};
|
|
let data = match frame.into_data() {
|
|
Ok(data) => data,
|
|
Err(error) => {
|
|
warn!("dropping malformed Noise relay data frame: {error}");
|
|
streams.remove(&stream_id);
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
continue;
|
|
}
|
|
};
|
|
if let Err(error) = stream.receive_data(data) {
|
|
warn!("failed to process Noise relay payload: {error}");
|
|
streams.remove(&stream_id);
|
|
send_reset(&physical_outgoing_tx, stream_id);
|
|
}
|
|
}
|
|
RelayFrameBodyKind::Reset => {
|
|
pending_handshakes.remove(&stream_id);
|
|
if let Some(stream) = streams.remove(&stream_id) {
|
|
// The reset reason is unauthenticated, so do not log it.
|
|
stream.disconnect(/*reason*/ None);
|
|
}
|
|
}
|
|
RelayFrameBodyKind::Ack
|
|
| RelayFrameBodyKind::Resume
|
|
| RelayFrameBodyKind::Heartbeat => {}
|
|
}
|
|
}
|
|
|
|
for (_stream_id, stream) in streams {
|
|
stream.disconnect(/*reason*/ None);
|
|
}
|
|
// Dropping the JoinSet aborts any registry checks still running.
|
|
if !physical_writer_task.is_finished() {
|
|
physical_writer_task.abort();
|
|
let _ = physical_writer_task.await;
|
|
}
|
|
}
|
|
|
|
/// Charge one failed authenticated-channel attempt to this physical relay.
|
|
///
|
|
/// Closing after a small fixed budget prevents a peer that has not been
|
|
/// authorized from triggering unbounded hybrid handshakes or registry checks.
|
|
fn failed_handshake_budget_exhausted(failed_handshakes: &mut usize) -> bool {
|
|
*failed_handshakes += 1;
|
|
*failed_handshakes >= MAX_FAILED_NOISE_HANDSHAKES
|
|
}
|
|
|
|
/// Responder state held while registry authorization is pending.
|
|
struct PendingHandshake {
|
|
validation_id: u64,
|
|
handshake: PendingResponderHandshake,
|
|
}
|
|
|
|
/// `validation_id` prevents an old check from completing a reused `stream_id`.
|
|
struct HarnessKeyValidationResult {
|
|
stream_id: String,
|
|
validation_id: u64,
|
|
result: Result<(), ExecServerError>,
|
|
}
|
|
|
|
/// Queue a best-effort reset without blocking the shared websocket loop.
|
|
/// Reset reasons are relay control data and are not treated as trusted text.
|
|
fn send_reset(physical_outgoing_tx: &mpsc::Sender<Vec<u8>>, stream_id: String) {
|
|
let reset = RelayMessageFrame::reset(stream_id, NOISE_RELAY_RESET_REASON.to_string());
|
|
let _ = physical_outgoing_tx.try_send(encode_relay_message_frame(&reset));
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "relay_noise_tests.rs"]
|
|
mod noise_tests;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::pin::Pin;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::AtomicBool;
|
|
use std::sync::atomic::Ordering;
|
|
use std::task::Context;
|
|
use std::task::Poll;
|
|
use std::time::Duration;
|
|
|
|
use codex_app_server_protocol::JSONRPCRequest;
|
|
use codex_app_server_protocol::RequestId;
|
|
use futures::Sink;
|
|
use futures::Stream;
|
|
use futures::channel::mpsc as futures_mpsc;
|
|
use futures::task::AtomicWaker;
|
|
use pretty_assertions::assert_eq;
|
|
use tokio::net::TcpListener;
|
|
use tokio::time::timeout;
|
|
use tokio_tungstenite::WebSocketStream;
|
|
use tokio_tungstenite::accept_async;
|
|
use tokio_tungstenite::connect_async;
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn harness_connection_sends_keepalive_and_receives_relay_data() -> anyhow::Result<()> {
|
|
let (client_websocket, mut server_websocket) = websocket_pair().await?;
|
|
let mut connection =
|
|
harness_connection_from_websocket(client_websocket, "test".to_string());
|
|
let stream_id = read_resume_stream_id(&mut server_websocket).await?;
|
|
read_keepalive_ping(&mut server_websocket).await?;
|
|
server_websocket
|
|
.send(Message::Pong(b"keepalive".to_vec().into()))
|
|
.await?;
|
|
let message = test_jsonrpc_message();
|
|
|
|
server_websocket
|
|
.send(Message::Binary(
|
|
encode_relay_message_frame(&RelayMessageFrame::data(
|
|
stream_id,
|
|
/*seq*/ 0,
|
|
jsonrpc_payload(&message)?,
|
|
))
|
|
.into(),
|
|
))
|
|
.await?;
|
|
assert!(matches!(
|
|
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
|
|
Some(JsonRpcConnectionEvent::Message(actual)) if actual == message
|
|
));
|
|
|
|
drop(connection);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn multiplexed_environment_sends_keepalive() -> anyhow::Result<()> {
|
|
let (client_websocket, mut server_websocket) = websocket_pair().await?;
|
|
let runtime_paths = crate::ExecServerRuntimePaths::new(
|
|
std::env::current_exe()?,
|
|
/*codex_linux_sandbox_exe*/ None,
|
|
)
|
|
.map_err(anyhow::Error::from)?;
|
|
let environment_task = tokio::spawn(run_multiplexed_environment(
|
|
client_websocket,
|
|
ConnectionProcessor::new(runtime_paths),
|
|
"test-environment".to_string(),
|
|
"test-registration".to_string(),
|
|
NoiseChannelIdentity::generate()?,
|
|
AllowHarnessKeyValidator,
|
|
));
|
|
|
|
read_keepalive_ping(&mut server_websocket).await?;
|
|
|
|
environment_task.abort();
|
|
let _ = environment_task.await;
|
|
Ok(())
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct AllowHarnessKeyValidator;
|
|
|
|
impl HarnessKeyValidator for AllowHarnessKeyValidator {
|
|
async fn validate_harness_key(
|
|
&self,
|
|
_harness_public_key: &NoiseChannelPublicKey,
|
|
_authorization: &str,
|
|
) -> Result<(), ExecServerError> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_event_with_keepalive_pings_while_incoming_queue_is_full() -> anyhow::Result<()> {
|
|
let (mut websocket, _control, mut outbound_rx) =
|
|
ControlledWebSocket::new(/*write_ready*/ true);
|
|
let (incoming_tx, mut incoming_rx) = mpsc::channel(/*buffer*/ 1);
|
|
let message = test_jsonrpc_message();
|
|
let expected_message = message.clone();
|
|
incoming_tx
|
|
.send(JsonRpcConnectionEvent::MalformedMessage {
|
|
reason: "first".to_string(),
|
|
})
|
|
.await?;
|
|
let mut keepalive = tokio::time::interval_at(
|
|
tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL,
|
|
WEBSOCKET_KEEPALIVE_INTERVAL,
|
|
);
|
|
keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
|
|
let send_task = tokio::spawn(async move {
|
|
send_event_with_keepalive(
|
|
&mut websocket,
|
|
&mut keepalive,
|
|
&incoming_tx,
|
|
JsonRpcConnectionEvent::Message(message),
|
|
)
|
|
.await
|
|
});
|
|
|
|
assert!(matches!(
|
|
timeout(Duration::from_secs(1), outbound_rx.next()).await?,
|
|
Some(Message::Ping(_))
|
|
));
|
|
assert!(matches!(
|
|
incoming_rx.recv().await,
|
|
Some(JsonRpcConnectionEvent::MalformedMessage { reason }) if reason == "first"
|
|
));
|
|
assert!(matches!(
|
|
timeout(Duration::from_secs(1), send_task).await??,
|
|
Ok(())
|
|
));
|
|
assert!(matches!(
|
|
incoming_rx.recv().await,
|
|
Some(JsonRpcConnectionEvent::Message(actual)) if actual == expected_message
|
|
));
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn harness_connection_reports_text_frames_as_malformed() -> anyhow::Result<()> {
|
|
let (client_websocket, mut server_websocket) = websocket_pair().await?;
|
|
let mut connection =
|
|
harness_connection_from_websocket(client_websocket, "test".to_string());
|
|
|
|
read_resume_stream_id(&mut server_websocket).await?;
|
|
server_websocket.send(Message::Text("nope".into())).await?;
|
|
assert!(matches!(
|
|
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
|
|
Some(JsonRpcConnectionEvent::MalformedMessage { reason })
|
|
if reason == "relay exec-server transport expects binary protobuf frames"
|
|
));
|
|
|
|
drop(connection);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn harness_connection_reports_server_close() -> anyhow::Result<()> {
|
|
let (client_websocket, mut server_websocket) = websocket_pair().await?;
|
|
let mut connection =
|
|
harness_connection_from_websocket(client_websocket, "test".to_string());
|
|
|
|
read_resume_stream_id(&mut server_websocket).await?;
|
|
server_websocket.close(None).await?;
|
|
assert!(matches!(
|
|
timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?,
|
|
Some(JsonRpcConnectionEvent::Disconnected { reason: None })
|
|
));
|
|
|
|
drop(connection);
|
|
Ok(())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn harness_connection_keeps_outbound_frame_while_send_is_backpressured()
|
|
-> anyhow::Result<()> {
|
|
let (websocket, control, mut outbound_rx) =
|
|
ControlledWebSocket::new(/*write_ready*/ true);
|
|
let mut connection = harness_connection_from_websocket(websocket, "test".to_string());
|
|
let Message::Binary(resume_payload) = timeout(Duration::from_secs(1), outbound_rx.next())
|
|
.await?
|
|
.expect("resume frame")
|
|
else {
|
|
anyhow::bail!("expected relay resume frame");
|
|
};
|
|
let stream_id = decode_relay_message_frame(resume_payload.as_ref())?.stream_id;
|
|
let message = test_jsonrpc_message();
|
|
|
|
control.set_write_blocked();
|
|
connection.outgoing_tx.send(message.clone()).await?;
|
|
control.wait_for_blocked_write().await?;
|
|
control.send_inbound(Message::Pong(b"check".to_vec().into()))?;
|
|
assert!(
|
|
timeout(Duration::from_millis(50), connection.incoming_rx.recv())
|
|
.await
|
|
.is_err()
|
|
);
|
|
|
|
control.set_write_ready();
|
|
let Message::Binary(data_payload) = timeout(Duration::from_secs(1), outbound_rx.next())
|
|
.await?
|
|
.expect("data frame")
|
|
else {
|
|
anyhow::bail!("expected relay data frame");
|
|
};
|
|
let frame = decode_relay_message_frame(data_payload.as_ref())?;
|
|
assert_eq!(frame.stream_id, stream_id);
|
|
assert_eq!(frame.into_jsonrpc_message()?, message);
|
|
drop(connection);
|
|
Ok(())
|
|
}
|
|
|
|
async fn websocket_pair() -> anyhow::Result<(
|
|
WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
|
|
WebSocketStream<tokio::net::TcpStream>,
|
|
)> {
|
|
let listener = TcpListener::bind("127.0.0.1:0").await?;
|
|
let websocket_url = format!("ws://{}", listener.local_addr()?);
|
|
let server_task = tokio::spawn(async move {
|
|
let (stream, _) = listener.accept().await?;
|
|
accept_async(stream).await.map_err(anyhow::Error::from)
|
|
});
|
|
let (client_websocket, _) = connect_async(websocket_url).await?;
|
|
let server_websocket = server_task.await??;
|
|
Ok((client_websocket, server_websocket))
|
|
}
|
|
|
|
async fn read_resume_stream_id(
|
|
websocket: &mut WebSocketStream<tokio::net::TcpStream>,
|
|
) -> anyhow::Result<String> {
|
|
let message = timeout(Duration::from_secs(1), websocket.next())
|
|
.await?
|
|
.expect("websocket should stay open")?;
|
|
let Message::Binary(payload) = message else {
|
|
anyhow::bail!("expected relay resume frame, got {message:?}");
|
|
};
|
|
let frame = decode_relay_message_frame(payload.as_ref())?;
|
|
assert_eq!(frame.validate()?, RelayFrameBodyKind::Resume);
|
|
Ok(frame.stream_id)
|
|
}
|
|
|
|
async fn read_keepalive_ping(
|
|
websocket: &mut WebSocketStream<tokio::net::TcpStream>,
|
|
) -> anyhow::Result<()> {
|
|
loop {
|
|
let Some(message) = timeout(Duration::from_secs(1), websocket.next()).await? else {
|
|
anyhow::bail!("websocket closed before keepalive ping");
|
|
};
|
|
match message? {
|
|
Message::Ping(_) => return Ok(()),
|
|
Message::Binary(_) | Message::Text(_) | Message::Pong(_) | Message::Frame(_) => {}
|
|
Message::Close(_) => anyhow::bail!("websocket closed before keepalive ping"),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn test_jsonrpc_message() -> JSONRPCMessage {
|
|
JSONRPCMessage::Request(JSONRPCRequest {
|
|
id: RequestId::Integer(1),
|
|
method: "test".to_string(),
|
|
params: None,
|
|
trace: None,
|
|
})
|
|
}
|
|
|
|
struct ControlledWebSocket {
|
|
inbound_rx: futures_mpsc::UnboundedReceiver<Result<Message, std::convert::Infallible>>,
|
|
outbound_tx: futures_mpsc::UnboundedSender<Message>,
|
|
write_ready: Arc<AtomicBool>,
|
|
write_blocked: Arc<AtomicBool>,
|
|
write_blocked_waker: Arc<AtomicWaker>,
|
|
write_waker: Arc<AtomicWaker>,
|
|
}
|
|
|
|
struct ControlledWebSocketHandle {
|
|
inbound_tx: futures_mpsc::UnboundedSender<Result<Message, std::convert::Infallible>>,
|
|
write_ready: Arc<AtomicBool>,
|
|
write_blocked: Arc<AtomicBool>,
|
|
write_blocked_waker: Arc<AtomicWaker>,
|
|
write_waker: Arc<AtomicWaker>,
|
|
}
|
|
|
|
impl ControlledWebSocket {
|
|
fn new(
|
|
write_ready: bool,
|
|
) -> (
|
|
Self,
|
|
ControlledWebSocketHandle,
|
|
futures_mpsc::UnboundedReceiver<Message>,
|
|
) {
|
|
let (inbound_tx, inbound_rx) = futures_mpsc::unbounded();
|
|
let (outbound_tx, outbound_rx) = futures_mpsc::unbounded();
|
|
let write_ready = Arc::new(AtomicBool::new(write_ready));
|
|
let write_blocked = Arc::new(AtomicBool::new(false));
|
|
let write_blocked_waker = Arc::new(AtomicWaker::new());
|
|
let write_waker = Arc::new(AtomicWaker::new());
|
|
(
|
|
Self {
|
|
inbound_rx,
|
|
outbound_tx,
|
|
write_ready: Arc::clone(&write_ready),
|
|
write_blocked: Arc::clone(&write_blocked),
|
|
write_blocked_waker: Arc::clone(&write_blocked_waker),
|
|
write_waker: Arc::clone(&write_waker),
|
|
},
|
|
ControlledWebSocketHandle {
|
|
inbound_tx,
|
|
write_ready,
|
|
write_blocked,
|
|
write_blocked_waker,
|
|
write_waker,
|
|
},
|
|
outbound_rx,
|
|
)
|
|
}
|
|
}
|
|
|
|
impl ControlledWebSocketHandle {
|
|
fn send_inbound(&self, message: Message) -> anyhow::Result<()> {
|
|
self.inbound_tx
|
|
.unbounded_send(Ok(message))
|
|
.map_err(anyhow::Error::from)
|
|
}
|
|
|
|
fn set_write_blocked(&self) {
|
|
self.write_ready.store(false, Ordering::Release);
|
|
}
|
|
|
|
fn set_write_ready(&self) {
|
|
self.write_ready.store(true, Ordering::Release);
|
|
self.write_waker.wake();
|
|
}
|
|
|
|
async fn wait_for_blocked_write(&self) -> anyhow::Result<()> {
|
|
timeout(
|
|
Duration::from_secs(1),
|
|
futures::future::poll_fn(|cx| {
|
|
if self.write_blocked.load(Ordering::Acquire) {
|
|
Poll::Ready(())
|
|
} else {
|
|
self.write_blocked_waker.register(cx.waker());
|
|
Poll::Pending
|
|
}
|
|
}),
|
|
)
|
|
.await?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Sink<Message> for ControlledWebSocket {
|
|
type Error = std::convert::Infallible;
|
|
|
|
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
if self.write_ready.load(Ordering::Acquire) {
|
|
Poll::Ready(Ok(()))
|
|
} else {
|
|
self.write_blocked.store(true, Ordering::Release);
|
|
self.write_blocked_waker.wake();
|
|
self.write_waker.register(cx.waker());
|
|
Poll::Pending
|
|
}
|
|
}
|
|
|
|
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
|
|
self.outbound_tx
|
|
.unbounded_send(item)
|
|
.expect("test outbound receiver should stay open");
|
|
Ok(())
|
|
}
|
|
|
|
fn poll_flush(
|
|
self: Pin<&mut Self>,
|
|
_cx: &mut Context<'_>,
|
|
) -> Poll<Result<(), Self::Error>> {
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
|
|
fn poll_close(
|
|
self: Pin<&mut Self>,
|
|
_cx: &mut Context<'_>,
|
|
) -> Poll<Result<(), Self::Error>> {
|
|
Poll::Ready(Ok(()))
|
|
}
|
|
}
|
|
|
|
impl Stream for ControlledWebSocket {
|
|
type Item = Result<Message, std::convert::Infallible>;
|
|
|
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
|
Pin::new(&mut self.inbound_rx).poll_next(cx)
|
|
}
|
|
}
|
|
}
|