launcher/src/ssh.rs

use anyhow::{Context, Result};
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use russh::client;
use russh::keys::{PrivateKeyWithHashAlg, load_secret_key as russh_load_secret_key};
use russh::{ChannelMsg, Error as SshError};
use sha2::{Digest, Sha256};
use std::path::Path;
use std::sync::Arc;
use tokio::time::{Duration, sleep};
use tracing::warn;

use crate::config::{Config, Node};

const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);

/// Establishes an SSH session to a node using the config credentials, with a 5s connect
/// and inactivity timeout.  Use [`ssh_session_follow`] for long-lived streaming commands.
pub async fn ssh_session(config: &Config, node: &Node) -> Result<client::Handle<VerifyingHandler>> {
    ssh_session_with_timeout(config, node, Some(CONNECT_TIMEOUT)).await
}

/// Like [`ssh_session`] but with no inactivity timeout — suitable for long-lived
/// streaming commands such as `docker logs -f`.
pub async fn ssh_session_follow(
    config: &Config,
    node: &Node,
) -> Result<client::Handle<VerifyingHandler>> {
    ssh_session_with_timeout(config, node, None).await
}

async fn ssh_session_with_timeout(
    config: &Config,
    node: &Node,
    inactivity_timeout: Option<Duration>,
) -> Result<client::Handle<VerifyingHandler>> {
    if config.ssh_user.is_empty() || config.ssh_key_path.as_os_str().is_empty() {
        anyhow::bail!(
            "ssh_user/ssh_key_path not set — no ./config.toml, and neither the global \
             config nor a cluster source provides them (see `sm121 config path`)"
        );
    }
    let key = cached_secret_key(&config.ssh_key_path)?;
    let ssh_config = russh::client::Config {
        inactivity_timeout,
        ..Default::default()
    };

    // Refuse to connect without a pinned host key — a missing `host_key`
    // would otherwise skip verification entirely (MITM vector). Run
    // `sm121 init-host-keys` to discover and pin keys, or set `host_key`
    // manually in config.toml. (`fetch_host_key` / init-host-keys uses a
    // separate capturing handler and is intentionally exempt.)
    if node.host_key.is_none() {
        anyhow::bail!(
            "No host_key set for node '{}' ({}). Refusing to connect without \
             host-key verification. Run `sm121 init-host-keys --cluster <name>` \
             to discover and pin host keys, or set `host_key` manually in config.toml.",
            node.name,
            node.ssh_ip
        );
    }
    let expected_key = Arc::new(std::sync::Mutex::new(node.host_key.clone()));

    let handler = VerifyingHandler { expected_key };
    let mut session = tokio::time::timeout(
        CONNECT_TIMEOUT,
        russh::client::connect(Arc::new(ssh_config), (node.ssh_ip.as_str(), 22), handler),
    )
    .await
    .context(format!(
        "SSH connect to {}:{} timed out after {}s",
        node.ssh_ip,
        22,
        CONNECT_TIMEOUT.as_secs()
    ))??;

    session
        .authenticate_publickey(&config.ssh_user, PrivateKeyWithHashAlg::new(key, None))
        .await?;

    Ok(session)
}

/// Runs a command on a node via SSH and returns stdout.
///
/// The remote command's exit status is tracked: if it is non-zero,
/// an Err is returned whose `to_string()` includes the captured output
/// for diagnosis.  Callers that don't care about exit status can ignore
/// the Result and inspect stdout via the returned `String`.
pub async fn run_ssh(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
) -> Result<String> {
    if dry_run {
        println!("  ssh {}@{}: {}", config.ssh_user, node.ssh_ip, command);
        return Ok(String::new());
    }

    let session = ssh_session(config, node).await?;
    let mut channel = session.channel_open_session().await?;
    channel.exec(true, command.as_bytes()).await?;

    let mut output = String::new();
    let mut exit_status: Option<u32> = None;
    while let Some(msg) = channel.wait().await {
        match msg {
            ChannelMsg::Data { data } => {
                output.push_str(&String::from_utf8_lossy(&data));
            }
            ChannelMsg::ExtendedData { data, .. } => {
                output.push_str(&String::from_utf8_lossy(&data));
            }
            ChannelMsg::ExitStatus { exit_status: s } => {
                exit_status = Some(s);
            }
            _ => {}
        }
    }

    match exit_status {
        Some(0) | None => Ok(output),
        Some(code) => Err(anyhow::anyhow!(
            "ssh {}@{}: command exited with status {}\n  command: {}\n  output:\n{}",
            config.ssh_user,
            node.ssh_ip,
            code,
            command,
            output
        )),
    }
}

