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

//! Properties over sequences of calls, through `broker::handle` with many threads at once:
//! the log verifies and no `seq` repeats; taint never goes down; every `Result` follows the
//! record that let its call run; the runtime sees a call only after `allowed` or an approval.
//! 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::collections::BTreeMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use brokerd::ledger::Answer;
use client::{Serve, call};
use proto::{
    AuditEvent, AuditRecord, CallId, DataClass, DecisionRecord, Message, SessionId, Timestamp,
    ToolResponse,
};
use rig::{Rig, grant_text, request};
use runtime::Recording;

fn grants(rig: &Rig) {
    let with = |class: &str, max: &str| {
        grant_text("read_file", "auto", "", "")
            .replace(
                "result_class = \"private\"",
                &format!("result_class = \"{class}\""),
            )
            .replace("max_taint = \"secret\"", &format!("max_taint = \"{max}\""))
    };
    rig.grant("notes", &(with("private", "secret") + "paths = [\"/n\"]\n"));
    rig.grant("keys", &(with("secret", "secret") + "paths = [\"/k\"]\n"));
    // Stops applying once a session has read a secret.
    rig.grant(
        "public",
        &(with("public", "private") + "paths = [\"/p\"]\n"),
    );
    rig.grant(
        "never",
        &grant_text("read_file", "deny", "", "paths = [\"/d\"]"),
    );
    rig.grant(
        "asked",
        &grant_text("read_file", "ask", "", "paths = [\"/a\"]"),
    );
}

/// Which call ran on what authority: (session, call) of every Decision that allowed and every
/// Approval whose re-decision let the call run, keyed by the decision's seq.
fn check_sequence(records: &[AuditRecord]) -> usize {
    let mut may_run: BTreeMap<u64, (SessionId, CallId)> = BTreeMap::new();
    let mut taint: BTreeMap<String, DataClass> = BTreeMap::new();
    let mut results = 0;
    for (i, r) in records.iter().enumerate() {
        assert_eq!(r.seq, i as u64, "no seq repeats or skips");
        match &r.event {
            AuditEvent::Decision {
                session,
                call,
                outcome: DecisionRecord::Allowed {},
                ..
            } => {
                may_run.insert(r.seq, (session.clone(), *call));
            }
            AuditEvent::Approval {
                session,
                call,
                decision,
                outcome: DecisionRecord::Allowed {} | DecisionRecord::Ask {},
                ..
            } => {
                may_run.insert(*decision, (session.clone(), *call));
            }
            AuditEvent::Result {
                session,
                call,
                decision,
                taint_after,
                ..
            } => {
                results += 1;
                let allowed = may_run.remove(decision);
                assert_eq!(
                    allowed,
                    Some((session.clone(), *call)),
                    "seq {}: a Result with no record letting its call run",
                    r.seq
                );
                let before = taint.insert(session.as_str().to_string(), *taint_after);
                assert!(
                    before.is_none_or(|b| b <= *taint_after),
                    "seq {}: taint went down",
                    r.seq
                );
            }
            _ => {}
        }
    }
    assert!(
        may_run.is_empty(),
        "calls allowed but never finished: {may_run:?}"
    );
    results
}

#[test]
fn eight_threads_two_sessions_fifty_calls_each() {
    let rig = Rig::new("sequence-many");
    grants(&rig);
    let rt = Recording::answering("content");
    let broker = rig.broker(&rt);
    let paths = ["/n/x", "/k/x", "/p/x", "/d/x", "/none/x"];
    let threads: Vec<_> = (0..8u64)
        .map(|t| {
            let broker = Arc::clone(&broker);
            std::thread::spawn(move || {
                let session = if t % 2 == 0 { "s-even" } else { "s-odd" };
                for i in 0..50u64 {
                    let path = paths[((t + i) % 5) as usize];
                    let args = format!(r#"{{"path":"{path}"}}"#);
                    let req = request(session, t * 1000 + i, "read_file", &args);
                    let frames = call(&broker, req);
                    assert_eq!(frames.len(), 1);
                }
            })
        })
        .collect();
    for t in threads {
        t.join().unwrap();
    }
    let report = brokerd::audit::verify_dir(&rig.cfg.audit_dir()).unwrap();
    assert!(report.failure.is_none(), "{:?}", report.failure);
    assert!(report.unfinished.is_empty());
    let records = rig.records();
    let decisions = records
        .iter()
        .filter(|r| matches!(r.event, AuditEvent::Decision { .. }))
        .count();
    assert_eq!(decisions, 400);
    let results = check_sequence(&records);
    assert_eq!(
        rt.count(),
        results,
        "the runtime saw exactly the allowed calls"
    );
    assert!(results > 0 && results < 400);
    // Both sessions read a secret.
    for s in ["s-even", "s-odd"] {
        let state = rig.state().read(&SessionId::new(s).unwrap()).unwrap();
        assert_eq!(state.taint, DataClass::Secret);
    }
}

#[test]
fn approved_and_refused_calls_run_only_after_an_approval() {
    let rig = Rig::new("sequence-ask");
    grants(&rig);
    let rt = Recording::answering("content");
    let broker = rig.broker(&rt);
    // An owner who approves even ids and refuses odd ones, as fast as they appear.
    let stop = Arc::new(AtomicBool::new(false));
    let owner = {
        let (broker, stop) = (Arc::clone(&broker), Arc::clone(&stop));
        std::thread::spawn(move || {
            while !stop.load(Ordering::SeqCst) {
                for item in broker.table().list() {
                    let Some(entry) = broker.table().take(item.approval) else {
                        continue;
                    };
                    let answer = if item.approval % 2 == 0 {
                        Answer::Approved { by: None }
                    } else {
                        Answer::Refused {
                            by: None,
                            reason: None,
                        }
                    };
                    let grants = broker.grants();
                    let done = broker.ledger().answer(
                        entry.ask,
                        item.approval,
                        answer,
                        &grants,
                        Timestamp::now(),
                    );
                    entry.reply.send(done.verdict).unwrap();
                }
                std::thread::sleep(Duration::from_millis(2));
            }
        })
    };
    let threads: Vec<_> = (0..4u64)
        .map(|t| {
            let broker = Arc::clone(&broker);
            std::thread::spawn(move || {
                for i in 0..10u64 {
                    let path = if i % 2 == 0 { "/a/x" } else { "/n/x" };
                    let args = format!(r#"{{"path":"{path}"}}"#);
                    let frames = call(&broker, request("s1", t * 100 + i, "read_file", &args));
                    let last = frames.last().unwrap();
                    assert!(
                        matches!(
                            &last.msg,
                            Message::ToolResponse(
                                ToolResponse::Result { .. } | ToolResponse::Denied { .. }
                            )
                        ),
                        "{last:?}"
                    );
                }
            })
        })
        .collect();
    for t in threads {
        t.join().unwrap();
    }
    stop.store(true, Ordering::SeqCst);
    owner.join().unwrap();
    let records = rig.records();
    let approvals = records
        .iter()
        .filter(|r| matches!(r.event, AuditEvent::Approval { .. }))
        .count();
    assert_eq!(approvals, 20, "one Approval for each of the twenty asks");
    let results = check_sequence(&records);
    assert_eq!(rt.count(), results);
    let report = brokerd::audit::verify_dir(&rig.cfg.audit_dir()).unwrap();
    assert!(report.failure.is_none() && report.abandoned.is_empty());
}