boxmaker/docs/plans/M3a/files/crates/brokerd/tests/broker.rs

//! One `broker.sock` connection: decision, record, run, answer; and every way it fails closed.
//! Do not edit.

#[path = "support/client.rs"]
mod client;
#[path = "support/rig.rs"]
mod rig;
#[path = "support/runtime.rs"]
mod runtime;
#[path = "support/sink.rs"]
mod sink;
#[path = "support/tmp.rs"]
mod tmp;

use std::io::Write;
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixStream;
use std::sync::Arc;

use brokerd::broker;
use brokerd::ledger::NOT_RECORDED;
use brokerd::runner::{REFUSING, RunError};
use client::{Serve, call, next, open};
use proto::{
    ApprovalList, Approve, ApproveResult, AuditEvent, DataClass, DecisionRecord, DenyReason, Empty,
    Envelope, ErrorCode, GrantsReport, Message, Refuse, ResultStatus, SessionId, ToolResponse,
    Turn, TurnDone, TurnEvent, Usage, WireError,
};
use rig::{Rig, grant_text, request};
use runtime::Recording;

const NOTES: &str = r#"{"path":"/n/a"}"#;

fn read_notes(call_id: u64) -> proto::ToolRequest {
    request("s1", call_id, "read_file", NOTES)
}