/// Runs a command via SSH with exponential backoff retry.
pub async fn run_ssh_with_retry(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
    max_retries: u32,
) -> Result<String> {
    let mut last_err = None;
    for attempt in 0..=max_retries {
        match run_ssh(config, node, command.clone(), dry_run).await {
            Ok(output) => return Ok(output),
            Err(e) => {
                warn!(
                    "SSH attempt {}/{} to {} failed: {}",
                    attempt + 1,
                    max_retries + 1,
                    node.name,
                    e
                );
                last_err = Some(e);
                if attempt < max_retries {
                    sleep(Duration::from_secs(2u64.pow(attempt))).await;
                }
            }
        }
    }
    Err(last_err.unwrap())
}

/// Runs a command via SSH and streams output to stdout (for follow mode).
/// Uses a session with no inactivity timeout so long-lived streams are not killed.
pub async fn run_ssh_follow(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
) -> Result<()> {
    if dry_run {
        println!("  ssh {}@{}: {} -f", config.ssh_user, node.ssh_ip, command);
        return Ok(());
    }

    let session = ssh_session_follow(config, node).await?;
    let mut channel = session.channel_open_session().await?;
    channel.exec(true, command.as_bytes()).await?;

    while let Some(msg) = channel.wait().await {
        match msg {
            ChannelMsg::Data { data } | ChannelMsg::ExtendedData { data, .. } => {
                print!("{}", String::from_utf8_lossy(&data));
            }
            _ => {}
        }
    }
    Ok(())
}

/// Runs a command via SSH with no inactivity timeout — suitable for long,
/// streaming commands like `docker save | ssh docker load`.
pub async fn run_ssh_long_running(
    config: &Config,
    node: &Node,
    command: String,
    dry_run: bool,
) -> Result<String> {
    if dry_run {
        println!(
            "  ssh {}@{}: {} (long-running)",
            config.ssh_user, node.ssh_ip, command
        );
        return Ok(String::new());
    }
    let session = ssh_session_follow(config, node).await?;
    let mut channel = session.channel_open_session().await?;
    channel.exec(true, command.as_bytes()).await?;
    let mut output = String::new();
    let mut exit_status: Option<u32> = None;
    while let Some(msg) = channel.wait().await {
        match msg {
            ChannelMsg::Data { data } => output.push_str(&String::from_utf8_lossy(&data)),
            ChannelMsg::ExtendedData { data, .. } => {
                output.push_str(&String::from_utf8_lossy(&data))
            }
            ChannelMsg::ExitStatus { exit_status: s } => exit_status = Some(s),
            _ => {}
        }
    }
    match exit_status {
        Some(0) | None => Ok(output),
        Some(code) => Err(anyhow::anyhow!(
            "ssh {}@{}: command exited with status {}\n  command: {}",
            config.ssh_user,
            node.ssh_ip,
            code,
            command
        )),
    }
}

/// Fetches the SSH host key fingerprint from a node.
pub async fn fetch_host_key(config: &Config, node: &Node) -> Result<String> {
    let blob = fetch_host_key_data(config, node).await?;
    let hash = Sha256::digest(&blob);
    Ok(format!("SHA256:{}", STANDARD.encode(hash)))
}

/// Fetches a node's full public key line (e.g. `ssh-ed25519 AAAA...`).
/// Consulted over a connection that already verifies the pinned host key,
/// so the returned key can be trusted to write into a `known_hosts` file.
pub async fn fetch_host_pubkey(config: &Config, node: &Node) -> Result<String> {
    let blob = fetch_host_key_data(config, node).await?;
    Ok(String::from_utf8_lossy(&blob).trim().to_string())
}

