ai-app/server/src/ssh.rs

use std::path::{Path, PathBuf};
use std::process::Command;

use crate::config::SshConfig;

/// The second half of what a transport is (PLAN.md's SSH section): "run this"
/// plus "reach this port". Locally the two numbers are one and nothing is
/// forwarded; over ssh the connection carries an `-L` tunnel, so a model server
/// binds loopback on the far machine and is never exposed to its network.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Forward {
    pub there: u16,
    pub here: u16,
}

const SSH_OPTIONS: [&str; 3] = [
    "BatchMode=yes",
    "ServerAliveInterval=30",
    "ServerAliveCountMax=3",
];

/// Stdio is left alone: how the streams are connected is the caller's decision
/// and differs by more than the transport does -- a probe wants pipes it will
/// drain, a session wants files that outlive this server.
///
/// A plain [`std::process::Command`], which `tokio` converts from, because not
/// every caller is async: the usage fetch is blocking by nature and should not
/// have to build an ssh invocation of its own.
pub fn command(
    remote: Option<&SshConfig>,
    program: &str,
    args: &[String],
    cwd: Option<&Path>,
    forward: Option<Forward>,
) -> Command {
    let Some(ssh) = remote else {
        let mut command = Command::new(program);
        command.args(args);
        if let Some(cwd) = cwd {
            command.current_dir(expand_home(cwd));
        }
        return command;
    };

    let mut command = Command::new("ssh");
    if let Some(forward) = forward {
        command.arg("-tt");
        command.args([
            "-L",
            &format!("127.0.0.1:{}:127.0.0.1:{}", forward.here, forward.there),
        ]);
        // Without this a forward that cannot be set up is a warning on stderr
        // and a session that runs anyway, answering nothing -- which would
        // arrive as "the model never became ready".
        command.args(["-o", "ExitOnForwardFailure=yes"]);
    } else {
        command.arg("-T");
    }
    for option in SSH_OPTIONS {
        command.args(["-o", option]);
    }
    for option in &ssh.options {
        command.args(["-o", option]);
    }
    if let Some(port) = ssh.port {
        command.args(["-p", &port.to_string()]);
    }
    if let Some(identity) = &ssh.identity_file {
        command.arg("-i").arg(identity);
        command.args(["-o", "IdentitiesOnly=yes"]);
    }
    command.arg(&ssh.address);
    command.arg(remote_script(program, args, cwd));
    command
}

fn remote_script(program: &str, args: &[String], cwd: Option<&Path>) -> String {
    let mut script = String::new();
    if let Some(cwd) = cwd {
        script.push_str("cd ");
        script.push_str(&quote_path(&cwd.to_string_lossy()));
        script.push_str(" && ");
    }
    script.push_str("exec ");
    script.push_str(&quote(program));
    for arg in args {
        script.push(' ');
        script.push_str(&quote(arg));
    }
    script
}

pub(crate) fn expand_home(path: &Path) -> PathBuf {
    let Some(rest) = path.to_str().and_then(|p| {
        if p == "~" {
            Some("")
        } else {
            p.strip_prefix("~/")
        }
    }) else {
        return path.to_path_buf();
    };
    match std::env::home_dir() {
        Some(home) => home.join(rest),
        None => path.to_path_buf(),
    }
}

/// `"$HOME"` rather than handing the tilde to the shell unquoted: the variable
/// is expanded, the expansion is not re-split or globbed because it is
/// double-quoted, and everything after it stays single-quoted and literal.
/// `$HOME` is set by every shell this can land in, including the fish login
/// shell on the dev VM, so this does not depend on the remote shell being POSIX.
pub(crate) fn quote_path(path: &str) -> String {
    if path == "~" {
        return "\"$HOME\"".to_string();
    }
    match path.strip_prefix("~/") {
        Some(rest) => format!("\"$HOME\"/{}", quote(rest)),
        None => quote(path),
    }
}

pub(crate) fn quote(word: &str) -> String {
    format!("'{}'", word.replace('\'', r"'\''"))
}

#[cfg(test)]
mod tests {
    use super::*;

    fn args<const N: usize>(args: [&str; N]) -> Vec<String> {
        args.iter().map(|arg| arg.to_string()).collect()
    }

    fn argv(command: &Command) -> Vec<String> {
        std::iter::once(command.get_program())
            .chain(command.get_args())
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect()
    }

    fn bare_host() -> SshConfig {
        SshConfig {
            address: "vm".to_string(),
            port: None,
            identity_file: None,
            options: vec![],
            models_dir: None,
            attachments_dir: None,
        }
    }

    #[test]
    fn a_session_with_no_host_runs_the_command_directly() {
        let command = command(
            None,
            "claude",
            &args(["-p", "--verbose"]),
            Some(Path::new("/tmp/x")),
            None,
        );
        assert_eq!(argv(&command), ["claude", "-p", "--verbose"]);
        assert_eq!(command.get_current_dir(), Some(Path::new("/tmp/x")));
    }

