agent-sandbox/tests/e2e/common.rs

use std::fs::{self, File};
use std::io::{self, Read, Write};
use std::ops::{Deref, DerefMut};
use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::thread;
use std::time::{Duration, Instant};

use tempfile::TempDir;

// Blacklist mode globs the host filesystem at sandbox startup to find
// sensitive paths to mask. A matching host file that vanishes mid-startup
// makes bwrap fail in any concurrent blacklist sandbox.
pub struct HostGlobsLock;

impl HostGlobsLock {
    /// Acquire shared access. Hold while running a blacklist sandbox.
    pub fn for_scan() -> RwLockReadGuard<'static, ()> {
        LOCK.read().unwrap()
    }

    /// Acquire exclusive access. Hold while creating or deleting a host
    /// path that may match a blacklist glob (e.g. `/tmp/ssh-*`).
    pub fn for_mutation() -> RwLockWriteGuard<'static, ()> {
        LOCK.write().unwrap()
    }
}

static LOCK: RwLock<()> = RwLock::new(());

#[allow(dead_code)]
enum HostGlobsGuard {
    Scan(RwLockReadGuard<'static, ()>),
    Mutation(RwLockWriteGuard<'static, ()>),
}

pub struct Sandbox {
    cmd: Command,
    _guard: Option<HostGlobsGuard>,
}

impl Sandbox {
    pub fn new(extra_args: &[&str]) -> Self {
        Self::build(&["--no-config"], extra_args, scan_guard_for(extra_args))
    }

    pub fn new_with_config(extra_args: &[&str]) -> Self {
        Self::build(&[], extra_args, scan_guard_for(extra_args))
    }

    pub fn new_for_host_mutation(extra_args: &[&str]) -> Self {
        debug_assert!(
            extra_args.contains(&"--blacklist"),
            "new_for_host_mutation is only meaningful for blacklist sandboxes"
        );
        Self::build(
            &["--no-config"],
            extra_args,
            Some(HostGlobsGuard::Mutation(HostGlobsLock::for_mutation())),
        )
    }

    fn build(prefix: &[&str], extra_args: &[&str], guard: Option<HostGlobsGuard>) -> Self {
        let mut cmd = Command::new(env!("CARGO_BIN_EXE_agent-sandbox"));
        cmd.args(prefix);
        cmd.args(extra_args);
        Self { cmd, _guard: guard }
    }
}

fn scan_guard_for(extra_args: &[&str]) -> Option<HostGlobsGuard> {
    extra_args
        .contains(&"--blacklist")
        .then(|| HostGlobsGuard::Scan(HostGlobsLock::for_scan()))
}

impl Deref for Sandbox {
    type Target = Command;
    fn deref(&self) -> &Command {
        &self.cmd
    }
}

impl DerefMut for Sandbox {
    fn deref_mut(&mut self) -> &mut Command {
        &mut self.cmd
    }
}

pub struct ConfigFile {
    _dir: TempDir,
    path: String,
}

impl ConfigFile {
    pub fn new(content: &str) -> Self {
        let dir = TempDir::new().unwrap();
        let path = dir.path().join("config.toml");
        fs::write(&path, content).expect("failed to write config");
        Self {
            _dir: dir,
            path: path.to_str().unwrap().to_string(),
        }
    }
}

impl std::ops::Deref for ConfigFile {
    type Target = str;
    fn deref(&self) -> &str {
        &self.path
    }
}

pub fn wait_for_file(path: &Path) {
    let deadline = Instant::now() + Duration::from_secs(10);
    while !path.exists() {
        assert!(Instant::now() < deadline, "timed out waiting for {path:?}");
        thread::sleep(Duration::from_millis(10));
    }
}

pub struct Terminal {
    master: Option<File>,
    output: String,
    session_leader: Child,
    _guard: RwLockReadGuard<'static, ()>,
}

impl Terminal {
    pub fn start(command_line: &[&str]) -> Self {
        let guard = HostGlobsLock::for_scan();
        let (master, slave) = open_pty();

        let mut command = Command::new(command_line[0]);
        command
            .args(&command_line[1..])
            .env("TERM", "dumb")
            .stdin(Stdio::from(slave.try_clone().unwrap()))
            .stdout(Stdio::from(slave.try_clone().unwrap()))
            .stderr(Stdio::from(slave));
        unsafe { command.pre_exec(take_terminal) };
        let session_leader = command.spawn().expect("failed to start session leader");

        Self {
            master: Some(master),
            output: String::new(),
            session_leader,
            _guard: guard,
        }
    }

    pub fn type_text(&mut self, text: &str) {
        self.master().write_all(text.as_bytes()).unwrap();
    }

    pub fn type_line(&mut self, line: &str) {
        self.type_text(&format!("{line}\n"));
    }

    pub fn read_output(&mut self) -> &str {
        self.read_available_output();
        &self.output
    }

    pub fn wait_for_output(&mut self, text: &str) {
        self.wait_until(&format!("output {text:?}"), |terminal| {
            terminal.output.contains(text)
        });
    }

    pub fn wait_until_shell_owns_terminal(&mut self) {
        self.wait_until("shell in foreground", Self::is_shell_in_foreground);
    }

    pub fn wait_until_job_owns_terminal(&mut self) {
        self.wait_until("job in foreground", |terminal| {
            !terminal.is_shell_in_foreground()
        });
    }

    pub fn wait_until_job_stopped(&mut self) {
        self.wait_until("job stopped", Self::has_stopped_job);
    }

    pub fn hang_up(&mut self) {
        self.master = None;
    }

    fn master(&mut self) -> &mut File {
        self.master.as_mut().expect("terminal is hung up")
    }

    fn read_available_output(&mut self) {
        let Some(master) = &mut self.master else {
            return;
        };
        let mut buffer = [0; 4096];
        while let Ok(count @ 1..) = master.read(&mut buffer) {
            self.output
                .push_str(&String::from_utf8_lossy(&buffer[..count]));
        }
    }

    fn wait_until(&mut self, description: &str, condition: impl Fn(&Self) -> bool) {
        let deadline = Instant::now() + Duration::from_secs(10);
        loop {
            self.read_available_output();
            if condition(self) {
                return;
            }
            assert!(
                Instant::now() < deadline,
                "timed out waiting for {description}, terminal output:\n{}",
                self.output
            );
            thread::sleep(Duration::from_millis(10));
        }
    }

    fn is_shell_in_foreground(&self) -> bool {
        let master = self.master.as_ref().expect("terminal is hung up");
        let foreground = unsafe { libc::tcgetpgrp(master.as_raw_fd()) };
        foreground == self.session_leader.id() as i32
    }

    fn has_stopped_job(&self) -> bool {
        let session = self.session_leader.id().to_string();
        fs::read_dir("/proc")
            .unwrap()
            .filter_map(|entry| fs::read_to_string(entry.ok()?.path().join("stat")).ok())
            .any(|stat| {
                let fields_after_comm = &stat[stat.rfind(')').unwrap() + 1..];
                let fields: Vec<&str> = fields_after_comm.split_whitespace().collect();
                fields[0] == "T" && fields[3] == session
            })
    }
}

fn open_pty() -> (File, OwnedFd) {
    let flags = libc::O_RDWR | libc::O_NOCTTY | libc::O_CLOEXEC;
    let master = unsafe { libc::posix_openpt(flags | libc::O_NONBLOCK) };
    assert!(
        master >= 0,
        "posix_openpt failed: {}",
        io::Error::last_os_error()
    );
    let master = unsafe { File::from_raw_fd(master) };

    assert_eq!(
        unsafe { libc::unlockpt(master.as_raw_fd()) },
        0,
        "unlockpt failed"
    );
    let slave = unsafe { libc::ioctl(master.as_raw_fd(), libc::TIOCGPTPEER, flags) };
    assert!(
        slave >= 0,
        "TIOCGPTPEER failed: {}",
        io::Error::last_os_error()
    );
    (master, unsafe { OwnedFd::from_raw_fd(slave) })
}

fn take_terminal() -> io::Result<()> {
    if unsafe { libc::setsid() } == -1 || unsafe { libc::ioctl(0, libc::TIOCSCTTY, 0) } == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

impl Drop for Terminal {
    fn drop(&mut self) {
        let _ = self.session_leader.kill();
        let _ = self.session_leader.wait();
    }
}