/// Raw host-key bytes captured via a fresh connection (used to derive the
/// fingerprint or the full public key line).
async fn fetch_host_key_data(config: &Config, node: &Node) -> Result<Vec<u8>> {
    use std::sync::Mutex;

    let captured = Arc::new(Mutex::new(None::<Vec<u8>>));
    let captured_clone = captured.clone();

    struct CaptureHandler {
        captured: Arc<Mutex<Option<Vec<u8>>>>,
    }
    impl client::Handler for CaptureHandler {
        type Error = SshError;
        async fn check_server_key(
            &mut self,
            key: &russh::keys::PublicKey,
        ) -> Result<bool, Self::Error> {
            *self.captured.lock().unwrap() = Some(key.to_string().into_bytes());
            Ok(true)
        }
    }

    let key = cached_secret_key(&config.ssh_key_path)?;
    let ssh_config = russh::client::Config::default();
    let handler = CaptureHandler {
        captured: captured_clone,
    };
    // Apply the same connect timeout as `ssh_session_with_timeout` so an
    // unreachable node fails fast instead of hanging `health --fingerprint`.
    let mut session = tokio::time::timeout(
        CONNECT_TIMEOUT,
        russh::client::connect(Arc::new(ssh_config), (node.ssh_ip.as_str(), 22), handler),
    )
    .await
    .context(format!(
        "SSH connect to {}:{} timed out after {}s",
        node.ssh_ip,
        22,
        CONNECT_TIMEOUT.as_secs()
    ))??;

    session
        .authenticate_publickey(&config.ssh_user, PrivateKeyWithHashAlg::new(key, None))
        .await?;

    captured
        .lock()
        .unwrap()
        .take()
        .context("No server key captured")
}

pub(crate) fn shq(s: &str) -> String {
    // A value that is already a complete single-quoted shell word (e.g. a recipe
    // arg the operator pre-quoted, like `'{"json":1}'`) is self-delimiting and
    // needs no escaping: single quotes already suppress every metachar/interp,
    // so wrapping it again would corrupt the token (the old `if starts_with('''`
    // special case). Everything else is wrapped in single quotes with any
    // interior quote escaped.
    if is_single_quoted_token(s) {
        return s.to_string();
    }
    let mut out = String::with_capacity(s.len() + 2);
    out.push('\'');
    for c in s.chars() {
        if c == '\'' {
            out.push_str(r"'\''");
        } else {
            out.push(c);
        }
    }
    out.push('\'');
    out
}

/// True when `s` is a well-formed single-quoted shell word: it starts and ends
/// with a single quote and its interior contains no (unescaped) quote. A valid
/// single-quoted token cannot inject — the shell treats its whole body as one
/// literal argument.
fn is_single_quoted_token(s: &str) -> bool {
    let b = s.as_bytes();
    b.len() >= 2 && b[0] == b'\'' && b[b.len() - 1] == b'\'' && !s[1..s.len() - 1].contains('\'')
}

/// Upload an in-memory script body to a node and make it executable.
///
/// Cheaper than [`upload_dir`]: no tar, just base64-encode the body and
/// pipe through SSH for decoding into `remote_path`, then `chmod +x`. Used
/// by the built-in script runner (`before_serve` / `after_stop`) to deploy
/// a baked-in script to each node before invoking it via `bash`.
pub async fn upload_script(
    config: &Config,
    node: &Node,
    body: &str,
    remote_path: &str,
    dry_run: bool,
) -> Result<()> {
    if dry_run {
        println!(
            "  ssh {}@{}: upload script ({} bytes) -> {}",
            config.ssh_user,
            node.ssh_ip,
            body.len(),
            remote_path
        );
        return Ok(());
    }
    let b64 = STANDARD.encode(body.as_bytes());
    // `remote_path` is shell-quoted; the base64 body is safe inside single
    // quotes (charset A-Za-z0-9+/=).
    let cmd = format!(
        "printf '%s' '{}' | base64 -d > {} && chmod +x {}",
        b64,
        shq(remote_path),
        shq(remote_path)
    );
    run_ssh(config, node, cmd, false).await?;
    Ok(())
}

