//! A home for ledger and broker tests: grants, audit log, session state, a flaky audit sink and //! a log to read. Do not edit. //! //! Included with `#[path = "support/rig.rs"] mod rig;`, beside `tmp` and `sink`. The broker //! tests add `client`. #![allow(dead_code)] // each test file uses a different part of this module use std::path::PathBuf; use brokerd::audit::Writer; use brokerd::config::{Approvals, Config, Paths, Sockets}; use brokerd::ledger::Ledger; use brokerd::state::StateStore; use proto::{AuditEvent, AuditRecord, CallId, SessionId, ToolRequest}; use crate::sink::{Flaky, Lines, Switch}; use crate::tmp::TempDir; pub struct Rig { pub dir: TempDir, pub cfg: Config, pub switch: Switch, pub lines: Lines, } impl Rig { pub fn new(tag: &str) -> Rig { Rig::with_ttl(tag, 900_000) } pub fn with_ttl(tag: &str, ttl_ms: u64) -> Rig { let dir = TempDir::new(tag); let grants = dir.path().join("grants"); std::fs::create_dir_all(&grants).unwrap(); let cfg = Config { paths: Paths { home: dir.path().to_path_buf(), grants, }, sockets: Sockets::default(), approvals: Approvals { ttl_ms }, }; Rig { dir, cfg, switch: Switch::default(), lines: Lines::default(), } } pub fn state(&self) -> StateStore { StateStore::new(&self.cfg.state_dir()) } /// Opens the audit log (once: the writer holds its lock) behind the flaky sink. pub fn ledger(&self) -> Ledger { let opened = Writer::open(&self.cfg.audit_dir(), false).unwrap(); let sink = Flaky { writer: opened.writer, switch: self.switch.clone(), }; Ledger::new(Box::new(sink), self.state(), self.lines.sink()) } /// Writes `grants/.toml`. pub fn grant(&self, id: &str, text: &str) { std::fs::write(self.cfg.paths.grants.join(format!("{id}.toml")), text).unwrap(); } pub fn remove_grant(&self, id: &str) { std::fs::remove_file(self.cfg.paths.grants.join(format!("{id}.toml"))).unwrap(); } pub fn state_file(&self, session: &str) -> PathBuf { self.cfg.state_dir().join(format!("{session}.json")) } /// Every record in the audit log, in order. pub fn records(&self) -> Vec { let dir = self.cfg.audit_dir(); let mut names: Vec = std::fs::read_dir(&dir) .unwrap() .map(|e| e.unwrap().file_name().into_string().unwrap()) .filter(|n| n.ends_with(".jsonl")) .collect(); names.sort(); let mut out = Vec::new(); for name in names { let text = std::fs::read_to_string(dir.join(name)).unwrap(); for line in text.lines() { out.push(serde_json::from_str(line).unwrap()); } } out } pub fn events(&self) -> Vec { self.records().into_iter().map(|r| r.event).collect() } } /// A grant file's text. `extra` goes before `[constraints]`, `constraints` after it. pub fn grant_text(tool: &str, mode: &str, extra: &str, constraints: &str) -> String { format!( "tool = \"{tool}\"\nmode = \"{mode}\"\nmax_taint = \"secret\"\nresult_class = \"private\"\n\ untrusted = false\n{extra}\n[constraints]\n{constraints}\n" ) } pub fn request(session: &str, call: u64, tool: &str, arguments: &str) -> ToolRequest { ToolRequest { session: SessionId::new(session).unwrap(), call: CallId(call), tool: tool.to_string(), arguments: arguments.to_string(), } }