boxmaker/docs/plans/M3b/files/crates/brokerd/tests/serve_runner.rs

//! `brokerd serve` with a `[runner]` section runs allowed calls through Podman (here a fake), and
//! says which runtime it uses. Without the section it refuses every call, as in M3a. Do not edit.

#[path = "support/fake_podman.rs"]
mod fake_podman;

use std::io::Read;
use std::os::unix::net::UnixStream;
use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};

use fake_podman::{Fake, IMAGE, serial};
use proto::{
    CallId, DataClass, Envelope, Message, PROTOCOL_VERSION, SessionId, ToolRequest, ToolResponse,
    read_frame, write_frame,
};

struct Running(Child);

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

impl Running {
    fn stop(mut self) -> String {
        let _ = self.0.kill();
        let _ = self.0.wait();
        let mut err = String::new();
        if let Some(mut stderr) = self.0.stderr.take() {
            let _ = stderr.read_to_string(&mut err);
        }
        err
    }
}

fn start(home: &Path, runner: &str) -> Running {
    std::fs::create_dir_all(home.join("grants")).unwrap();
    std::fs::write(
        home.join("grants/notes.toml"),
        "tool = \"read_file\"\nmode = \"auto\"\nmax_taint = \"secret\"\n[constraints]\npaths = [\"/n\"]\n",
    )
    .unwrap();
    let config = home.join("brokerd.toml");
    std::fs::write(
        &config,
        format!(
            "[paths]\nhome = \"{0}\"\ngrants = \"{0}/grants\"\n{runner}\n",
            home.display()
        ),
    )
    .unwrap();
    let child = Command::new(env!("CARGO_BIN_EXE_brokerd"))
        .args(["serve", "--config"])
        .arg(&config)
        .stdout(Stdio::null())
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    let running = Running(child);
    let until = Instant::now() + Duration::from_secs(10);
    while UnixStream::connect(home.join("run/loop-broker/broker.sock")).is_err() {
        assert!(Instant::now() < until, "brokerd never listened");
        std::thread::sleep(Duration::from_millis(20));
    }
    running
}

fn read_call(home: &Path) -> ToolResponse {
    let mut stream = UnixStream::connect(home.join("run/loop-broker/broker.sock")).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    let request = ToolRequest {
        session: SessionId::new("s1").unwrap(),
        call: CallId(1),
        tool: "read_file".to_string(),
        arguments: r#"{"path":"/n/a.md"}"#.to_string(),
    };
    let envelope = Envelope {
        v: PROTOCOL_VERSION,
        id: 1,
        r#final: true,
        msg: Message::ToolRequest(request),
    };
    write_frame(&mut stream, &envelope).unwrap();
    match read_frame(&mut stream).unwrap().msg {
        Message::ToolResponse(response) => response,
        other => panic!("{other:?}"),
    }
}

#[test]
fn with_a_runner_an_allowed_call_runs_in_a_container() {
    let _s = serial();
    let fake = Fake::new(
        "serve",
        "cat > /dev/null; printf 'from the container'; exit 0",
    );
    let home = fake.dir.join("home");
    let running = start(
        &home,
        &format!(
            "[runner]\npodman = \"{}\"\nimage = \"{IMAGE}\"\n",
            fake.script.display()
        ),
    );
    assert_eq!(
        read_call(&home),
        ToolResponse::Result {
            content: "from the container".to_string(),
            class: DataClass::Private,
            untrusted: true,
            truncated: false,
        }
    );
    let printed = running.stop();
    assert!(
        printed.contains(&format!("tools run in containers from {IMAGE}")),
        "{printed}"
    );
    assert_eq!(fake.calls().len(), 1);
}

#[test]
fn without_a_runner_every_call_is_refused_and_it_says_so() {
    let _s = serial();
    let fake = Fake::new("serve-none", "exit 0");
    let home = fake.dir.join("home");
    let running = start(&home, "");
    assert_eq!(
        read_call(&home),
        ToolResponse::Failed {
            message: brokerd::runner::REFUSING.to_string()
        }
    );
    let printed = running.stop();
    assert!(printed.contains("no [runner] section"), "{printed}");
    assert!(fake.calls().is_empty());
}