ai-app/server/src/ssh.rs

//! Building the command a driver actually spawns -- locally, or wrapped in
//! `ssh` when the session names a host to run on.
//!
//! The whole point of the session design is that a driver speaks JSONL over
//! a child process's stdio and doesn't care what that child is. A remote
//! session is therefore the identical command with `ssh host …` in front:
//! stdio doesn't care, so nothing downstream of here changes.
//!
//! Uses the system `ssh` client rather than a Rust SSH library, so
//! `~/.ssh/config`, agents, and jump hosts all keep working and there is
//! only one place to configure connections (PLAN.md, rule 23).

use std::path::Path;
use std::process::Stdio;

use tokio::process::Command;

use crate::config::SshConfig;

/// Options forced onto every connection. `BatchMode` makes a missing key
/// fail immediately with a readable message instead of hanging on a
/// password prompt that nothing can answer; the keepalives turn a silently
/// dropped link into a process exit, which the session reports as `exited`
/// rather than appearing to hang forever.
const SSH_OPTIONS: [&str; 3] = [
    "BatchMode=yes",
    "ServerAliveInterval=30",
    "ServerAliveCountMax=3",
];

/// Builds the child process for `program args…`, run in `cwd`, either on
/// this machine (`ssh` absent) or on the machine it describes.
pub fn command(
    remote: Option<&SshConfig>,
    program: &str,
    args: &[String],
    cwd: Option<&Path>,
) -> Command {
    let Some(ssh) = remote else {
        let mut command = Command::new(program);
        command.args(args);
        if let Some(cwd) = cwd {
            command.current_dir(cwd);
        }
        return configure(command);
    };

    let mut command = Command::new("ssh");
    // -T: no pty. This carries JSONL, and a pty would rewrite it (echo,
    // CRLF translation, ^C handling) into something the parser can't read.
    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);
        // Without this, ssh may offer an agent key first and authenticate
        // as somebody else entirely -- silently, and with different
        // permissions than intended.
        command.args(["-o", "IdentitiesOnly=yes"]);
    }
    command.arg(&ssh.address);
    command.arg(remote_script(program, args, cwd));
    configure(command)
}

fn configure(mut command: Command) -> Command {
    command
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .kill_on_drop(true);
    command
}

/// The single argument handed to the remote login shell.
///
/// `exec` so the CLI replaces that shell: the process the connection is
/// attached to is then the CLI itself, and dropping the connection takes
/// it down rather than leaving an orphan behind a live wrapper.
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(&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
}

/// Single-quotes one word for a POSIX shell.
///
/// Everything crossing to the remote side goes through here: paths, model
/// names, and prompts-as-arguments are all attacker-adjacent input in a
/// server whose whole job is running commands, and unquoted they would be
/// shell syntax rather than data.
fn quote(word: &str) -> String {
    // Inside single quotes every character is literal except `'` itself,
    // which is closed, escaped, and reopened.
    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()
    }

    /// The rendered argv, for asserting on what would actually run.
    fn argv(command: &Command) -> Vec<String> {
        let std = command.as_std();
        std::iter::once(std.get_program())
            .chain(std.get_args())
            .map(|arg| arg.to_string_lossy().into_owned())
            .collect()
    }

    /// A host with nothing configured but a name to dial, so `~/.ssh/config`
    /// decides everything else -- the case that proves this adds no flags of
    /// its own when it was not told to.
    fn bare_host() -> SshConfig {
        SshConfig {
            address: "vm".to_string(),
            port: None,
            identity_file: None,
            options: vec![],
        }
    }

    #[test]
    fn a_session_with_no_host_runs_the_command_directly() {
        let command = command(
            None,
            "claude",
            &args(["-p", "--verbose"]),
            Some(Path::new("/tmp/x")),
        );
        assert_eq!(argv(&command), ["claude", "-p", "--verbose"]);
        assert_eq!(
            command.as_std().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()],
        };
        let rendered = argv(&command(
            Some(&ssh),
            "claude",
            &args(["-p", "--model", "haiku"]),
            Some(Path::new("/home/bob/work")),
        ));

        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()));
        // The host, then exactly one argument: the remote script.
        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));
        assert_eq!(rendered.last().unwrap(), "exec 'claude' '-p'");
        // No -i means no IdentitiesOnly: ~/.ssh/config decides instead.
        assert!(!rendered.contains(&"IdentitiesOnly=yes".to_string()));
    }

    #[test]
    fn shell_metacharacters_cross_as_data_not_syntax() {
        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'");

        // The end-to-end version of the same worry: a working directory
        // that tries to close the quote and start a new command.
        let ssh = bare_host();
        let evil = Path::new("/tmp/'; touch /tmp/pwned; '");
        let rendered = argv(&command(Some(&ssh), "claude", &[], Some(evil)));
        let script = rendered.last().unwrap();
        assert_eq!(
            script,
            r"cd '/tmp/'\''; touch /tmp/pwned; '\''' && exec 'claude'"
        );
        assert!(!script.contains("; touch /tmp/pwned; '\" "));
    }
}