agent-sandbox/tests/e2e/terminal.rs

use std::fs;
use std::path::Path;
use std::thread;
use std::time::Duration;

use tempfile::TempDir;

use crate::common::*;

const BASH: &[&str] = &["bash", "--norc", "--noprofile", "-i"];
const ZSH: &[&str] = &["zsh", "-f", "-i"];

#[test]
fn ctrl_c_lets_entrypoint_clean_up_in_bash() {
    assert_ctrl_c_lets_entrypoint_clean_up(BASH);
}

#[test]
fn ctrl_c_lets_entrypoint_clean_up_in_zsh() {
    assert_ctrl_c_lets_entrypoint_clean_up(ZSH);
}

fn assert_ctrl_c_lets_entrypoint_clean_up(shell: &[&str]) {
    let dir = TempDir::new().expect("failed to create temp dir");
    let script = "trap 'sleep 0.5; touch cleaned; exit 0' INT
                  touch ready
                  sleep 60 & wait";

    let mut terminal = Terminal::start(shell);
    terminal.type_line(&build_sandbox_command(dir.path(), script));
    wait_for_file(&dir.path().join("ready"));
    terminal.wait_until_job_owns_terminal();
    terminal.type_text(CTRL_C);
    terminal.type_line("echo STATUS-$?");
    terminal.wait_for_output("STATUS-0");

    assert!(
        dir.path().join("cleaned").exists(),
        "cleanup did not finish"
    );
}

const CTRL_C: &str = "\x03";
const CTRL_Z: &str = "\x1a";

fn build_sandbox_command(dir: &Path, script: &str) -> String {
    let args = build_sandbox_args(dir, script);
    shlex::try_join(args.iter().map(String::as_str)).unwrap()
}

fn build_sandbox_args(dir: &Path, script: &str) -> Vec<String> {
    let dir = dir.to_str().unwrap();
    [
        env!("CARGO_BIN_EXE_agent-sandbox"),
        "--no-config",
        "--blacklist",
        "--chdir",
        dir,
        "--",
        "sh",
        "-c",
        script,
    ]
    .map(String::from)
    .to_vec()
}

#[test]
fn ctrl_z_and_fg_resume_entrypoint_in_bash() {
    assert_ctrl_z_and_fg_resume_entrypoint(BASH);
}

#[test]
fn ctrl_z_and_fg_resume_entrypoint_in_zsh() {
    assert_ctrl_z_and_fg_resume_entrypoint(ZSH);
}

fn assert_ctrl_z_and_fg_resume_entrypoint(shell: &[&str]) {
    let dir = TempDir::new().expect("failed to create temp dir");
    let script = "touch ready
                  read line
                  echo \"got-$line\"";

    let mut terminal = Terminal::start(shell);
    terminal.type_line(&build_sandbox_command(dir.path(), script));
    wait_for_file(&dir.path().join("ready"));
    terminal.wait_until_job_owns_terminal();
    terminal.type_text(CTRL_Z);
    terminal.wait_until_shell_owns_terminal();
    terminal.type_line("fg");
    terminal.wait_until_job_owns_terminal();
    terminal.type_line("hello");
    terminal.wait_for_output("got-hello");
    terminal.type_line("echo STATUS-$?");

    terminal.wait_for_output("STATUS-0");
}

#[test]
fn ctrl_c_ends_surrounding_bash_script() {
    let dir = TempDir::new().expect("failed to create temp dir");
    let script = ": > ready
                  sleep 60";
    let surrounding_script = format!(
        "{}; echo AFTER-$((1 + 1))",
        build_sandbox_command(dir.path(), script)
    );

    let mut terminal = Terminal::start(BASH);
    terminal.type_line(&format!(
        "bash -c {}",
        shlex::try_quote(&surrounding_script).unwrap()
    ));
    wait_for_file(&dir.path().join("ready"));
    terminal.wait_until_job_owns_terminal();
    terminal.type_text(CTRL_C);
    terminal.type_line("echo STATUS-$?");
    terminal.wait_for_output("STATUS-130");

    assert!(
        !terminal.read_output().contains("AFTER-2"),
        "surrounding script continued after Ctrl-C"
    );
}

#[test]
fn ctrl_z_and_bg_continue_entrypoint_in_bash() {
    assert_ctrl_z_and_bg_continue_entrypoint(BASH);
}

#[test]
fn ctrl_z_and_bg_continue_entrypoint_in_zsh() {
    assert_ctrl_z_and_bg_continue_entrypoint(ZSH);
}