fn only(frames: &[Envelope]) -> &ToolResponse {
    assert_eq!(frames.len(), 1, "{frames:?}");
    assert!(frames[0].r#final);
    match &frames[0].msg {
        Message::ToolResponse(r) => r,
        other => panic!("{other:?}"),
    }
}

fn denied(reason: DenyReason) -> ToolResponse {
    ToolResponse::Denied { reason }
}

fn auto_notes(rig: &Rig, class: &str) {
    let text = grant_text("read_file", "auto", "", "paths = [\"/n\"]").replace(
        "result_class = \"private\"",
        &format!("result_class = \"{class}\""),
    );
    rig.grant("notes", &text);
}

#[test]
fn a_call_no_grant_allows_is_denied_and_recorded() {
    let rig = Rig::new("broker-nogrant");
    let rt = Recording::answering("x");
    let broker = rig.broker(&rt);
    let frames = call(&broker, read_notes(41));
    assert_eq!(frames[0].id, 41, "the answer carries the request's id");
    assert_eq!(only(&frames), &denied(DenyReason::NoGrant));
    assert_eq!(rt.count(), 0);
    match rig.events().as_slice() {
        [AuditEvent::Decision { outcome, .. }] => assert_eq!(
            *outcome,
            DecisionRecord::Denied {
                reason: DenyReason::NoGrant
            }
        ),
        other => panic!("{other:?}"),
    }
}

#[test]
fn an_allowed_call_runs_and_its_result_raises_the_state() {
    let rig = Rig::new("broker-allowed");
    auto_notes(&rig, "secret");
    let rt = Recording::answering("the notes");
    let broker = rig.broker(&rt);
    let frames = call(&broker, read_notes(7));
    assert_eq!(
        only(&frames),
        &ToolResponse::Result {
            content: "the notes".to_string(),
            class: DataClass::Secret,
            untrusted: false,
            truncated: false,
        }
    );
    assert_eq!(rt.count(), 1);
    let events = rig.events();
    assert!(matches!(
        events.as_slice(),
        [
            AuditEvent::Decision { .. },
            AuditEvent::Result {
                decision: 0,
                status: ResultStatus::Result,
                taint_after: DataClass::Secret,
                ..
            }
        ]
    ));
    let state = rig.state().read(&SessionId::new("s1").unwrap()).unwrap();
    assert_eq!(state.taint, DataClass::Secret);
}

#[test]
fn a_runtime_failure_is_passed_on_and_recorded_as_failed() {
    let rig = Rig::new("broker-refusing");
    auto_notes(&rig, "secret");
    let rt = Recording::with(Err(RunError::Unavailable(REFUSING.to_string())));
    let broker = rig.broker(&rt);
    let frames = call(&broker, read_notes(1));
    assert_eq!(
        only(&frames),
        &ToolResponse::Failed {
            message: REFUSING.to_string()
        }
    );
    match &rig.events()[1] {
        AuditEvent::Result {
            status,
            taint_after,
            ..
        } => assert_eq!(
            (*status, *taint_after),
            (ResultStatus::Failed, DataClass::Private)
        ),
        other => panic!("{other:?}"),
    }
    assert!(!rig.state_file("s1").exists(), "a failure changes no state");
}

fn every_kind_but_tool_request() -> Vec<Message> {
    let usage = Usage {
        cache_n: 0,
        prompt_n: 0,
        predicted_n: 0,
        reasoning_tokens: 0,
        thinking_capped: false,
    };
    vec![
        Message::ToolResponse(denied(DenyReason::NoGrant)),
        Message::Error(WireError {
            code: ErrorCode::Internal,
            detail: String::new(),
        }),
        Message::Turn(Turn {
            session: SessionId::new("s1").unwrap(),
            content: String::new(),
            resume: false,
        }),
        Message::TurnEvent(TurnEvent::Content {
            text: String::new(),
        }),
        Message::TurnDone(TurnDone {
            content: String::new(),
            usage,
        }),
        Message::Approvals(Empty {}),
        Message::ApprovalList(ApprovalList { items: Vec::new() }),
        Message::Approve(Approve { approval: 0 }),
        Message::ApproveResult(ApproveResult {
            outcome: DecisionRecord::Allowed {},
        }),
        Message::Refuse(Refuse {
            approval: 0,
            reason: None,
        }),
        Message::Ok(Empty {}),
        Message::CheckGrants(Empty {}),
        Message::GrantsReport(GrantsReport {
            problems: Vec::new(),
        }),
    ]
}

#[test]
fn every_other_kind_on_broker_sock_is_forbidden() {
    let rig = Rig::new("broker-forbidden");
    let rt = Recording::answering("x");
    let broker = rig.broker(&rt);
    let kinds = every_kind_but_tool_request();
    assert_eq!(kinds.len(), 13, "every Message variant but tool_request");
    for (i, msg) in kinds.into_iter().enumerate() {
        let id = 100 + i as u64;
        let mut stream = open(&broker, broker::handle, id, msg);
        let env = next(&mut stream);
        assert_eq!((env.id, env.r#final), (id, true));
        match env.msg {
            Message::Error(e) => assert_eq!(e.code, ErrorCode::Forbidden),
            other => panic!("{other:?}"),
        }
    }
    let lines = rig.lines.with("broker.sock");
    assert_eq!(lines.len(), 13, "{:?}", rig.lines.all());
    for line in &lines {
        assert!(
            line.ends_with("\nsee docs/runbook.md#socket-forbidden"),
            "{line}"
        );
    }
    assert_eq!(rig.lines.with("kind approve on broker.sock").len(), 1);
    assert!(
        rig.events().is_empty(),
        "nothing is recorded for a refused kind"
    );
}

#[test]
fn a_frame_that_is_not_json_is_answered_with_bad_message() {
    let rig = Rig::new("broker-badframe");
    let rt = Recording::answering("x");
    let broker = rig.broker(&rt);
    let (mut client, server) = UnixStream::pair().unwrap();
    let b = Arc::clone(&broker);
    std::thread::spawn(move || broker::handle(server, &b));
    client.write_all(&3u32.to_be_bytes()).unwrap();
    client.write_all(b"{{{").unwrap();
    let env = next(&mut client);
    assert_eq!((env.id, env.r#final), (0, true));
    match env.msg {
        Message::Error(e) => assert_eq!(e.code, ErrorCode::BadMessage),
        other => panic!("{other:?}"),
    }
    assert!(rig.events().is_empty());
}

#[test]
fn one_invalid_grant_file_denies_a_call_a_valid_file_would_allow() {
    let rig = Rig::new("broker-invalid");
    auto_notes(&rig, "private");
    rig.grant("bad", "tool = \"read_file\"\nmdoe = \"auto\"\n");
    let rt = Recording::answering("x");
    let broker = rig.broker(&rt);
    for id in [1, 2] {
        let frames = call(&broker, read_notes(id));
        assert_eq!(only(&frames), &denied(DenyReason::GrantsInvalid));
    }
    assert_eq!(rt.count(), 0);
    // Printed once for the two calls, with the pointer.
    let printed = rig.lines.with("bad.toml");
    assert_eq!(printed.len(), 1, "{:?}", rig.lines.all());
    assert!(
        printed[0].ends_with("see docs/runbook.md#grants-invalid"),
        "{}",
        printed[0]
    );
    rig.remove_grant("bad");
    let frames = call(&broker, read_notes(3));
    assert!(matches!(only(&frames), ToolResponse::Result { .. }));
    assert_eq!(rt.count(), 1);
}

#[test]
fn a_damaged_session_state_is_state_unreadable() {
    let rig = Rig::new("broker-state");
    auto_notes(&rig, "private");
    std::fs::create_dir_all(rig.cfg.state_dir()).unwrap();
    std::fs::write(rig.state_file("s1"), "{\"taint\":\"loud\"}").unwrap();
    let rt = Recording::answering("x");
    let broker = rig.broker(&rt);
    let frames = call(&broker, read_notes(1));
    assert_eq!(only(&frames), &denied(DenyReason::StateUnreadable));
    assert_eq!(rt.count(), 0);
    let lines = rig.lines.with("s1.json");
    assert!(
        lines
            .iter()
            .any(|l| l.ends_with("see docs/runbook.md#broker-state-damaged")),
        "{lines:?}"
    );
}

#[test]
fn a_failed_audit_write_runs_nothing_now_or_later() {
    let rig = Rig::new("broker-audit");
    auto_notes(&rig, "private");
    let rt = Recording::answering("x");
    let broker = rig.broker(&rt);
    rig.switch.fail(true);
    let frames = call(&broker, read_notes(1));
    assert_eq!(only(&frames), &denied(DenyReason::AuditUnavailable));
    rig.switch.fail(false);
    let frames = call(&broker, read_notes(2));
    assert_eq!(only(&frames), &denied(DenyReason::AuditUnavailable));
    assert_eq!(rt.count(), 0);
    assert!(
        !rig.lines
            .with("see docs/runbook.md#audit-unavailable")
            .is_empty()
    );
}

#[test]
fn a_state_that_cannot_be_written_withholds_the_content() {
    if tmp::running_as_root("a_state_that_cannot_be_written_withholds_the_content") {
        return;
    }
    let rig = Rig::new("broker-ro");
    auto_notes(&rig, "secret");
    let dir = rig.cfg.state_dir();
    std::fs::create_dir_all(&dir).unwrap();
    std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o500)).unwrap();
    let rt = Recording::answering("the secret");
    let broker = rig.broker(&rt);
    let frames = call(&broker, read_notes(1));
    std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o700)).unwrap();
    assert_eq!(
        only(&frames),
        &ToolResponse::Failed {
            message: NOT_RECORDED.to_string()
        }
    );
    assert_eq!(rt.count(), 1, "it ran; its content is what is withheld");
    let text = format!("{frames:?}");
    assert!(!text.contains("the secret"), "{text}");
}