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>> {
    let key = load_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(Arc::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::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 {
        if let ChannelMsg::Data { data } = msg {
            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::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> {
    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 = load_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(Arc::new(key), None),
        )
        .await?;

    let blob = captured
        .lock()
        .unwrap()
        .take()
        .context("No server key captured")?;
    let hash = Sha256::digest(&blob);
    let fp = format!("SHA256:{}", STANDARD.encode(hash));
    Ok(fp)
}

pub(crate) fn shq(s: &str) -> 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
}

/// 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.
    // `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} && printf '%s' '{b64}' | base64 -d | tar xzf -",
        dir = shq(remote_dir),
        b64 = b64
    );

    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")
}

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