//! 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 MiB
pub 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),
})
})
}
}