File size: 7,445 Bytes
ea39c0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | //! One executor-side virtual stream after the Noise handshake.
//!
//! The environment loop owns reads and a per-stream task owns writes. They share
//! `NoiseTransport` because its send and receive nonces live in the same value;
//! the mutex is never held across `.await`.
use std::sync::Arc;
use std::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tracing::warn;
use crate::ExecServerError;
use crate::connection::CHANNEL_CAPACITY;
use crate::noise_channel::NoiseTransport;
use crate::noise_relay::message_framing::MessageDecoder;
use crate::noise_relay::message_framing::NOISE_RECORD_PLAINTEXT_LEN;
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::telemetry::ExecutorRegistration;
/// Identifies one completed virtual-stream instance.
///
/// Stream IDs are supplied by the untrusted relay peer and may be reused. The
/// instance ID prevents a delayed writer notification from removing a newer
/// stream that happens to use the same routing ID.
pub(crate) struct ClosedNoiseVirtualStream {
pub(crate) stream_id: String,
pub(crate) instance_id: u64,
}
/// 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<H: NoiseStreamHandler> {
incoming_tx: mpsc::Sender<H::Incoming>,
disconnected_tx: watch::Sender<bool>,
transport: Arc<Mutex<NoiseTransport>>,
inbound_ciphertexts: OrderedCiphertextFrames,
inbound_decoder: MessageDecoder,
pub(crate) instance_id: u64,
}
impl<H: NoiseStreamHandler> NoiseVirtualStream<H> {
pub(crate) fn disconnect(self) {
let _ = self.disconnected_tx.send(true);
}
/// 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)? {
let plaintext = {
let mut transport = self
.transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.decrypt(&ciphertext).map_err(|error| {
ExecServerError::Protocol(format!("Noise relay decryption failed: {error}"))
})?
};
for message in self.inbound_decoder.push(&plaintext)? {
self.incoming_tx
.try_send(H::decode(message)?)
.map_err(|_| {
ExecServerError::Protocol(
"Noise virtual stream inbound queue is full or closed".to_string(),
)
})?;
}
}
Ok(())
}
}
/// 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<H: NoiseStreamHandler>(
stream_id: String,
instance_id: u64,
handler: H,
physical_outgoing_tx: mpsc::Sender<Vec<u8>>,
closed_stream_tx: mpsc::Sender<ClosedNoiseVirtualStream>,
transport: NoiseTransport,
executor_registration: Option<Arc<ExecutorRegistration>>,
) -> 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 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) = outgoing_rx.recv().await {
let message = match H::encode(message) {
Ok(message) => message,
Err(error) => {
warn!("failed to encode Noise virtual stream payload: {error}");
break;
}
};
// 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) => {
warn!("Noise virtual stream sequence exhausted: {error}");
break 'writer;
}
};
let ciphertext = {
let mut transport = writer_transport
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
transport.encrypt(plaintext_record)
};
let ciphertext = match ciphertext {
Ok(ciphertext) => ciphertext,
Err(error) => {
warn!("failed to encrypt Noise virtual stream payload: {error}");
break 'writer;
}
};
let frame = RelayMessageFrame::data(
writer_stream_id.clone(),
seq,
ciphertext,
trace.take(),
);
if physical_outgoing_tx
.send(encode_relay_message_frame(&frame))
.await
.is_err()
{
break 'writer;
}
}
}
// The physical relay owns reset delivery and rejects stale instance IDs.
let closed_stream = ClosedNoiseVirtualStream {
stream_id: writer_stream_id,
instance_id,
};
let _ = closed_stream_tx.send(closed_stream).await;
});
let connection = NoiseStreamConnection {
outgoing_tx,
incoming_rx,
disconnected_rx,
writer_task,
executor_registration,
};
tokio::spawn(async move {
handler.run_connection(connection).await;
let _ = owner_closed_stream_tx
.send(ClosedNoiseVirtualStream {
stream_id: owner_stream_id,
instance_id,
})
.await;
});
NoiseVirtualStream {
incoming_tx,
disconnected_tx,
transport,
inbound_ciphertexts: OrderedCiphertextFrames::default(),
inbound_decoder: MessageDecoder::default(),
instance_id,
}
}
#[cfg(test)]
#[path = "executor_stream_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "forward_stream_tests.rs"]
mod forward_tests;
|