//! Standalone remote worker relay launcher.
//!
//! Executed on the target host (e.g. over SSH) by the klbr host runtime.
//! Establishes a private local UnixListener on the target, spawns the local Python worker,
//! and multiplexes framed control, child stdout, and child stderr back over stdout/stdin.
use klbr_runtime::{
remote::relay::{write_preamble, write_relay_frame, Channel, RemoteProcessEnd},
target::TargetIdentity,
};
use std::{
path::{Path, PathBuf},
process::Stdio,
sync::Arc,
time::Duration,
};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::UnixListener,
process::Command,
sync::{mpsc, Mutex},
};
#[derive(Debug)]
struct RelayArgs {
session: String,
generation: String,
role: String,
token: String,
cwd: PathBuf,
python: PathBuf,
runner: PathBuf,
nonce: String,
bundle_dir: Option<PathBuf>,
behavior: Option<PathBuf>,
revision: Option<String>,
target: Option<String>,
}
fn parse_args() -> Result<RelayArgs, String> {
let mut session = None;
let mut generation = None;
let mut role = None;
let mut token = None;
let mut cwd = None;
let mut python = None;
let mut runner = None;
let mut nonce = std::env::var("KLBR_RELAY_NONCE").ok();
let mut bundle_dir = None;
let mut behavior = None;
let mut revision = None;
let mut target = None;
let args: Vec<String> = std::env::args().skip(1).collect();
let mut i = 0;
while i < args.len() {
match args[i].as_str() {
"--session" => {
i += 1;
session = args.get(i).cloned();
}
"--generation" => {
i += 1;
generation = args.get(i).cloned();
}
"--role" => {
i += 1;
role = args.get(i).cloned();
}
"--token" => {
i += 1;
token = args.get(i).cloned();
}
"--cwd" => {
i += 1;
cwd = args.get(i).map(PathBuf::from);
}
"--python" => {
i += 1;
python = args.get(i).map(PathBuf::from);
}
"--runner" => {
i += 1;
runner = args.get(i).map(PathBuf::from);
}
"--nonce" => {
i += 1;
nonce = args.get(i).cloned();
}
"--bundle-dir" => {
i += 1;
bundle_dir = args.get(i).map(PathBuf::from);
}
"--behavior" => {
i += 1;
behavior = args.get(i).map(PathBuf::from);
}
"--revision" => {
i += 1;
revision = args.get(i).cloned();
}
"--target" => {
i += 1;
target = args.get(i).cloned();
}
other => return Err(format!("unknown argument: {other}")),
}
i += 1;
}
Ok(RelayArgs {
session: session.ok_or("missing --session")?,
generation: generation.ok_or("missing --generation")?,
role: role.ok_or("missing --role")?,
token: token.ok_or("missing --token")?,
cwd: cwd.ok_or("missing --cwd")?,
python: python.ok_or("missing --python")?,
runner: runner.ok_or("missing --runner")?,
nonce: nonce.ok_or("missing --nonce or KLBR_RELAY_NONCE")?,
bundle_dir,
behavior,
revision,
target,
})
}
fn gather_identity(cwd: &Path, python: &Path) -> Result<TargetIdentity, String> {
let account = std::env::var("USER")
.or_else(|_| std::env::var("LOGNAME"))
.unwrap_or_else(|_| "unknown".into());
let machine = match std::process::Command::new("uname").arg("-n").output() {
Ok(out) if out.status.success() => String::from_utf8_lossy(&out.stdout).trim().to_string(),
_ => "unknown-machine".into(),
};
let platform = match std::process::Command::new("uname").arg("-s").output() {
Ok(out) if out.status.success() => String::from_utf8_lossy(&out.stdout).trim().to_string(),
_ => std::env::consts::OS.to_string(),
};
let arch = match std::process::Command::new("uname").arg("-m").output() {
Ok(out) if out.status.success() => String::from_utf8_lossy(&out.stdout).trim().to_string(),
_ => std::env::consts::ARCH.to_string(),
};
let canonical_cwd = std::fs::canonicalize(cwd)
.map_err(|e| format!("failed to canonicalize cwd {}: {e}", cwd.display()))?
.display()
.to_string();
let canonical_python = std::fs::canonicalize(python)
.map_err(|e| format!("failed to canonicalize python {}: {e}", python.display()))?
.display()
.to_string();
let py_ver = match std::process::Command::new(python).arg("--version").output() {
Ok(out) if out.status.success() => {
let combined = format!(
"{} {}",
String::from_utf8_lossy(&out.stdout),
String::from_utf8_lossy(&out.stderr)
);
combined.trim().to_string()
}
_ => "Python 3.unknown".into(),
};
TargetIdentity::new(
account,
machine,
platform,
arch,
canonical_python,
py_ver,
canonical_cwd,
None,
1,
)
.map_err(|e| format!("failed to construct TargetIdentity: {e}"))
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = match parse_args() {
Ok(a) => a,
Err(e) => {
eprintln!("klbr-relay argument error: {e}");
std::process::exit(102);
}
};
// Strict workspace law: configured workspace MUST exist and be a directory.
// NEVER fall back to ~ or local directory.
if !args.cwd.is_dir() {
eprintln!(
"klbr-relay error: configured workspace {:?} does not exist or is not a directory",
args.cwd
);
std::process::exit(103);
}
if !args.python.is_file() {
eprintln!(
"klbr-relay error: configured python binary {:?} does not exist",
args.python
);
std::process::exit(104);
}
if !args.runner.is_file() {
eprintln!(
"klbr-relay error: configured runner script {:?} does not exist",
args.runner
);
std::process::exit(105);
}
// 1. Anti-forgery preamble: emit nonce-bearing line to stdout first
let mut stdout = tokio::io::stdout();
write_preamble(&mut stdout, &args.nonce).await?;
// 2. Gather verified target identity and send StartupReport frame
let identity = match gather_identity(&args.cwd, &args.python) {
Ok(id) => id,
Err(e) => {
eprintln!("klbr-relay identity error: {e}");
std::process::exit(106);
}
};
let _bundle = &args.bundle_dir;
let identity_bytes = serde_json::to_vec(&identity)?;
write_relay_frame(&mut stdout, Channel::StartupReport, &identity_bytes).await?;
// 3. Create private temp directory for private UnixListener control socket
let tempdir = tempfile::Builder::new().prefix("klbr-relay-").tempdir()?;
let sock_path = tempdir.path().join("control.sock");
let listener = UnixListener::bind(&sock_path)?;
// 4. Spawn child Python worker process
let mut cmd = Command::new(&args.python);
cmd.args(["-I", "-B", "-u"]).arg(&args.runner).args([
"--socket",
&sock_path.display().to_string(),
"--role",
&args.role,
"--session",
&args.session,
"--generation",
&args.generation,
"--token",
&args.token,
]);
if let Some(ref b) = args.behavior {
cmd.arg("--behavior").arg(b);
}
if let Some(ref rev) = args.revision {
cmd.arg("--revision").arg(rev);
}
if let Some(ref tgt) = args.target {
cmd.arg("--target").arg(tgt);
}
// An auxiliary remote namespace is confined to the verified workspace root; the
// dispatcher refuses paths outside it. Passing the already-validated cwd keeps a
// single source of truth for "where this namespace may act".
if args.role == "remote" {
cmd.arg("--workspace").arg(&args.cwd);
}
cmd.current_dir(&args.cwd)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = cmd.spawn()?;
let child_stdout = child.stdout.take().expect("child stdout");
let child_stderr = child.stderr.take().expect("child stderr");
// 5. Accept connection from Python worker with 15-second startup timeout
let accept_res = tokio::time::timeout(Duration::from_secs(15), listener.accept()).await;
let (unix_stream, _) = match accept_res {
Ok(Ok(pair)) => pair,
Ok(Err(e)) => {
eprintln!("klbr-relay accept error: {e}");
let _ = child.kill().await;
std::process::exit(107);
}
Err(_) => {
eprintln!("klbr-relay startup timeout waiting for worker socket connection");
let _ = child.kill().await;
std::process::exit(108);
}
};
// 6. Multiplex streams
let (outbound_tx, mut outbound_rx) = mpsc::channel::<(Channel, Vec<u8>)>(64);
let writer = Arc::new(Mutex::new(tokio::io::stdout()));
let writer_clone = writer.clone();
// Outbound serialization task: the only writer of relay frames.
let serializer = tokio::spawn(async move {
while let Some((ch, data)) = outbound_rx.recv().await {
let mut w = writer_clone.lock().await;
if write_relay_frame(&mut *w, ch, &data).await.is_err() {
break;
}
}
});
// Task: worker UnixStream -> relay stdout (Channel::Control)
let tx_ctrl = outbound_tx.clone();
let (mut unix_reader, mut unix_writer) = unix_stream.into_split();
let control_task = tokio::spawn(async move {
let mut buf = vec![0u8; 16384];
loop {
match unix_reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if tx_ctrl
.send((Channel::Control, buf[..n].to_vec()))
.await
.is_err()
{
break;
}
}
Err(_) => break,
}
}
});
// Task: worker stdout -> relay stdout (Channel::Stdout)
let tx_stdout = outbound_tx.clone();
let mut child_stdout_r = child_stdout;
let stdout_task = tokio::spawn(async move {
let mut buf = vec![0u8; 16384];
loop {
match child_stdout_r.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if tx_stdout
.send((Channel::Stdout, buf[..n].to_vec()))
.await
.is_err()
{
break;
}
}
Err(_) => break,
}
}
});
// Task: worker stderr -> relay stdout (Channel::Stderr)
let tx_stderr = outbound_tx.clone();
let mut child_stderr_r = child_stderr;
let stderr_task = tokio::spawn(async move {
let mut buf = vec![0u8; 16384];
loop {
match child_stderr_r.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if tx_stderr
.send((Channel::Stderr, buf[..n].to_vec()))
.await
.is_err()
{
break;
}
}
Err(_) => break,
}
}
});
// Task: relay stdin -> worker UnixStream / Kill
let (kill_tx, mut kill_rx) = mpsc::channel::<()>(1);
let mut stdin_reader = tokio::io::stdin();
tokio::spawn(async move {
loop {
match klbr_runtime::remote::relay::read_relay_frame(&mut stdin_reader).await {
Ok((Channel::Control, data)) => {
if unix_writer.write_all(&data).await.is_err() {
break;
}
}
Ok((Channel::Kill, _)) => {
let _ = kill_tx.send(()).await;
break;
}
_ => break,
}
}
});
// Wait for worker process to exit or kill request from host
let exit_status = tokio::select! {
res = child.wait() => res,
_ = kill_rx.recv() => {
let _ = child.kill().await;
child.wait().await
}
};
let end = match exit_status {
Ok(s) if s.code().is_some() => RemoteProcessEnd::Exited(s.code().unwrap() as i64),
Ok(s) => RemoteProcessEnd::Signaled(format!("{s}")),
Err(_e) => RemoteProcessEnd::Unknown,
};
// The child is gone, but its pipes can still hold the last thing it said: a reader task may be
// mid-read, and a frame may be queued behind it. The terminal frame is written by the same
// serializer as everything else, after the readers drain, or the host hears the end of the
// process before it hears the process.
let _ = tokio::time::timeout(Duration::from_secs(5), async {
let _ = tokio::join!(control_task, stdout_task, stderr_task);
})
.await;
let end_bytes = serde_json::to_vec(&end)?;
let _ = outbound_tx.send((Channel::ProcessEnd, end_bytes)).await;
drop(outbound_tx);
let _ = serializer.await;
let _ = writer.lock().await.flush().await;
Ok(())
}