ive harnessed the harness
Something went wrong. Try again.
33 kB · 910 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911//! Remote relay protocol, framing, and multiplexing over SSH.//!//! Enforces://! 1. Anti-forgery nonce preamble: prevents shell startup output/banners from//! being mistaken for control messages or worker Hello.//! 2. Strict binary multiplex framing with NO resync: one corrupt byte or oversized//! frame invalidates the transport.//! 3. Clean multiplexing of worker control frames, worker stdout, worker stderr,//! process termination events, and kill requests.
use crate::target::TargetIdentity;use serde::{Deserialize, Serialize};use std::{ fmt, sync::{ atomic::{AtomicBool, Ordering}, Arc, },};use tokio::{ io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream}, sync::{mpsc, Mutex, Notify},};
pub const RELAY_PREAMBLE_PREFIX: &[u8] = b"KLBR-RELAY/1 ";pub const MAX_PREAMBLE_SCAN_BYTES: usize = 65_536;pub const MAX_FRAME_PAYLOAD_SIZE: usize = 1_048_576; // 1 MiBpub const DUPLEX_BUFFER_SIZE: usize = 65_536;
/// Multiplex channel tags.#[derive(Debug, Clone, Copy, PartialEq, Eq)]#[repr(u8)]pub enum Channel { /// Framed worker control message (bidirectional). Control = 0x01, /// Raw worker stdout chunks (target -> host). Stdout = 0x02, /// Raw worker stderr chunks (target -> host). Stderr = 0x03, /// Verified target identity report (target -> host, sent once at startup). StartupReport = 0x04, /// Remote worker process termination status (target -> host). ProcessEnd = 0x05, /// Request from host to kill remote worker process (host -> target). Kill = 0x06,}
impl Channel { pub fn from_u8(b: u8) -> Result<Self, RelayError> { match b { 0x01 => Ok(Self::Control), 0x02 => Ok(Self::Stdout), 0x03 => Ok(Self::Stderr), 0x04 => Ok(Self::StartupReport), 0x05 => Ok(Self::ProcessEnd), 0x06 => Ok(Self::Kill), other => Err(RelayError::UnknownChannel(other)), } }}
/// Remote process termination outcome.#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]pub enum RemoteProcessEnd { Exited(i64), Signaled(String), Unknown,}
/// Typed relay errors.#[derive(Debug)]pub enum RelayError { PreambleMissing(usize), NonceMismatch { expected: String, found: String }, FrameOversized(usize, usize), MalformedFrame(String), UnknownChannel(u8), TransportClosed, Io(std::io::Error), Json(serde_json::Error), StartupFailed(String),}
impl fmt::Display for RelayError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::PreambleMissing(scanned) => write!( f, "relay preamble missing or corrupt after scanning {scanned} bytes" ), Self::NonceMismatch { expected, found } => { write!(f, "nonce mismatch: expected '{expected}', found '{found}'") } Self::FrameOversized(len, max) => { write!(f, "oversized frame length {len} exceeds limit {max}") } Self::MalformedFrame(msg) => write!(f, "malformed frame: {msg}"), Self::UnknownChannel(tag) => write!(f, "unknown channel tag: 0x{tag:02x}"), Self::TransportClosed => write!(f, "transport closed unexpectedly"), Self::Io(err) => write!(f, "io error: {err}"), Self::Json(err) => write!(f, "json error: {err}"), Self::StartupFailed(err) => write!(f, "relay startup failed: {err}"), } }}
impl std::error::Error for RelayError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { Self::Io(err) => Some(err), Self::Json(err) => Some(err), _ => None, } }}
impl From<std::io::Error> for RelayError { fn from(err: std::io::Error) -> Self { Self::Io(err) }}
impl From<serde_json::Error> for RelayError { fn from(err: serde_json::Error) -> Self { Self::Json(err) }}
/// Writes the anti-forgery preamble line carrying the connection nonce.pub async fn write_preamble<W: AsyncWrite + Unpin>( writer: &mut W, nonce: &str,) -> Result<(), std::io::Error> { let line = format!("KLBR-RELAY/1 {nonce}\n"); writer.write_all(line.as_bytes()).await?; writer.flush().await?; Ok(())}
/// Result of scanning for the anti-forgery preamble on stdout.#[derive(Debug)]pub struct PreambleOutcome { /// Any banner/motd bytes emitted before the preamble line. pub diagnostic_banner: Vec<u8>, /// The nonce value verified on the wire. pub matched_nonce: String,}
/// Scans the reader until the `KLBR-RELAY/1 <nonce>\n` line is observed./// Captures any preceding bytes as a diagnostic banner (up to MAX_PREAMBLE_SCAN_BYTES)./// Fails with typed errors if the banner exceeds the bound or if nonce does not match.pub async fn read_preamble<R: AsyncRead + Unpin>( reader: &mut R, expected_nonce: &str,) -> Result<PreambleOutcome, RelayError> { let mut banner = Vec::new(); let prefix = RELAY_PREAMBLE_PREFIX; let mut match_idx = 0;
while banner.len() < MAX_PREAMBLE_SCAN_BYTES { let mut b = [0u8; 1]; let n = reader.read(&mut b).await?; if n == 0 { return Err(RelayError::PreambleMissing(banner.len())); }
let byte = b[0]; if byte == prefix[match_idx] { match_idx += 1; if match_idx == prefix.len() { // Found prefix! Read the rest of the line (the nonce) let mut nonce_bytes = Vec::new(); loop { let mut nb = [0u8; 1]; let r = reader.read(&mut nb).await?; if r == 0 { return Err(RelayError::PreambleMissing(banner.len())); } if nb[0] == b'\n' { break; } if nb[0] != b'\r' { nonce_bytes.push(nb[0]); } if nonce_bytes.len() > 256 { return Err(RelayError::MalformedFrame( "nonce too long in preamble".into(), )); } } let found_nonce = String::from_utf8_lossy(&nonce_bytes).trim().to_string(); if found_nonce != expected_nonce { return Err(RelayError::NonceMismatch { expected: expected_nonce.to_string(), found: found_nonce, }); } return Ok(PreambleOutcome { diagnostic_banner: banner, matched_nonce: found_nonce, }); } } else { if match_idx > 0 { // Partial match was false; push matched prefix bytes into banner banner.extend_from_slice(&prefix[..match_idx]); match_idx = 0; // Check if current byte matches start of prefix if byte == prefix[0] { match_idx = 1; continue; } } banner.push(byte); } }
Err(RelayError::PreambleMissing(banner.len()))}
/// Writes a single framed message: `[channel: u8][len: u32 BE][payload: len bytes]`.pub async fn write_relay_frame<W: AsyncWrite + Unpin>( writer: &mut W, channel: Channel, payload: &[u8],) -> Result<(), RelayError> { if payload.len() > MAX_FRAME_PAYLOAD_SIZE { return Err(RelayError::FrameOversized( payload.len(), MAX_FRAME_PAYLOAD_SIZE, )); } writer.write_u8(channel as u8).await?; writer.write_u32(payload.len() as u32).await?; writer.write_all(payload).await?; writer.flush().await?; Ok(())}
/// Reads a single framed message with strict parsing and NO resync.pub async fn read_relay_frame<R: AsyncRead + Unpin>( reader: &mut R,) -> Result<(Channel, Vec<u8>), RelayError> { let channel_byte = match reader.read_u8().await { Ok(b) => b, Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { return Err(RelayError::TransportClosed); } Err(e) => return Err(RelayError::Io(e)), }; let channel = Channel::from_u8(channel_byte)?; let len = match reader.read_u32().await { Ok(l) => l as usize, Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { return Err(RelayError::TransportClosed); } Err(e) => return Err(RelayError::Io(e)), }; if len > MAX_FRAME_PAYLOAD_SIZE { return Err(RelayError::FrameOversized(len, MAX_FRAME_PAYLOAD_SIZE)); } let mut buf = vec![0u8; len]; reader.read_exact(&mut buf).await?; Ok((channel, buf))}
/// An established host-side remote relay session./// Cloneable injection handle allowing gate tests to inject synthetic, late, or expired/// frames directly into the production inbound read path of a live session's Worker control channel.#[cfg(debug_assertions)]#[derive(Clone, Debug)]pub struct RelayControlInjector { inject_tx: mpsc::Sender<(Channel, Vec<u8>)>,}
#[cfg(debug_assertions)]impl RelayControlInjector { /// Injects raw control bytes into the inbound stream seen by the host Worker control channel. pub async fn inject_inbound_control_frame(&self, raw_bytes: &[u8]) -> Result<(), RelayError> { self.inject_tx .send((Channel::Control, raw_bytes.to_vec())) .await .map_err(|_| RelayError::TransportClosed) }
/// Serializes a protocol Envelope and injects it into the production read path of the host Worker. pub async fn inject_inbound_envelope( &self, envelope: &crate::protocol::Envelope, ) -> Result<(), RelayError> { let mut buf = Vec::new(); crate::protocol::write_frame(&mut buf, envelope) .await .map_err(|e| RelayError::MalformedFrame(format!("serialize envelope: {e}")))?; self.inject_inbound_control_frame(&buf).await }}
pub struct RelaySession { pub control: DuplexStream, pub stdout: DuplexStream, pub stderr: DuplexStream, pub identity: TargetIdentity, pub process_end: Arc<Mutex<Option<RemoteProcessEnd>>>, pub process_end_notify: Arc<Notify>, kill_tx: mpsc::Sender<()>, #[cfg(debug_assertions)] inject_tx: mpsc::Sender<(Channel, Vec<u8>)>, is_dead: Arc<AtomicBool>,}
impl RelaySession { /// Sends a kill request to the remote relay. pub async fn send_kill(&self) { let _ = self.kill_tx.send(()).await; }
/// Returns a control injector handle. Only available in debug builds (`#[cfg(debug_assertions)]`). #[cfg(debug_assertions)] pub fn injector(&self) -> RelayControlInjector { RelayControlInjector { inject_tx: self.inject_tx.clone(), } }
/// Checks whether the remote process has completed. pub async fn poll_process_end(&self) -> Option<RemoteProcessEnd> { self.process_end.lock().await.clone() }
/// Checks whether the transport or process has terminated. pub fn is_dead(&self) -> bool { self.is_dead.load(Ordering::SeqCst) }
/// Awaits remote process termination. pub async fn wait_process_end(&self) -> RemoteProcessEnd { loop { if let Some(end) = self.process_end.lock().await.clone() { return end; } self.process_end_notify.notified().await; } }}
/// Connects to a remote relay over an already-open async reader/writer (e.g. SSH stdin/stdout).////// Performs:/// 1. Nonce-matched preamble scan, capturing any diagnostic banner./// 2. Reads the initial `StartupReport` frame to verify target identity./// 3. Spawns the background bidirectional multiplexer task./// 4. Returns the `RelaySession` with framed duplex streams.pub async fn establish_relay_session<R, W>( mut reader: R, mut writer: W, expected_nonce: &str,) -> Result<(RelaySession, Vec<u8>), RelayError>where R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static,{ // 1. Anti-forgery preamble scan let preamble = read_preamble(&mut reader, expected_nonce).await?;
// 2. Initial StartupReport frame let (ch, payload) = read_relay_frame(&mut reader).await?; if ch != Channel::StartupReport { return Err(RelayError::MalformedFrame(format!( "expected StartupReport frame first, got channel {:?}", ch ))); } let identity: TargetIdentity = serde_json::from_slice(&payload)?;
// 3. Setup multiplex duplex channels let (host_ctrl, relay_ctrl) = tokio::io::duplex(DUPLEX_BUFFER_SIZE); let (mut relay_ctrl_read, mut relay_ctrl_write) = tokio::io::split(relay_ctrl); let (host_stdout, mut relay_stdout) = tokio::io::duplex(DUPLEX_BUFFER_SIZE); let (host_stderr, mut relay_stderr) = tokio::io::duplex(DUPLEX_BUFFER_SIZE);
let process_end = Arc::new(Mutex::new(None)); let process_end_notify = Arc::new(Notify::new()); let (kill_tx, mut kill_rx) = mpsc::channel::<()>(8); #[cfg(debug_assertions)] let (inject_tx, mut inject_rx) = mpsc::channel::<(Channel, Vec<u8>)>(16); let is_dead = Arc::new(AtomicBool::new(false));
let pe_clone = process_end.clone(); let notify_clone = process_end_notify.clone(); let is_dead_clone = is_dead.clone();
// Spawn outbound writer task: reads from relay_ctrl and kill_rx, writes frames to remote let (outbound_tx, mut outbound_rx) = mpsc::channel::<(Channel, Vec<u8>)>(64); let outbound_tx_ctrl = outbound_tx.clone();
// Task to pump host control bytes into outbound frames let pump_task = tokio::spawn(async move { let mut buf = vec![0u8; 16384]; loop { match relay_ctrl_read.read(&mut buf).await { Ok(0) => break, // closed Ok(n) => { if outbound_tx_ctrl .send((Channel::Control, buf[..n].to_vec())) .await .is_err() { break; } } Err(_) => break, } } });
// Task to handle kill requests let outbound_tx_kill = outbound_tx.clone(); tokio::spawn(async move { while kill_rx.recv().await.is_some() { let _ = outbound_tx_kill.send((Channel::Kill, Vec::new())).await; } });
// Task to serialize all outbound frames into writer tokio::spawn(async move { while let Some((ch, data)) = outbound_rx.recv().await { if write_relay_frame(&mut writer, ch, &data).await.is_err() { break; } } });
#[cfg(debug_assertions)] let inject_tx_clone = inject_tx.clone();
// Task to read inbound frames from remote reader and route to duplexes tokio::spawn(async move { loop { #[cfg(debug_assertions)] let inbound = tokio::select! { res = read_relay_frame(&mut reader) => res, Some((ch, data)) = inject_rx.recv() => { if ch == Channel::Control { if relay_ctrl_write.write_all(&data).await.is_err() { break; } let _ = relay_ctrl_write.flush().await; } continue; } };
#[cfg(not(debug_assertions))] let inbound = read_relay_frame(&mut reader).await;
match inbound { Ok((Channel::Control, data)) => { if relay_ctrl_write.write_all(&data).await.is_err() { break; } } Ok((Channel::Stdout, data)) => { if relay_stdout.write_all(&data).await.is_err() { break; } } Ok((Channel::Stderr, data)) => { if relay_stderr.write_all(&data).await.is_err() { break; } } Ok((Channel::ProcessEnd, data)) => { let end: RemoteProcessEnd = serde_json::from_slice(&data).unwrap_or(RemoteProcessEnd::Unknown); *pe_clone.lock().await = Some(end); notify_clone.notify_waiters(); is_dead_clone.store(true, Ordering::SeqCst); pump_task.abort(); drop(relay_ctrl_write); drop(relay_stdout); drop(relay_stderr); break; } Ok((Channel::StartupReport, _)) => { // Unexpected duplicate startup report - strict protocol violation is_dead_clone.store(true, Ordering::SeqCst); pump_task.abort(); drop(relay_ctrl_write); drop(relay_stdout); drop(relay_stderr); break; } Ok((Channel::Kill, _)) => { // Host should not receive kill from target } Err(_) => { // Transport dropped or framing violated let mut guard = pe_clone.lock().await; if guard.is_none() { *guard = Some(RemoteProcessEnd::Unknown); notify_clone.notify_waiters(); } is_dead_clone.store(true, Ordering::SeqCst); pump_task.abort(); drop(relay_ctrl_write); drop(relay_stdout); drop(relay_stderr); break; } } } });
let session = RelaySession { control: host_ctrl, stdout: host_stdout, stderr: host_stderr, identity, process_end, process_end_notify, kill_tx, #[cfg(debug_assertions)] inject_tx: inject_tx_clone, is_dead, };
Ok((session, preamble.diagnostic_banner))}
use crate::{ remote::ssh::{SshConfig, SshTransport}, target::{RemoteTarget, TargetMismatch}, worker::{ drain, BackendStartError, BoxFuture, ControlChannel, DrainTasks, LaunchedWorker, ProcessEnd, ProcessHandle, Tails, WorkerBackend, WorkerLaunch, },};use std::{path::PathBuf, time::Duration};
/// Process handle for a remote worker process executing behind klbr-relay.use crate::remote::ssh::SshProcess;
pub struct RemoteProcessHandle { pub kill_tx: mpsc::Sender<()>, pub process_end: Arc<Mutex<Option<RemoteProcessEnd>>>, pub process_end_notify: Arc<Notify>, pub ssh_process: Arc<Mutex<SshProcess>>,}
impl ProcessHandle for RemoteProcessHandle { fn kill(&mut self) -> BoxFuture<'_, ()> { let kill_tx = self.kill_tx.clone(); let ssh_process = self.ssh_process.clone(); Box::pin(async move { let _ = kill_tx.send(()).await; let mut guard = ssh_process.lock().await; let _ = guard.kill().await; }) }
fn wait(&mut self) -> BoxFuture<'_, ProcessEnd> { let process_end = self.process_end.clone(); let notify = self.process_end_notify.clone(); let ssh_process = self.ssh_process.clone(); Box::pin(async move { let end = loop { if let Some(end) = process_end.lock().await.clone() { break end; } notify.notified().await; }; let mut guard = ssh_process.lock().await; let _ = guard.wait().await; match end { RemoteProcessEnd::Exited(c) => ProcessEnd::Exited(c), RemoteProcessEnd::Signaled(s) => ProcessEnd::Signaled(s), RemoteProcessEnd::Unknown => ProcessEnd::Unknown, } }) }}
/// Remote worker backend: connects over SSH to target, launches klbr-relay,/// verifies target identity, and demultiplexes private control stream.pub struct RemoteBackend { pub config: SshConfig, pub target: RemoteTarget, pub relay_bin: PathBuf, pub python: PathBuf, pub runner: PathBuf, pub bundle_dir: Option<PathBuf>, #[cfg(debug_assertions)] pub active_injector: Arc<Mutex<Option<RelayControlInjector>>>,}
const RELAY_FAILURE_DIAGNOSTIC_GRACE: Duration = Duration::from_secs(1);
async fn remote_start_error( error: RelayError, target: &RemoteTarget, mut process: SshProcess,) -> BackendStartError { if tokio::time::timeout(RELAY_FAILURE_DIAGNOSTIC_GRACE, process.wait()) .await .is_err() { let _ = process.kill().await; let _ = process.wait().await; } let stderr = process.stderr_string().await; let missing_workspace = format!( "klbr-relay error: configured workspace {:?} does not exist or is not a directory", PathBuf::from(&target.workspace.path) ); if stderr.lines().any(|line| line == missing_workspace) { return BackendStartError::MissingWorkspace { workspace: target.workspace.path.clone(), }; }
let reason = match &error { RelayError::PreambleMissing(_) | RelayError::NonceMismatch { .. } => { format!("preamble handshake rejected: {error}") } _ => error.to_string(), }; let diagnostic = stderr.trim(); BackendStartError::ChannelFailed { detail: if diagnostic.is_empty() { reason } else { format!("{reason}; remote stderr: {diagnostic}") }, }}
impl RemoteBackend { pub fn new( config: SshConfig, target: RemoteTarget, relay_bin: PathBuf, python: PathBuf, runner: PathBuf, ) -> Self { Self { config, target, relay_bin, python, runner, bundle_dir: None, #[cfg(debug_assertions)] active_injector: Arc::new(Mutex::new(None)), } }
pub fn with_bundle_dir(mut self, bundle_dir: PathBuf) -> Self { self.bundle_dir = Some(bundle_dir); self }
/// Provides the live control injector handle from the most recent launch. /// Absent from release builds (`#[cfg(debug_assertions)]`). #[cfg(debug_assertions)] pub async fn active_injector(&self) -> Option<RelayControlInjector> { self.active_injector.lock().await.clone() }}
impl WorkerBackend for RemoteBackend { fn launch( &self, launch: WorkerLaunch, ) -> BoxFuture<'_, std::result::Result<LaunchedWorker, BackendStartError>> { Box::pin(async move { let nonce = uuid::Uuid::new_v4().simple().to_string(); let bundle_arg = self .bundle_dir .as_ref() .map(|b| format!(" --bundle-dir {}", b.display())) .unwrap_or_default();
let behavior_arg = if let Some(ref b) = launch.behavior { let path_buf = match &b.location { crate::worker::BehaviorLocation::LocalPath(p) => p.clone(), crate::worker::BehaviorLocation::RemoteBundle(bundle_id) => { if let Some(ref bdir) = self.bundle_dir { if bdir.join("behavior/manifest.json").exists() { bdir.join("behavior") } else { bdir.clone() } } else { std::path::PathBuf::from(bundle_id) } } }; format!( " --behavior {} --revision {}", path_buf.display(), b.revision ) } else { String::new() };
let target_arg = if let Some(ref stamp) = launch.binding { let json_str = serde_json::to_string(stamp).unwrap_or_default(); format!(" --target {}", SshTransport::shell_escape(&json_str)) } else { String::new() };
let remote_cmd = format!( "{} --session {} --generation {} --role {} --token {} --cwd {} --python {} --runner {} --nonce {}{}{}{}", self.relay_bin.display(), launch.session_id, launch.generation, match launch.role { crate::protocol::Role::Workbench => "workbench", crate::protocol::Role::Hooks => "hooks", crate::protocol::Role::Remote => "remote", }, launch.token, self.target.workspace.path, self.python.display(), self.runner.display(), nonce, bundle_arg, behavior_arg, target_arg );
let ssh_child = SshTransport::spawn(&self.config, &remote_cmd).map_err(|e| { BackendStartError::UnavailableTarget { target: self.target.id.to_string(), reason: format!("ssh spawn failed: {e}"), } })?;
let (ssh_stdin, ssh_stdout, mut ssh_process) = ssh_child.split();
let relay = tokio::time::timeout( launch.startup_deadline, establish_relay_session(ssh_stdout, ssh_stdin, &nonce), ) .await; let (session, _banner) = match relay { Ok(Ok(session)) => session, Ok(Err(error)) => { return Err(remote_start_error(error, &self.target, ssh_process).await) } Err(_) => { let _ = ssh_process.kill().await; let _ = ssh_process.wait().await; return Err(BackendStartError::Deadline); } };
// Strict target identity validation if let Err(mismatch) = self.target.verify_identity(&session.identity) { if let Some(tm) = mismatch.downcast_ref::<TargetMismatch>() { return Err(match tm { TargetMismatch::AccountMismatch { configured, resolved, .. } => BackendStartError::UnavailableTarget { target: self.target.id.to_string(), reason: format!( "account mismatch: configured {configured}, resolved {resolved}" ), }, TargetMismatch::PlatformMismatch { configured, resolved, .. } => BackendStartError::PlatformMismatch { required: configured.clone(), actual: resolved.clone(), }, TargetMismatch::WorkspaceUnusable { workspace, .. } => { BackendStartError::MissingWorkspace { workspace: workspace.clone(), } } }); } return Err(BackendStartError::UnavailableTarget { target: self.target.id.to_string(), reason: format!("{mismatch}"), }); }
#[cfg(debug_assertions)] { let injector = session.injector(); *self.active_injector.lock().await = Some(injector); }
let session_identity = session.identity.clone(); let process_handle = RemoteProcessHandle { kill_tx: session.kill_tx.clone(), process_end: session.process_end.clone(), process_end_notify: session.process_end_notify.clone(), ssh_process: Arc::new(Mutex::new(ssh_process)), };
let tails = Arc::new(std::sync::Mutex::new(Tails::default())); let drains = DrainTasks(vec![ tokio::spawn(drain(session.stdout, tails.clone(), false)), tokio::spawn(drain(session.stderr, tails.clone(), true)), ]);
let mut session_control = session.control;
let control_stream: Box<dyn ControlChannel> = if let Some(ref stamp) = launch.binding { use crate::protocol::{ read_frame, write_frame, HelloBinding, Message, RuntimeIdentity, WorkerIdentity, RELAY_PROTOCOL_VERSION, }; let envelope = read_frame(&mut session_control).await.map_err(|e| { BackendStartError::ChannelFailed { detail: format!("failed to read hello from remote relay: {e}"), } })?;
// Capabilities may ONLY come from a measurement. A remote runner probes // its own target and reports them at hello; any other role performs no // such probe, so the honest answer for it is NONE rather than a plausible // list. Fabricating capabilities here would advertise a shell or a // filesystem that was never observed - the exact "advertised but not // measured" failure that Capability::is_usable() exists to prevent. let capabilities = match &envelope.message { Message::Hello { identity: Some(identity), .. } => identity.runtime.capabilities.clone(), _ => Vec::new(), }; let worker_identity = WorkerIdentity { protocol_version: 1, relay_protocol_version: RELAY_PROTOCOL_VERSION, target: Some(HelloBinding { stamp: stamp.clone(), identity: Some(session_identity.clone()), }), runtime: RuntimeIdentity { platform: session_identity.platform.clone(), arch: session_identity.arch.clone(), interpreter_path: session_identity.interpreter_path.clone(), interpreter_version: session_identity.interpreter_version.clone(), environment_fingerprint: session_identity.environment_fingerprint.clone(), capabilities, }, };
let mut enriched = envelope; enriched.target = Some(stamp.clone()); if let Message::Hello { token, role, revision, slots, skills, .. } = enriched.message { enriched.message = Message::Hello { token, role, revision, slots, skills, identity: Some(worker_identity), }; }
let (host_side, mut bridge_side) = tokio::io::duplex(DUPLEX_BUFFER_SIZE); write_frame(&mut bridge_side, &enriched) .await .map_err(|e| BackendStartError::ChannelFailed { detail: format!("failed to write enriched hello: {e}"), })?;
tokio::spawn(async move { let (mut s_r, mut s_w) = tokio::io::split(session_control); let (mut b_r, mut b_w) = tokio::io::split(bridge_side); let _ = tokio::join!( tokio::io::copy(&mut s_r, &mut b_w), tokio::io::copy(&mut b_r, &mut s_w) ); });
Box::new(host_side) } else { Box::new(session_control) };
Ok(LaunchedWorker { control: control_stream, process: Box::new(process_handle), tails, drains, resolved: Some(session_identity), }) }) }}