120 lines
3.7 KiB
Rust
120 lines
3.7 KiB
Rust
//! 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 },
|
|
runner: None,
|
|
};
|
|
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/<id>.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<AuditRecord> {
|
|
let dir = self.cfg.audit_dir();
|
|
let mut names: Vec<String> = 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<AuditEvent> {
|
|
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(),
|
|
}
|
|
}
|