fn assert_ctrl_z_and_bg_continue_entrypoint(shell: &[&str]) {
    let dir = TempDir::new().expect("failed to create temp dir");
    let script = "touch ready
                  while [ ! -e go ]; do sleep 0.01; done
                  touch done";

    let mut terminal = Terminal::start(shell);
    terminal.type_line(&build_sandbox_command(dir.path(), script));
    wait_for_file(&dir.path().join("ready"));
    terminal.wait_until_job_owns_terminal();
    terminal.type_text(CTRL_Z);
    terminal.wait_until_shell_owns_terminal();
    fs::write(dir.path().join("go"), "").unwrap();
    thread::sleep(Duration::from_millis(200));
    assert!(!dir.path().join("done").exists(), "entrypoint did not stop");
    terminal.type_line("bg");

    wait_for_file(&dir.path().join("done"));
}

#[test]
fn background_job_reads_terminal_after_fg_in_bash() {
    assert_background_job_reads_terminal_after_fg(BASH);
}

#[test]
fn background_job_reads_terminal_after_fg_in_zsh() {
    assert_background_job_reads_terminal_after_fg(ZSH);
}

fn assert_background_job_reads_terminal_after_fg(shell: &[&str]) {
    let dir = TempDir::new().expect("failed to create temp dir");
    let script = "touch ready
                  read line
                  echo \"got-$line\"";

    let mut terminal = Terminal::start(shell);
    terminal.type_line(&format!("{} &", build_sandbox_command(dir.path(), script)));
    wait_for_file(&dir.path().join("ready"));
    terminal.wait_until_job_stopped();
    terminal.type_line("fg");
    terminal.wait_until_job_owns_terminal();
    terminal.type_line("hello");
    terminal.wait_for_output("got-hello");
    terminal.type_line("echo STATUS-$?");

    terminal.wait_for_output("STATUS-0");
}

#[test]
fn ctrl_c_reaches_both_sandboxes_in_a_pipeline() {
    let first = TempDir::new().expect("failed to create temp dir");
    let second = TempDir::new().expect("failed to create temp dir");
    let script = "trap 'touch cleaned; exit 0' INT
                  touch ready
                  sleep 60 & wait";

    let mut terminal = Terminal::start(BASH);
    terminal.type_line(&format!(
        "{} | {}",
        build_sandbox_command(first.path(), script),
        build_sandbox_command(second.path(), script)
    ));
    wait_for_file(&first.path().join("ready"));
    wait_for_file(&second.path().join("ready"));
    terminal.wait_until_job_owns_terminal();
    terminal.type_text(CTRL_C);
    terminal.type_line("echo STATUS-$?");
    terminal.wait_for_output("STATUS-0");

    assert!(
        first.path().join("cleaned").exists(),
        "first sandbox did not clean up"
    );
    assert!(
        second.path().join("cleaned").exists(),
        "second sandbox did not clean up"
    );
}

#[test]
fn hangup_of_session_leader_reaches_entrypoint() {
    let dir = TempDir::new().expect("failed to create temp dir");
    let args = build_sandbox_args(dir.path(), HANGUP_SCRIPT);

    let mut terminal = Terminal::start(&args.iter().map(String::as_str).collect::<Vec<_>>());
    wait_for_file(&dir.path().join("ready"));
    terminal.hang_up();

    wait_for_file(&dir.path().join("hungup"));
}

const HANGUP_SCRIPT: &str = "trap 'touch hungup; exit 0' HUP
                             touch ready
                             sleep 60 & wait";

#[test]
fn shell_hangup_reaches_entrypoint_in_bash() {
    assert_shell_hangup_reaches_entrypoint(BASH);
}

#[test]
fn shell_hangup_reaches_entrypoint_in_zsh() {
    assert_shell_hangup_reaches_entrypoint(ZSH);
}

fn assert_shell_hangup_reaches_entrypoint(shell: &[&str]) {
    let dir = TempDir::new().expect("failed to create temp dir");

    let mut terminal = Terminal::start(shell);
    terminal.type_line(&build_sandbox_command(dir.path(), HANGUP_SCRIPT));
    wait_for_file(&dir.path().join("ready"));
    terminal.hang_up();

    wait_for_file(&dir.path().join("hungup"));
}

#[test]
fn nohup_keeps_entrypoint_alive_after_hangup() {
    let dir = TempDir::new().expect("failed to create temp dir");
    let script = "touch ready
                  while [ ! -e go ]; do sleep 0.01; done
                  touch survived";

    let mut terminal = Terminal::start(BASH);
    terminal.type_line(&format!(
        "nohup {} >/dev/null 2>&1",
        build_sandbox_command(dir.path(), script)
    ));
    wait_for_file(&dir.path().join("ready"));
    terminal.hang_up();
    thread::sleep(Duration::from_millis(200));
    fs::write(dir.path().join("go"), "").unwrap();

    wait_for_file(&dir.path().join("survived"));
}