boxmaker/crates/brokerd/tests/memory_port.rs

//! Decided memory calls go to `memoryd` with the session's state, everything else to the other
//! runtime; `memoryd`'s label can only make a result stricter; and a container runtime never runs
//! a memory tool (M5a spec, section 8). Do not edit.

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

use std::os::unix::net::UnixListener;
use std::path::Path;
use std::sync::mpsc;
use std::time::Duration;

use brokerd::container::{IN_MEMORYD, Podman};
use brokerd::memory_port::{MemoryPort, Routed, UNREACHABLE};
use brokerd::policy::{Outcome, SessionState, decide};
use brokerd::runner::run;
use build::{grant, now, read, request, set};
use bxcore::{
    DataClass, Envelope, MemoryCall, Message, Mode, ToolRequest, ToolResponse, read_frame,
    write_frame,
};
use fake_podman::{Fake, Lines, serial};
use runtime::{Recording, Shared};
use tmp::TempDir;

/// A fake `memoryd`: answers one call with `answer` and hands back what it was sent.
fn fake_memoryd(socket: &Path, answer: ToolResponse) -> mpsc::Receiver<MemoryCall> {
    let listener = UnixListener::bind(socket).unwrap();
    let (tx, rx) = mpsc::channel();
    std::thread::spawn(move || {
        let (mut stream, _) = listener.accept().unwrap();
        let env = read_frame(&mut stream).unwrap();
        let Message::MemoryCall(call) = env.msg else {
            panic!("not a memory call: {:?}", env.msg)
        };
        tx.send(call).unwrap();
        let reply = Envelope {
            v: 1,
            id: env.id,
            r#final: true,
            msg: Message::ToolResponse(answer),
        };
        write_frame(&mut stream, &reply).unwrap();
    });
    rx
}

fn result(class: DataClass, untrusted: bool) -> ToolResponse {
    ToolResponse::Result {
        content: "[1] memory · …".to_string(),
        class,
        untrusted,
        truncated: false,
    }
}

fn decide_in(req: ToolRequest, state: SessionState) -> brokerd::policy::Decision {
    let grants = set(vec![grant("notes", "read_file", Mode::Auto).paths(&["/n"])]);
    match decide(req, &grants, state, now()) {
        Outcome::Allowed(d) => d,
        other => panic!("not allowed: {other:?}"),
    }
}

fn routed(socket: &Path, other: &std::sync::Arc<Recording>) -> Routed {
    Routed {
        memory: MemoryPort::new(socket.to_path_buf(), Duration::from_secs(5)),
        other: Box::new(Shared(std::sync::Arc::clone(other))),
    }
}

#[test]
fn a_memory_call_goes_to_memoryd_with_the_session_state() {
    let dir = TempDir::new("mp-route");
    let socket = dir.path().join("memory.sock");
    let seen = fake_memoryd(&socket, result(DataClass::Private, false));
    let other = Recording::answering("container");
    let state = SessionState {
        taint: DataClass::Secret,
        untrusted: true,
    };
    let got = run(
        decide_in(request("memory", r#"{"text":"x","action":"note"}"#), state),
        &routed(&socket, &other),
    );
    assert!(matches!(got, ToolResponse::Result { .. }), "{got:?}");
    let call = seen.recv_timeout(Duration::from_secs(5)).unwrap();
    assert_eq!((call.taint, call.untrusted), (DataClass::Secret, true));
    assert_eq!(call.request.tool, "memory");
    assert_eq!(
        call.request.arguments, r#"{"action":"note","text":"x"}"#,
        "canonical"
    );
    assert_eq!(other.count(), 0, "nothing reached the other runtime");
}

#[test]
fn other_tools_go_to_the_other_runtime() {
    let dir = TempDir::new("mp-other");
    let other = Recording::answering("the file");
    let got = run(
        decide_in(read("/n/a"), SessionState::default()),
        &routed(&dir.path().join("none.sock"), &other),
    );
    assert!(matches!(got, ToolResponse::Result { ref content, .. } if content == "the file"));
    assert_eq!(other.count(), 1);
}

#[test]
fn memoryd_can_make_a_result_stricter_but_never_laxer() {
    let dir = TempDir::new("mp-label");
    let other = Recording::answering("unused");
    for (n, (answer, want)) in [
        ((DataClass::Secret, true), (DataClass::Secret, true)),
        ((DataClass::Public, false), (DataClass::Private, false)),
        ((DataClass::Private, true), (DataClass::Private, true)),
    ]
    .into_iter()
    .enumerate()
    {
        let socket = dir.path().join(format!("m-{n}.sock"));
        let _seen = fake_memoryd(&socket, result(answer.0, answer.1));
        let got = run(
            decide_in(
                request("recall", r#"{"query":"x"}"#),
                SessionState::default(),
            ),
            &routed(&socket, &other),
        );
        let ToolResponse::Result {
            class, untrusted, ..
        } = got
        else {
            panic!("{got:?}")
        };
        assert_eq!((class, untrusted), want, "memoryd said {answer:?}");
    }
}

#[test]
fn failures_and_silence_are_failed_results() {
    let dir = TempDir::new("mp-fail");
    let socket = dir.path().join("memory.sock");
    let _seen = fake_memoryd(
        &socket,
        ToolResponse::Failed {
            message: "memory.md already has this entry".to_string(),
        },
    );
    let other = Recording::answering("unused");
    let add = r#"{"action":"add","file":"memory","text":"x"}"#;
    let got = run(
        decide_in(request("memory", add), SessionState::default()),
        &routed(&socket, &other),
    );
    assert_eq!(
        got,
        ToolResponse::Failed {
            message: "memory.md already has this entry".to_string()
        }
    );
    let got = run(
        decide_in(request("memory", add), SessionState::default()),
        &routed(&dir.path().join("nobody.sock"), &other),
    );
    assert_eq!(
        got,
        ToolResponse::Failed {
            message: UNREACHABLE.to_string()
        }
    );
}

#[test]
fn a_container_runtime_never_runs_a_memory_tool() {
    let _s = serial();
    let fake = Fake::new("memory", "exit 0");
    let log = Lines::default();
    let podman = Podman::new(fake.runner(""), fake.dir.join("egress"), log.sink());
    let got = run(
        decide_in(
            request("recall", r#"{"query":"x"}"#),
            SessionState::default(),
        ),
        &podman,
    );
    assert_eq!(
        got,
        ToolResponse::Failed {
            message: IN_MEMORYD.to_string()
        }
    );
    assert!(fake.calls().is_empty(), "podman was never called");
}