    #[test]
    fn a_session_with_a_host_wraps_the_same_command_in_ssh() {
        let ssh = SshConfig {
            address: "bob@10.0.2.15".to_string(),
            port: Some(2222),
            identity_file: Some("/home/me/.ssh/id_ai".into()),
            options: vec!["StrictHostKeyChecking=accept-new".to_string()],
            models_dir: None,
            attachments_dir: None,
        };
        let rendered = argv(&command(
            Some(&ssh),
            "claude",
            &args(["-p", "--model", "haiku"]),
            Some(Path::new("/home/bob/work")),
            None,
        ));

        assert_eq!(rendered[0], "ssh");
        assert!(rendered.contains(&"-T".to_string()));
        assert!(rendered.contains(&"BatchMode=yes".to_string()));
        assert!(rendered.contains(&"StrictHostKeyChecking=accept-new".to_string()));
        assert!(rendered.contains(&"IdentitiesOnly=yes".to_string()));
        assert!(rendered.contains(&"2222".to_string()));
        assert!(rendered.contains(&"/home/me/.ssh/id_ai".to_string()));
        assert_eq!(rendered[rendered.len() - 2], "bob@10.0.2.15");
        assert_eq!(
            rendered[rendered.len() - 1],
            "cd '/home/bob/work' && exec 'claude' '-p' '--model' 'haiku'",
        );
    }

    #[test]
    fn a_remote_command_without_a_cwd_just_execs() {
        let ssh = bare_host();
        let rendered = argv(&command(Some(&ssh), "claude", &args(["-p"]), None, None));
        assert_eq!(rendered.last().unwrap(), "exec 'claude' '-p'");
        assert!(!rendered.contains(&"IdentitiesOnly=yes".to_string()));
    }

    #[test]
    fn a_forwarded_port_rides_the_same_connection_as_the_command() {
        let ssh = bare_host();
        let rendered = argv(&command(
            Some(&ssh),
            "llama-server",
            &args(["--port", "24242"]),
            None,
            Some(Forward {
                there: 24242,
                here: 41000,
            }),
        ));
        let forward = rendered
            .iter()
            .position(|arg| arg == "-L")
            .expect("a forward");
        assert_eq!(rendered[forward + 1], "127.0.0.1:41000:127.0.0.1:24242");
        assert!(rendered.contains(&"ExitOnForwardFailure=yes".to_string()));
        assert!(rendered.contains(&"-tt".to_string()));
        assert!(!rendered.contains(&"-T".to_string()));
        assert!(forward < rendered.len() - 2);
        assert_eq!(
            rendered.last().unwrap(),
            "exec 'llama-server' '--port' '24242'"
        );

        let plain = argv(&command(Some(&ssh), "claude", &args(["-p"]), None, None));
        assert!(!plain.contains(&"-L".to_string()));
        assert!(plain.contains(&"-T".to_string()));
        assert!(!plain.contains(&"-tt".to_string()));
    }

    #[test]
    fn a_leading_tilde_expands_and_nothing_else_does() {
        assert_eq!(quote_path("~"), "\"$HOME\"");
        assert_eq!(quote_path("~/repos/ai-app"), "\"$HOME\"/'repos/ai-app'");
        assert_eq!(quote_path("/tmp/~/x"), "'/tmp/~/x'");
        assert_eq!(quote_path("~user/x"), "'~user/x'");

        assert_eq!(
            remote_script("claude", &args(["-p"]), Some(Path::new("~/repos/ai-app"))),
            "cd \"$HOME\"/'repos/ai-app' && exec 'claude' '-p'",
        );
    }

    #[test]
    fn a_local_cwd_expands_its_tilde_the_same_way() {
        let Some(home) = std::env::home_dir() else {
            return;
        };
        assert_eq!(
            expand_home(Path::new("~/repos/ai-app")),
            home.join("repos/ai-app")
        );
        assert_eq!(expand_home(Path::new("~")), home);
        assert_eq!(expand_home(Path::new("/tmp/~/x")), Path::new("/tmp/~/x"));
        assert_eq!(expand_home(Path::new("~user/x")), Path::new("~user/x"));

        let local = command(
            None,
            "claude",
            &args(["-p"]),
            Some(Path::new("~/work")),
            None,
        );
        assert_eq!(local.get_current_dir(), Some(home.join("work").as_path()));
    }

    #[test]
    fn shell_metacharacters_cross_as_data_not_syntax() {
        assert_eq!(
            quote_path("~/'; touch /tmp/pwned; '"),
            r#""$HOME"/''\''; touch /tmp/pwned; '\'''"#,
        );

        assert_eq!(quote("plain"), "'plain'");
        assert_eq!(quote("with space"), "'with space'");
        assert_eq!(quote("; rm -rf /"), "'; rm -rf /'");
        assert_eq!(quote("$(whoami)"), "'$(whoami)'");
        assert_eq!(quote("it's"), r"'it'\''s'");

        let ssh = bare_host();
        let evil = Path::new("/tmp/'; touch /tmp/pwned; '");
        let rendered = argv(&command(Some(&ssh), "claude", &[], Some(evil), None));
        let script = rendered.last().unwrap();
        assert_eq!(
            script,
            r"cd '/tmp/'\''; touch /tmp/pwned; '\''' && exec 'claude'"
        );
        assert!(!script.contains("; touch /tmp/pwned; '\" "));
    }
}