path: root/tests/machine_setup.rs

#![cfg(unix)]

use std::fs;
use std::io::{Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::process::Command;
use std::sync::mpsc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use portable_pty::{native_pty_system, CommandBuilder, PtySize};

// No real SSH connection or server is started. Stop at startup after recording setup actions.
const SSH: &str = r#"#!/bin/sh
for arg do
    last=$arg
    if [ "$arg" = 'StrictHostKeyChecking=yes' ]; then strict_host_key_check=yes; fi
done
if [ "$FAKE_STRICT_HOST_KEY_FAILURE" = yes ] && [ "$strict_host_key_check" = yes ]; then
    echo 'No RSA host key is known for fake-host and you have requested strict checking.' >&2
    echo 'Host key verification failed.' >&2
    exit 255
fi
if [ "$last" = 'command -v herdr' ]; then
    if [ "$FAKE_INSTALLED" != missing ]; then echo /home/remote/.local/bin/herdr; fi
    exit 0
fi
case "$last" in
    'tee '*) cat >/dev/null; exit 0 ;;
esac
case "$last" in
    *'herdr-remote-output-ready:1'*) script=$last ;;
    *) script=$(cat) ;;
esac
printf '\n%s\n' 'herdr-remote-output-ready:1'
case "$script" in
    *'uname -s'*) echo platform >>"$FAKE_ROOT/probes"; uname -s; uname -m ;;
    *'version='*)
        echo candidates >>"$FAKE_ROOT/probes"
        if [ "$FAKE_INSTALLED" != missing ]; then echo /home/remote/.local/bin/herdr; fi ;;
    *'status client --json'*)
        if [ "$FAKE_INSTALLED" = new ] || [ -f "$FAKE_ROOT/installed" ]; then
            printf '%s\n' "$FAKE_CLIENT_STATUS"
        else
            echo '{"version":"0.8.2","protocol":20}'
        fi ;;
    *'status server --json'*)
        if [ -f "$FAKE_ROOT/stopped" ]; then
            echo '{"running":false}'
        elif [ "$FAKE_STRICT_HOST_KEY_FAILURE" = yes ]; then
            echo '{"running":true,"version":"0.8.2","capabilities":{"live_handoff":true,"detached_server_daemon":true,"endpoint_protocol_generation":1,"surface_interest":true,"health_check":true}}'
        else
            echo '{"running":true,"version":"0.8.2","capabilities":{"live_handoff":true,"detached_server_daemon":true}}'
        fi ;;
    *'server live-handoff'*) echo handoff >>"$FAKE_ROOT/actions"; exit 1 ;;
    *'server stop'*) echo stop >>"$FAKE_ROOT/actions"; touch "$FAKE_ROOT/stopped" ;;
    *'session list --json'*)
        echo sessions >>"$FAKE_ROOT/probes"
        if [ "$FAKE_SESSIONS" = failed ]; then echo 'session query denied' >&2; exit 1; fi
        printf '%s\n' "$FAKE_SESSIONS" ;;
    *'remote-client-bridge'*)
        printf '%s\n' "$script" >"$FAKE_ROOT/start-command"
        echo start >>"$FAKE_ROOT/actions"; echo 'test startup failure' >&2; exit 1 ;;
    *'mkdir -p'*) printf '/fake/tmp\000/fake/herdr\000' ;;
    *'chmod 755'*) echo install >>"$FAKE_ROOT/actions"; touch "$FAKE_ROOT/installed" ;;
    *'command -v herdr'*) if [ "$FAKE_INSTALLED" != missing ]; then echo /home/remote/.local/bin/herdr; fi ;;
    *'test -x '*) exit 1 ;;
    *) echo "unexpected fake SSH script: $script" >&2; exit 1 ;;
esac
"#;

struct SetupResult {
    output: String,
    actions: String,
    prompts: usize,
    success: bool,
    probes: String,
    start_command: String,
}

#[derive(Default)]
struct SetupOptions<'a> {
    sessions: Option<&'a str>,
    selection: &'a [u8],
    explicit_session: Option<&'a str>,
    noninteractive: bool,
}

fn setup(installed: &str, answer: &str, handoff: bool) -> SetupResult {
    setup_with_strict_host_key_failure(installed, answer, handoff, false)
}

fn setup_with_strict_host_key_failure(
    installed: &str,
    answer: &str,
    handoff: bool,
    strict_host_key_failure: bool,
) -> SetupResult {
    setup_options(
        installed,
        answer,
        handoff,
        strict_host_key_failure,
        SetupOptions::default(),
    )
}

