//! Harness side of the Noise relay.//!//! The rendezvous service routes frames by `stream_id`, but does not authenticate//! the executor or see JSON-RPC plaintext. We claim a stream, complete hybrid IK//! against the registry-provided executor key, and then expose the result as a//! normal `JsonRpcConnection`. Outbound JSON-RPC is framed and split into Noise//! records; inbound records are reordered before decryption and reassembly.use futures::Sink;use futures::SinkExt;use futures::Stream;use futures::StreamExt;use tokio::sync::mpsc;use tokio::sync::watch;use tokio_tungstenite::tungstenite::Message;use tracing::Instrument;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::noise_channel::InitiatorHandshake;use crate::noise_channel::NoiseChannelIdentity;use crate::noise_channel::NoiseChannelPublicKey;use crate::noise_channel::NoiseTransport;use crate::noise_channel::noise_channel_prologue;use crate::noise_relay::message_framing::JsonRpcMessageDecoder;use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN;use crate::noise_relay::message_framing::frame_jsonrpc_message;use crate::noise_relay::ordered_ciphertext::OrderedCiphertextFrames;use crate::noise_relay::take_next_sequence;use crate::relay::RelayFrameBodyKind;use crate::relay::decode_relay_message_frame;use crate::relay::encode_relay_message_frame;use crate::relay_proto::RelayData;use crate::relay_proto::RelayMessageFrame;/// Values that bind one harness websocket to the intended executor registration.////// These fields all come from the same registry response. Keeping them together/// makes that relationship visible at the call site and avoids mixing up the/// several string and key arguments used to start the handshake.pub(crate) struct NoiseHarnessConnectionArgs {    pub(crate) connection_label: String,    pub(crate) environment_id: String,    pub(crate) executor_registration_id: String,    pub(crate) identity: NoiseChannelIdentity,    pub(crate) responder_public_key: NoiseChannelPublicKey,    pub(crate) harness_key_authorization: String,}// Reset frames are cleartext relay control and are not authenticated by Noise.// Preserve the availability signal while replacing attacker-controlled reason// text before it reaches disconnect diagnostics.const NOISE_RELAY_RESET_DISCONNECT_REASON: &str = "Noise relay stream reset";/// Adapt one harness rendezvous websocket into an authenticated JSON-RPC connection.////// The returned connection is not usable until the background task completes/// hybrid IK against the registry-pinned exec-server key. Rendezvous can see/// stream metadata and ciphertext, but never JSON-RPC plaintext or either/// endpoint's private key. Failures close the connection rather than falling/// back to plaintext.pub(crate) fn noise_harness_connection_from_websocket<T, E>(    stream: T,    args: NoiseHarnessConnectionArgs,) -> JsonRpcConnectionwhere    T: Sink<Message, Error = E> + Stream<Item = Result<Message, E>> + Unpin + Send + 'static,    E: std::fmt::Display + Send + 'static,{    let NoiseHarnessConnectionArgs {        connection_label,        environment_id,        executor_registration_id,        identity,        responder_public_key,        harness_key_authorization,    } = args;    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 stream_span = tracing::debug_span!(        "noise_relay.stream",        noise_side = "harness",        environment_id = %environment_id,        executor_registration_id = %executor_registration_id,        stream_id = %stream_id,    );    let websocket_task = tokio::spawn(async move {        let mut websocket = stream;        // Bind the Noise transcript to the exact environment registration and        // virtual relay stream before emitting any handshake bytes. A captured        // handshake cannot be spliced onto a different routed connection.        let prologue =            noise_channel_prologue(&environment_id, &executor_registration_id, &stream_id);        let (initiator_handshake, request) = match InitiatorHandshake::start(            &identity,            &responder_public_key,            &prologue,            harness_key_authorization.as_bytes(),        ) {            Ok(handshake) => handshake,            Err(error) => {                send_disconnected(                    &incoming_tx,                    &disconnected_tx,                    format!("failed to start Noise relay handshake: {error}"),                )                .await;                return;            }        };        // Resume claims the stream ID at rendezvous; Handshake carries the        // opaque first IK message. No JSON-RPC data is sent before the        // responder proves possession of the pinned static key.        let resume = RelayMessageFrame::resume(stream_id.clone());        let handshake = RelayMessageFrame::handshake(stream_id.clone(), request);        if websocket            .send(Message::Binary(encode_relay_message_frame(&resume).into()))            .await            .is_err()            || websocket                .send(Message::Binary(                    encode_relay_message_frame(&handshake).into(),                ))                .await                .is_err()        {            let _ = disconnected_tx.send(true);            return;        }        // During the handshake, ignore unrelated routed streams and control        // frames, but reject data on our stream. Accepting early data would        // create a plaintext or unauthenticated application path.        let mut transport = loop {            let Some(incoming_message) = websocket.next().await else {                send_disconnected(                    &incoming_tx,                    &disconnected_tx,                    "Noise relay websocket ended during handshake".to_string(),                )                .await;                return;            };            let message = match incoming_message {                Ok(Message::Binary(payload)) => payload,                Ok(Message::Close(_)) => {                    send_disconnected(                        &incoming_tx,                        &disconnected_tx,                        "Noise relay websocket received close frame during handshake".to_string(),                    )                    .await;                    return;                }                Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue,                Ok(Message::Text(_)) => {                    send_disconnected(                        &incoming_tx,                        &disconnected_tx,                        "Noise relay transport expects binary protobuf frames".to_string(),                    )                    .await;                    return;                }                Err(error) => {                    send_disconnected(                        &incoming_tx,                        &disconnected_tx,                        format!(                            "failed to read Noise relay websocket from {connection_label}: {error}"                        ),                    )                    .await;                    return;                }            };            let frame = match decode_relay_message_frame(message.as_ref()) {                Ok(frame) => frame,                Err(error) => {                    send_disconnected(                        &incoming_tx,                        &disconnected_tx,                        format!("failed to parse Noise relay frame: {error}"),                    )                    .await;                    return;                }            };            if frame.stream_id != stream_id {                debug!("Noise relay ignored frame for unrelated stream during handshake");                continue;            }            match frame.validate() {                Ok(RelayFrameBodyKind::Handshake) => {                    let response = match frame.into_handshake_payload() {                        Ok(response) => response,                        Err(error) => {                            send_disconnected(                                &incoming_tx,                                &disconnected_tx,                                format!("invalid Noise relay handshake response: {error}"),                            )                            .await;                            return;                        }                    };                    match initiator_handshake.finish(&response) {                        Ok(transport) => {                            info!(                                noise_event = "handshake",                                noise_outcome = "ok",                                "Noise harness handshake completed"                            );                            break transport;                        }                        Err(error) => {                            send_disconnected(                                &incoming_tx,                                &disconnected_tx,                                format!("Noise relay handshake failed: {error}"),                            )                            .await;                            return;                        }                    }                }                Ok(RelayFrameBodyKind::Reset) => {                    send_disconnected(                        &incoming_tx,                        &disconnected_tx,                        NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(),                    )                    .await;                    return;                }                Ok(                    RelayFrameBodyKind::Ack                    | RelayFrameBodyKind::Resume                    | RelayFrameBodyKind::Heartbeat,                ) => {}                Ok(RelayFrameBodyKind::Data) | Err(_) => {                    send_disconnected(                        &incoming_tx,                        &disconnected_tx,                        "Noise relay received data before handshake completion".to_string(),                    )                    .await;                    return;                }            }        };        // After the handshake, each relay sequence maps to exactly one Noise        // transport record. Outbound records are encrypted once; inbound        // records are reordered and deduplicated before decryption.        let mut next_outbound_seq = 0u32;        let mut inbound_ciphertexts = OrderedCiphertextFrames::default();        let mut inbound_decoder = JsonRpcMessageDecoder::default();        'relay: loop {            tokio::select! {                maybe_message = outgoing_rx.recv() => {                    let Some(message) = maybe_message else {                        break;                    };                    let framed = match frame_jsonrpc_message(&message) {                        Ok(framed) => framed,                        Err(error) => {                            warn!("failed to frame JSON-RPC payload for Noise relay: {error}");                            break;                        }                    };                    for plaintext_record in framed.chunks(NOISE_RECORD_PLAINTEXT_LEN) {                        let seq = match take_next_sequence(&mut next_outbound_seq) {                            Ok(seq) => seq,                            Err(error) => {                                warn!("Noise relay sequence exhausted: {error}");                                break 'relay;                            }                        };                        let ciphertext = match transport.encrypt(plaintext_record) {                            Ok(ciphertext) => ciphertext,                            Err(error) => {                                warn!("failed to encrypt JSON-RPC payload for Noise relay: {error}");                                break 'relay;                            }                        };                        let frame = RelayMessageFrame::data(stream_id.clone(), seq, ciphertext);                        if let Err(error) = websocket                            .send(Message::Binary(encode_relay_message_frame(&frame).into()))                            .await                        {                            warn!("failed to write Noise relay websocket: {error}");                            break 'relay;                        }                    }                }                incoming_message = websocket.next() => {                    let Some(incoming_message) = incoming_message else {                        break;                    };                    match incoming_message {                        Ok(Message::Binary(payload)) => {                            let frame = match decode_relay_message_frame(payload.as_ref()) {                                Ok(frame) => frame,                                Err(error) => {                                    send_malformed(&incoming_tx, error.to_string()).await;                                    break;                                }                            };                            if frame.stream_id != stream_id {                                continue;                            }                            match frame.validate() {                                Ok(RelayFrameBodyKind::Data) => {                                    let data = match frame.into_data() {                                        Ok(data) => data,                                        Err(error) => {                                            send_malformed(&incoming_tx, error.to_string()).await;                                            break;                                        }                                    };                                    if let Err(error) = receive_data(                                        &mut inbound_ciphertexts,                                        &mut transport,                                        &mut inbound_decoder,                                        data,                                        &incoming_tx,                                    )                                    .await                                    {                                        send_malformed(&incoming_tx, error.to_string()).await;                                        break;                                    }                                }                                Ok(RelayFrameBodyKind::Reset) => {                                    let _ = incoming_tx                                        .send(JsonRpcConnectionEvent::Disconnected {                                            reason: Some(                                                NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(),                                            ),                                        })                                        .await;                                    break;                                }                                Ok(                                    RelayFrameBodyKind::Ack                                    | RelayFrameBodyKind::Resume                                    | RelayFrameBodyKind::Heartbeat,                                ) => {}                                Ok(RelayFrameBodyKind::Handshake) | Err(_) => {                                    send_malformed(                                        &incoming_tx,                                        "Noise relay received invalid post-handshake frame".to_string(),                                    )                                    .await;                                    break;                                }                            }                        }                        Ok(Message::Close(_)) => break,                        Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => {}                        Ok(Message::Text(_)) => {                            send_malformed(                                &incoming_tx,                                "Noise relay transport expects binary protobuf frames".to_string(),                            )                            .await;                            break;                        }                        Err(error) => {                            debug!("Noise relay websocket read failed: {error}");                            break;                        }                    }                }            }        }        let _ = disconnected_tx.send(true);    }    .instrument(stream_span));    JsonRpcConnection {        outgoing_tx,        incoming_rx,        disconnected_rx,        task_handles: vec![websocket_task],        transport: JsonRpcTransport::Plain,    }}/// Order and decrypt one relay frame, then emit any complete JSON-RPC messages./// Relay records and JSON-RPC messages do not share boundaries, so reassembly/// happens after decryption.async fn receive_data(    inbound_ciphertexts: &mut OrderedCiphertextFrames,    transport: &mut NoiseTransport,    decoder: &mut JsonRpcMessageDecoder,    data: RelayData,    incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,) -> Result<(), ExecServerError> {    // Ordering must happen before decryption because Noise transport nonces are    // implicit. A future or duplicate ciphertext passed directly to Clatter    // would desynchronize the channel.    for ciphertext in inbound_ciphertexts.push(data.seq, data.payload)? {        let plaintext = transport.decrypt(&ciphertext).map_err(|error| {            ExecServerError::Protocol(format!("Noise relay decryption failed: {error}"))        })?;        // The authenticated byte stream can carry partial or multiple JSON-RPC        // messages; emit only complete, successfully parsed messages.        for message in decoder.push(&plaintext)? {            incoming_tx                .send(JsonRpcConnectionEvent::Message(message))                .await                .map_err(|_| ExecServerError::Closed)?;        }    }    Ok(())}async fn send_malformed(incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>, reason: String) {    let _ = incoming_tx        .send(JsonRpcConnectionEvent::MalformedMessage { reason })        .await;}async fn send_disconnected(    incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,    disconnected_tx: &watch::Sender<bool>,    reason: String,) {    let _ = disconnected_tx.send(true);    let _ = incoming_tx        .send(JsonRpcConnectionEvent::Disconnected {            reason: Some(reason),        })        .await;}