/// Upload a local directory to a remote node by tarring, base64-encoding,
/// and piping the result through SSH for decoding and extraction.
///
/// The remote directory is created if it doesn't exist.
pub async fn upload_dir(
    config: &Config,
    node: &Node,
    local_dir: &Path,
    remote_dir: &str,
    dry_run: bool,
) -> Result<()> {
    if dry_run {
        println!(
            "  ssh {}@{}: tar -> base64 -> decode -> extract -> {} (dry-run)",
            config.ssh_user, node.ssh_ip, remote_dir
        );
        return Ok(());
    }

    // Tar only the *contents* of the directory (not the directory itself)
    // so that files extract directly into remote_dir.
    let output = std::process::Command::new("tar")
        .args(["czf", "-", "--dereference", "-C"])
        .arg(local_dir)
        .arg(".")
        .output()
        .context(format!("Failed to tar directory '{}'", local_dir.display()))?;
    if !output.status.success() {
        anyhow::bail!(
            "tar failed on '{}': {}",
            local_dir.display(),
            String::from_utf8_lossy(&output.stderr)
        );
    }
    let b64 = STANDARD.encode(&output.stdout);

    // Base64 is safe inside single quotes (charset: A-Za-z0-9+/=).
    // `printf '%s'` is portable and doesn't interpret escape sequences.
    //
    // The full payload can exceed the platform ARG_MAX (e.g. from-source build
    // dirs with large hash-lock files -> ~250KB base64 > the ~256KB limit the
    // remote `zsh -c` shell enforces), so we never inline it all in one argv.
    // Instead stream it in fixed-size chunks via `>>` appends to a temp file on
    // the node, then base64-decode + extract in a final command.
    let tmp = format!("/tmp/sm121-up-{}.b64", std::process::id());
    const CHUNK: usize = 60_000;
    let chunks: Vec<&str> = b64
        .as_bytes()
        .chunks(CHUNK)
        .map(|c| std::str::from_utf8(c).expect("base64 is ASCII"))
        .collect();
    for (i, piece) in chunks.iter().enumerate() {
        let op = if i == 0 { ">" } else { ">>" };
        let cmd = format!("printf '%s' '{}' {} {}", piece, op, shq(&tmp));
        run_ssh(config, node, cmd, false).await?;
    }

    // `remote_dir` is shell-quoted — a path with spaces or metachars would
    // otherwise break `mkdir -p` / `cd`.
    let cmd = format!(
        "mkdir -p {dir} && cd {dir} && base64 -d < {tmp} | tar xzf - && rm -f {tmp}",
        dir = shq(remote_dir),
        tmp = shq(&tmp)
    );

    run_ssh(config, node, cmd, false).await?;
    Ok(())
}

/// Loads an SSH private key from the given path.
pub fn load_secret_key(path: &Path) -> Result<russh::keys::PrivateKey> {
    russh_load_secret_key(path, None).context("Failed to load SSH key")
}

/// Parsed private keys, cached by canonical path. The web dashboard opens a
/// fresh SSH session per node every 30 s (the fleet probe) — re-reading and
/// re-parsing the key file on every session was pure waste; the key is
/// static per deployment, so parse it once and share the `Arc`.
static KEY_CACHE: std::sync::LazyLock<
    std::sync::Mutex<std::collections::HashMap<std::path::PathBuf, Arc<russh::keys::PrivateKey>>>,
> = std::sync::LazyLock::new(|| std::sync::Mutex::new(std::collections::HashMap::new()));

fn cached_secret_key(path: &Path) -> Result<Arc<russh::keys::PrivateKey>> {
    let mut cache = KEY_CACHE.lock().unwrap();
    let key_path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
    if let Some(k) = cache.get(&key_path) {
        return Ok(Arc::clone(k));
    }
    let k = Arc::new(load_secret_key(path)?);
    cache.insert(key_path, Arc::clone(&k));
    Ok(k)
}

/// Handler that verifies the server's host key against an expected fingerprint.
pub struct VerifyingHandler {
    pub expected_key: Arc<std::sync::Mutex<Option<String>>>,
}

impl client::Handler for VerifyingHandler {
    type Error = SshError;

    async fn check_server_key(
        &mut self,
        server_public_key: &russh::keys::PublicKey,
    ) -> Result<bool, Self::Error> {
        let expected = self.expected_key.lock().unwrap().clone();
        if let Some(ref expected_fp) = expected {
            let key_bytes = server_public_key.to_string().into_bytes();
            let hash = Sha256::digest(&key_bytes);
            let fp = format!("SHA256:{}", STANDARD.encode(hash));
            if &fp != expected_fp {
                return Err(SshError::Kex);
            }
        }
        Ok(true)
    }
}

#[cfg(test)]
#[path = "ssh_tests.rs"]
mod tests;