fn setup_options(
    installed: &str,
    answer: &str,
    handoff: bool,
    strict_host_key_failure: bool,
    options: SetupOptions<'_>,
) -> SetupResult {
    let root = std::path::PathBuf::from(format!(
        "/var/tmp/herdr-machine-setup-{}-{}-{}-{}-{}-{}",
        std::process::id(),
        installed,
        answer.trim().is_empty(),
        handoff,
        strict_host_key_failure,
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let app = if cfg!(debug_assertions) {
        "herdr-dev"
    } else {
        "herdr"
    };
    fs::create_dir_all(root.join("bin")).unwrap();
    fs::create_dir_all(root.join("config").join(app)).unwrap();
    fs::write(root.join("bin/ssh"), SSH).unwrap();
    fs::set_permissions(root.join("bin/ssh"), fs::Permissions::from_mode(0o700)).unwrap();
    fs::write(
        root.join("config").join(app).join("config.toml"),
        "onboarding = false\n[remote]\nmanage_ssh_config = false\n",
    )
    .unwrap();
    let status = Command::new(env!("CARGO_BIN_EXE_herdr"))
        .args(["status", "client", "--json"])
        .output()
        .unwrap();
    assert!(status.status.success());

    let pair = native_pty_system().openpty(PtySize::default()).unwrap();
    let mut command = CommandBuilder::new(env!("CARGO_BIN_EXE_herdr"));
    if handoff {
        command.args(["--remote", "fake-host", "--handoff"]);
    } else {
        command.args(["machine", "add", "fake-host", "--label", "VPS"]);
        if let Some(session) = options.explicit_session {
            command.args(["--remote-session", session]);
        }
    }
    command.env(
        "PATH",
        format!("{}:/usr/bin:/bin", root.join("bin").display()),
    );
    command.env("HOME", &root);
    command.env("XDG_CONFIG_HOME", root.join("config"));
    command.env("XDG_STATE_HOME", root.join("state"));
    command.env("XDG_RUNTIME_DIR", &root);
    command.env("FAKE_ROOT", &root);
    command.env("FAKE_INSTALLED", installed);
    command.env(
        "FAKE_SESSIONS",
        options
            .sessions
            .unwrap_or(r#"{"sessions":[{"name":"default","running":true}]}"#),
    );
    command.env(
        "FAKE_STRICT_HOST_KEY_FAILURE",
        if strict_host_key_failure { "yes" } else { "no" },
    );
    command.env(
        "FAKE_CLIENT_STATUS",
        String::from_utf8(status.stdout).unwrap(),
    );
    for name in [
        "HERDR_ENV",
        "HERDR_SESSION",
        "HERDR_SOCKET_PATH",
        "HERDR_CLIENT_SOCKET_PATH",
        "HERDR_REMOTE_BINARY",
        "HERDR_CONFIG_PATH",
    ] {
        command.env_remove(name);
    }
    if options.noninteractive {
        // A redirected stdin makes machine add non-interactive while retaining
        // the same isolated SSH fixture and captured terminal diagnostics.
        let mut shell = CommandBuilder::new("/bin/sh");
        shell.args(["-c", "exec \"$@\" </dev/null", "sh"]);
        shell.args(command.get_argv());
        shell.env_clear();
        for (key, value) in command.iter_full_env_as_str() {
            shell.env(key, value);
        }
        command = shell;
    }
    let mut child = pair.slave.spawn_command(command).unwrap();
    drop(pair.slave);
    let mut reader = pair.master.try_clone_reader().unwrap();
    let mut writer = pair.master.take_writer().unwrap();
    let (tx, rx) = mpsc::channel();
    let reading = std::thread::spawn(move || {
        let mut buffer = [0; 4096];
        while let Ok(len) = reader.read(&mut buffer) {
            if len == 0 || tx.send(buffer[..len].to_vec()).is_err() {
                break;
            }
        }
    });
    let mut output = String::new();
    let mut prompts = 0;
    let mut timed_out = false;
    let mut picked = false;
    loop {
        match rx.recv_timeout(Duration::from_secs(20)) {
            Ok(bytes) => {
                output.push_str(&String::from_utf8_lossy(&bytes));
                if !picked && output.contains("Esc cancel") {
                    writer.write_all(options.selection).unwrap();
                    writer.flush().unwrap();
                    picked = true;
                }
                let count = output.matches("[y/N]").count() + output.matches("[Y/n]").count();
                if count > prompts {
                    writer
                        .write_all(if prompts == 0 {
                            answer.as_bytes()
                        } else {
                            b"n\n"
                        })
                        .unwrap();
                    writer.flush().unwrap();
                    prompts = count;
                }
            }
            Err(mpsc::RecvTimeoutError::Disconnected) => break,
            Err(mpsc::RecvTimeoutError::Timeout) => {
                timed_out = true;
                child.kill().unwrap();
                break;
            }
        }
    }
    let status = child.wait().unwrap();
    drop(writer);
    drop(pair.master);
    reading.join().unwrap();
    let actions = fs::read_to_string(root.join("actions")).unwrap_or_default();
    let probes = fs::read_to_string(root.join("probes")).unwrap_or_default();
    let start_command = fs::read_to_string(root.join("start-command")).unwrap_or_default();
    // Every scenario either cancels or reaches the intentionally failed startup.
    let saved = root
        .join("state")
        .join(app)
        .join("client/endpoints.json")
        .exists();
    fs::remove_dir_all(root).unwrap();
    assert!(
        !saved,
        "failed or cancelled setup saved the machine: {output}"
    );
    assert!(!timed_out, "setup timed out: {output}");
    SetupResult {
        output,
        actions,
        prompts,
        success: status.success(),
        probes,
        start_command,
    }
}

#[test]
fn machine_add_accepts_help_argument_order() {
    let nonce = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_nanos();
    let root = std::env::temp_dir().join(format!(
        "herdr-machine-add-order-{}-{nonce}",
        std::process::id()
    ));
    fs::create_dir(&root).unwrap();
    let app = if cfg!(debug_assertions) {
        "herdr-dev"
    } else {
        "herdr"
    };
    fs::create_dir_all(root.join("config").join(app)).unwrap();
    fs::write(
        root.join("config").join(app).join("config.toml"),
        "onboarding = false\n[remote]\nmanage_ssh_config = false\n",
    )
    .unwrap();
    let mut command = Command::new(env!("CARGO_BIN_EXE_herdr"));
    command.args(["machine", "add", "--label", "coder", "workstation.coder"]);
    // Reach remote preparation, but never execute SSH or start a server.
    command.env("PATH", root.join("no-executables"));
    command.env("HOME", &root);
    command.env("XDG_CONFIG_HOME", root.join("config"));
    command.env("XDG_STATE_HOME", root.join("state"));
    command.env("XDG_RUNTIME_DIR", &root);
    for name in [
        "HERDR_ENV",
        "HERDR_SESSION",
        "HERDR_SOCKET_PATH",
        "HERDR_CLIENT_SOCKET_PATH",
        "HERDR_REMOTE_BINARY",
        "HERDR_CONFIG_PATH",
    ] {
        command.env_remove(name);
    }
    let output = command.output().unwrap();
    let saved = root
        .join("state")
        .join(app)
        .join("client/endpoints.json")
        .exists();
    fs::remove_dir_all(root).unwrap();
    let stderr = String::from_utf8_lossy(&output.stderr);
    assert_eq!(output.status.code(), Some(1), "{stderr}");
    assert!(stderr.contains("machine was not saved"), "{stderr}");
    assert!(!saved, "failed preparation must not save a machine");
}

#[test]
fn machine_add_fresh_host_offers_installation_and_cancellation_saves_nothing() {
    let result = setup("missing", "n\n", false);
    assert_eq!(result.probes, "platform\ncandidates\n", "{}", result.output);
    assert_eq!(result.prompts, 1, "{}", result.output);
    assert!(
        result.output.contains("installation cancelled"),
        "{}",
        result.output
    );
    assert!(result.actions.is_empty(), "{}", result.output);
}

#[test]
fn machine_add_discovers_sessions_once_and_prepares_the_selected_session() {
    for (sessions, selection, expected) in [
        (r#"{"sessions":[]}"#, &b""[..], "default"),
        (
            r#"{"extra":1,"sessions":[{"name":"stopped","running":false},{"name":"bad/name","running":true},{"name":"agents","running":true,"future":true}]}"#,
            &b""[..],
            "agents",
        ),
        (
            r#"{"sessions":[{"name":"default","running":true},{"name":"agents","running":true}]}"#,
            &b"\x1b[B\r"[..],
            "agents",
        ),
    ] {
        let result = setup_options(
            "new",
            "y\n",
            false,
            false,
            SetupOptions {
                sessions: Some(sessions),
                selection,
                ..Default::default()
            },
        );
        assert_eq!(
            result.probes, "platform\ncandidates\nsessions\n",
            "{}",
            result.output
        );
        assert_eq!(result.actions, "stop\nstart\n", "{}", result.output);
        if expected == "default" {
            assert!(
                !result.start_command.contains("--session"),
                "{}",
                result.start_command
            );
        } else {
            assert!(
                result.start_command.contains("--session agents"),
                "{}",
                result.start_command
            );
        }
        if !selection.is_empty() {
            let picker = result
                .output
                .split("Running sessions on fake-host:")
                .nth(1)
                .unwrap();
            assert!(
                picker.starts_with("\r\n"),
                "picker did not return to column zero: {picker:?}"
            );
            assert!(picker.contains("default\u{1b}[0m\r\n"), "{picker:?}");
            assert!(picker.contains("  agents\r\n"), "{picker:?}");
        }
    }
}

#[test]
fn machine_add_failed_discovery_and_cancel_never_prepare_or_save() {
    for (sessions, selection, expected) in [
        ("failed", &b""[..], "session query denied"),
        ("not json", &b""[..], "--remote-session"),
        (
            r#"{"sessions":[{"name":"agents"}]}"#,
            &b""[..],
            "--remote-session",
        ),
        (
            r#"{"sessions":[{"name":"default","running":true},{"name":"agents","running":true}]}"#,
            &b"\x1b"[..],
            "selection cancelled",
        ),
    ] {
        let result = setup_options(
            "new",
            "y\n",
            false,
            false,
            SetupOptions {
                sessions: Some(sessions),
                selection,
                ..Default::default()
            },
        );
        assert!(!result.success, "{}", result.output);
        assert!(result.actions.is_empty(), "{}", result.output);
        assert!(result.output.contains(expected), "{}", result.output);
        assert!(
            result.output.contains("machine was not saved"),
            "{}",
            result.output
        );
    }
}

#[test]
fn machine_add_skips_discovery_for_scripts_and_explicit_sessions() {
    for (noninteractive, explicit_session) in [
        (true, None),
        (true, Some("agents")),
        (false, Some("agents")),
    ] {
        let result = setup_options(
            "new",
            "y\n",
            false,
            false,
            SetupOptions {
                noninteractive,
                explicit_session,
                sessions: Some("failed"),
                ..Default::default()
            },
        );
        assert_eq!(result.probes, "platform\ncandidates\n", "{}", result.output);
        assert!(
            result.output.contains(if explicit_session.is_some() {
                "session agents"
            } else {
                "session default"
            }),
            "{}",
            result.output
        );
        if noninteractive {
            assert!(result.actions.is_empty(), "{}", result.output);
        } else {
            assert_eq!(result.actions, "stop\nstart\n", "{}", result.output);
        }
    }
}

#[test]
fn machine_add_reports_strict_host_key_failure() {
    let result = setup_with_strict_host_key_failure("new", "", false, true);
    assert!(!result.success, "{}", result.output);
    assert_eq!(result.prompts, 0, "{}", result.output);
    assert!(result.actions.is_empty(), "{}", result.output);
    assert!(
        result.output.contains("Host key verification failed"),
        "{}",
        result.output
    );
    assert!(
        result
            .output
            .contains("saved machines use strict host-key checking"),
        "{}",
        result.output
    );
    assert!(
        !result.output.contains("lost connection to server"),
        "{}",
        result.output
    );
}

#[test]
fn machine_setup_old_install_requires_one_explicit_stop_approval() {
    let result = setup("old", "y\n", false);
    assert_eq!(result.prompts, 1, "{}", result.output);
    assert!(
        result
            .output
            .contains("This stops active remote pane processes"),
        "{}",
        result.output
    );
    assert_eq!(
        result.actions, "install\nstop\nstart\n",
        "{}",
        result.output
    );
}

#[test]
fn machine_setup_old_install_enter_cancels_without_remote_changes() {
    let result = setup("old", "\n", false);
    assert!(result.actions.is_empty(), "{}", result.output);
    assert!(
        result.output.contains("machine was not saved"),
        "{}",
        result.output
    );
}

#[test]
fn machine_setup_new_install_old_server_enter_does_not_stop_or_handoff() {
    let result = setup("new", "\n", false);
    assert!(result.actions.is_empty(), "{}", result.output);
    assert!(
        result
            .output
            .contains("This stops active remote pane processes"),
        "{}",
        result.output
    );
}

#[test]
fn machine_setup_new_install_old_server_approval_stops_then_starts_without_installing() {
    let result = setup("new", "y\n", false);
    assert_eq!(result.prompts, 1, "{}", result.output);
    assert_eq!(result.actions, "stop\nstart\n", "{}", result.output);
}

#[test]
fn machine_setup_explicit_remote_handoff_remains_available_but_restart_fallback_defaults_to_no() {
    let result = setup("new", "\n", true);
    assert_eq!(result.actions, "handoff\n", "{}", result.output);
    assert!(
        result
            .output
            .contains("This stops active remote pane processes"),
        "{}",
        result.output
    );
}