ive harnessed the harness
Something went wrong. Try again.
16 kB · 478 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479//! Host-owned SSH transport for remote worker execution and bootstrap execution.//!//! Enforces://! 1. Genuine SSH transport with authentic known-hosts, keys, and target login shell.//! 2. Separation of SSH stderr into a background-drained bounded buffer so verbose//! banners, warnings, or connection errors cannot deadlock the channel.//! 3. Host ownership of all transport process creation (no unmanaged child SSH loops).//! 4. Bounded execution primitives for remote bootstrap and smoke verification.
use std::{path::PathBuf, process::Stdio, sync::Arc, time::Duration};use tokio::{ io::AsyncReadExt, process::{Child, ChildStderr, ChildStdin, ChildStdout, Command}, sync::{watch, Mutex},};
pub const MAX_STDERR_TAIL_BYTES: usize = 65_536;
/// SSH host key checking policy.#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum StrictHostKeyChecking { /// Strict checking against known_hosts: connection fails if host key is missing or changed. Yes, /// Automatically records new host keys on first connect, but strictly rejects changed keys. AcceptNew, /// Insecure: disables host key verification completely (for testing only). No,}
/// SSH configuration for connecting to a target.#[derive(Debug, Clone)]pub struct SshConfig { pub host: String, pub port: Option<u16>, pub user: Option<String>, pub identity_file: Option<PathBuf>, pub known_hosts_path: Option<PathBuf>, pub strict_host_key_checking: StrictHostKeyChecking, pub extra_args: Vec<String>,}
impl SshConfig { pub fn new(host: impl Into<String>) -> Self { Self { host: host.into(), port: None, user: None, identity_file: None, known_hosts_path: None, strict_host_key_checking: StrictHostKeyChecking::AcceptNew, extra_args: Vec::new(), } }
pub fn with_user(mut self, user: impl Into<String>) -> Self { self.user = Some(user.into()); self }
pub fn with_port(mut self, port: u16) -> Self { self.port = Some(port); self }
pub fn with_identity_file(mut self, path: impl Into<PathBuf>) -> Self { self.identity_file = Some(path.into()); self }
pub fn with_known_hosts(mut self, path: impl Into<PathBuf>) -> Self { self.known_hosts_path = Some(path.into()); self }
pub fn with_strict_host_key_checking(mut self, policy: StrictHostKeyChecking) -> Self { self.strict_host_key_checking = policy; self }}
/// An active SSH child process with isolated stdin/stdout streams and a separately-drained stderr tail.pub struct SshChild { pub stdin: ChildStdin, pub stdout: ChildStdout, pub process: SshProcess,}
/// The retained process and diagnostic handle after stdin/stdout are handed to the multiplexer.pub struct SshProcess { pub stderr_tail: Arc<Mutex<Vec<u8>>>, child: Child, stderr_drained: watch::Receiver<bool>,}
impl SshProcess { /// Kills the SSH transport child process. pub async fn kill(&mut self) -> Result<(), std::io::Error> { self.child.kill().await }
/// Awaits the SSH transport and its diagnostic stream. pub async fn wait(&mut self) -> Result<std::process::ExitStatus, std::io::Error> { let status = self.child.wait().await; let _ = self.stderr_drained.wait_for(|drained| *drained).await; status }
/// Retrieves the current drained stderr tail as string. pub async fn stderr_string(&self) -> String { let guard = self.stderr_tail.lock().await; String::from_utf8_lossy(&guard).to_string() }}
impl SshChild { pub fn split(self) -> (ChildStdin, ChildStdout, SshProcess) { (self.stdin, self.stdout, self.process) }
/// Kills the SSH transport child process. pub async fn kill(&mut self) -> Result<(), std::io::Error> { self.process.kill().await }
/// Awaits the SSH transport process exit. pub async fn wait(&mut self) -> Result<std::process::ExitStatus, std::io::Error> { self.process.wait().await }
/// Retrieves the current drained stderr tail as string. pub async fn stderr_string(&self) -> String { self.process.stderr_string().await }}
/// A remote command whose three streams are handed to the caller UNDRAINED.////// This is the evidence-complete counterpart to [`SshChild`]. `SshChild` drains stderr/// into a bounded tail in the background, which guarantees the channel cannot deadlock/// but makes complete stderr unrecoverable. When a caller must retain *complete*/// output as retrievable evidence, it takes ownership of the raw streams instead.////// # Hazard////// The caller MUST drain `stderr` concurrently with `stdout`. A remote process that/// fills its stderr pipe while the host reads only stdout will block forever. The/// bounded background drain that normally prevents this is deliberately absent here,/// so that responsibility transfers with the streams. The caller is likewise/// responsible for bounding what it retains and for reaping the child.pub struct SshStreams { pub stdin: ChildStdin, pub stdout: ChildStdout, pub stderr: ChildStderr, pub process: SshStreamingProcess,}
/// Process control for a streaming remote command.////// Deliberately does NOT expose `stderr_string()`: in streaming mode no background/// tail is collected, so such an accessor could only ever return an empty buffer and/// would misrepresent absent evidence as observed silence.pub struct SshStreamingProcess { child: Child,}
impl SshStreamingProcess { pub async fn kill(&mut self) -> Result<(), std::io::Error> { self.child.kill().await }
pub async fn wait(&mut self) -> Result<std::process::ExitStatus, std::io::Error> { self.child.wait().await }}
/// Result of a bounded remote command execution.#[derive(Debug, Clone)]pub struct RemoteExecResult { pub exit_code: i32, pub stdout: Vec<u8>, pub stderr: Vec<u8>, pub stdout_truncated: bool, pub stderr_truncated: bool,}
pub struct SshTransport;
impl SshTransport { /// Escapes an argument using POSIX shell single-quote wrapping. pub fn shell_escape(arg: &str) -> String { if arg.is_empty() { return "''".to_string(); } if arg .bytes() .all(|b| b.is_ascii_alphanumeric() || b"-_./:=+@".contains(&b)) { return arg.to_string(); } format!("'{}'", arg.replace('\'', "'\\''")) }
/// Builds the SSH argument list for launching a remote command. pub fn build_args(config: &SshConfig, remote_cmd: &str) -> Vec<String> { let host_key_check = match config.strict_host_key_checking { StrictHostKeyChecking::Yes => "yes", StrictHostKeyChecking::AcceptNew => "accept-new", StrictHostKeyChecking::No => "no", };
let mut args = vec![ "-o".to_string(), "BatchMode=yes".to_string(), "-o".to_string(), format!("StrictHostKeyChecking={host_key_check}"), ];
if let Some(port) = config.port { args.push("-p".to_string()); args.push(port.to_string()); }
if let Some(ref user) = config.user { args.push("-l".to_string()); args.push(user.clone()); }
if let Some(ref key) = config.identity_file { args.push("-i".to_string()); args.push(key.display().to_string()); }
if let Some(ref kh) = config.known_hosts_path { args.push("-o".to_string()); args.push(format!("UserKnownHostsFile={}", kh.display())); }
for extra in &config.extra_args { args.push(extra.clone()); }
args.push(config.host.clone()); args.push(remote_cmd.to_string());
args }
/// Spawns an SSH child process executing `remote_cmd` on the target host. pub fn spawn(config: &SshConfig, remote_cmd: &str) -> Result<SshChild, std::io::Error> { let args = Self::build_args(config, remote_cmd);
let mut cmd = Command::new("ssh"); cmd.args(&args); cmd.stdin(Stdio::piped()); cmd.stdout(Stdio::piped()); cmd.stderr(Stdio::piped());
let mut child = cmd.spawn()?; let stdin = child .stdin .take() .ok_or_else(|| std::io::Error::other("missing child stdin"))?; let stdout = child .stdout .take() .ok_or_else(|| std::io::Error::other("missing child stdout"))?; let mut stderr = child .stderr .take() .ok_or_else(|| std::io::Error::other("missing child stderr"))?;
let stderr_tail = Arc::new(Mutex::new(Vec::new())); let tail_clone = stderr_tail.clone(); let (stderr_drained_tx, stderr_drained) = watch::channel(false);
// Background task to drain SSH stderr separately so banners or error messages // never block the pipe or intermingle with stdout. tokio::spawn(async move { let mut buf = [0u8; 4096]; loop { match stderr.read(&mut buf).await { Ok(0) => break, Ok(n) => { let mut guard = tail_clone.lock().await; guard.extend_from_slice(&buf[..n]); if guard.len() > MAX_STDERR_TAIL_BYTES { let overflow = guard.len() - MAX_STDERR_TAIL_BYTES; guard.drain(0..overflow); } } Err(_) => break, } } let _ = stderr_drained_tx.send(true); });
Ok(SshChild { stdin, stdout, process: SshProcess { stderr_tail, child, stderr_drained, }, }) }
/// Executes an explicit argv command on the remote target with bounded output capture and timeout. /// /// Avoids unescaped shell string interpolation: escapes each element of `argv` safely. /// Bounded at `max_output_bytes` to prevent runaway memory allocation. /// Spawns an explicit argv command and hands the caller all three UNDRAINED streams. /// /// Escaping is identical to [`Self::spawn_argv`] and [`Self::exec_argv`]. Unlike those, /// no background stderr drain is installed and no output bound is applied, so complete /// stdout AND stderr can be retained as retrievable evidence. /// /// # Hazard /// /// The caller MUST drain stderr concurrently with stdout or the remote process can block /// on a full pipe. See [`SshStreams`]. pub fn spawn_argv_streaming( config: &SshConfig, argv: &[&str], ) -> Result<SshStreams, std::io::Error> { if argv.is_empty() { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "empty argv for remote spawn", )); }
let escaped_cmd = argv .iter() .map(|a| Self::shell_escape(a)) .collect::<Vec<_>>() .join(" "); let args = Self::build_args(config, &escaped_cmd);
let mut cmd = Command::new("ssh"); cmd.args(&args); cmd.stdin(Stdio::piped()); cmd.stdout(Stdio::piped()); cmd.stderr(Stdio::piped());
let mut child = cmd.spawn()?; let stdin = child .stdin .take() .ok_or_else(|| std::io::Error::other("missing child stdin"))?; let stdout = child .stdout .take() .ok_or_else(|| std::io::Error::other("missing child stdout"))?; let stderr = child .stderr .take() .ok_or_else(|| std::io::Error::other("missing child stderr"))?;
Ok(SshStreams { stdin, stdout, stderr, process: SshStreamingProcess { child }, }) }
/// Spawns an explicit argv command on the remote target, leaving stdin/stdout/stderr open. /// /// Identical escaping to [`Self::exec_argv`]: every element of `argv` is POSIX /// single-quote escaped and joined, so there is no shell interpolation of /// user-supplied values beyond per-argument escaping. Unlike `exec_argv` this /// performs no capture, bounding or timeout - the caller owns the streams and is /// responsible for draining, bounding and reaping them. pub fn spawn_argv(config: &SshConfig, argv: &[&str]) -> Result<SshChild, std::io::Error> { if argv.is_empty() { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "empty argv for remote spawn", )); }
let escaped_cmd = argv .iter() .map(|a| Self::shell_escape(a)) .collect::<Vec<_>>() .join(" ");
Self::spawn(config, &escaped_cmd) }
pub async fn exec_argv( config: &SshConfig, argv: &[&str], max_output_bytes: usize, timeout_duration: Duration, ) -> Result<RemoteExecResult, std::io::Error> { if argv.is_empty() { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "empty argv for remote exec", )); }
let escaped_cmd = argv .iter() .map(|a| Self::shell_escape(a)) .collect::<Vec<_>>() .join(" ");
let child = Self::spawn(config, &escaped_cmd)?;
let (_stdin, mut stdout, mut process) = child.split(); let stderr_tail = process.stderr_tail.clone();
let run_fut = async { let mut stdout_buf = Vec::new(); let mut stdout_truncated = false; let mut chunk = [0u8; 4096];
loop { let n = stdout.read(&mut chunk).await?; if n == 0 { break; } if stdout_buf.len() < max_output_bytes { let available = max_output_bytes - stdout_buf.len(); let to_take = n.min(available); stdout_buf.extend_from_slice(&chunk[..to_take]); if to_take < n { stdout_truncated = true; } } else { stdout_truncated = true; } }
let status = process.wait().await?; let stderr_bytes = { let guard = stderr_tail.lock().await; guard.clone() }; let stderr_truncated = stderr_bytes.len() >= MAX_STDERR_TAIL_BYTES;
Ok(RemoteExecResult { exit_code: status.code().unwrap_or(-1), stdout: stdout_buf, stderr: stderr_bytes, stdout_truncated, stderr_truncated, }) };
match tokio::time::timeout(timeout_duration, run_fut).await { Ok(res) => res, Err(_) => { let _ = process.kill().await; Err(std::io::Error::new( std::io::ErrorKind::TimedOut, "remote exec timed out", )) } } }
/// Convenience execution helper for bootstrap routines returning (ExitStatus, stdout, stderr_string). pub async fn execute( config: &SshConfig, remote_cmd: &str, ) -> Result<(std::process::ExitStatus, Vec<u8>, String), std::io::Error> { let mut child = Self::spawn(config, remote_cmd)?; let mut stdout_bytes = Vec::new(); let _ = child.stdout.read_to_end(&mut stdout_bytes).await?; let status = child.wait().await?; let stderr_str = child.stderr_string().await; Ok((status, stdout_bytes, stderr_str)) }}