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