use indexmap::IndexMap;
use lex_bytecode::program::{DeclaredEffect, EffectArg, Program};
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone, Default)]
pub struct Policy {
pub allow_effects: BTreeSet<String>,
pub allow_fs_read: Vec<PathBuf>,
pub allow_fs_write: Vec<PathBuf>,
pub allow_net_host: Vec<String>,
pub allow_proc: Vec<String>,
pub budget: Option<u64>,
}
impl Policy {
pub fn pure() -> Self { Self::default() }
pub fn permissive() -> Self {
let mut s = BTreeSet::new();
for k in [
"io", "net", "time", "rand", "llm", "proc", "panic",
"fs_read", "fs_write", "budget",
"llm_local", "llm_cloud", "a2a", "mcp",
"env",
"sql", "random", "chat", "log", "kv", "stream", "fs_walk", "concurrent", ] {
s.insert(k.to_string());
}
Self {
allow_effects: s,
allow_fs_read: Vec::new(),
allow_fs_write: Vec::new(),
allow_net_host: Vec::new(),
allow_proc: Vec::new(),
budget: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, thiserror::Error)]
#[error("policy violation: {kind} {detail}")]
pub struct PolicyViolation {
pub kind: String,
pub detail: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub effect: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub at: Option<String>,
}
impl PolicyViolation {
pub fn effect_not_allowed(effect: &str, at: impl Into<String>) -> Self {
Self {
kind: "effect_not_allowed".into(),
detail: format!("effect `{effect}` not in --allow-effects"),
effect: Some(effect.into()),
path: None,
at: Some(at.into()),
}
}
pub fn fs_path_not_allowed(effect: &str, path: &str, at: impl Into<String>) -> Self {
Self {
kind: "fs_path_not_allowed".into(),
detail: format!("path `{path}` outside --allow-{effect}"),
effect: Some(effect.into()),
path: Some(path.into()),
at: Some(at.into()),
}
}
pub fn budget_exceeded(declared: u64, ceiling: u64) -> Self {
Self {
kind: "budget_exceeded".into(),
detail: format!("declared budget {declared} exceeds ceiling {ceiling}"),
effect: Some("budget".into()),
path: None,
at: None,
}
}
}
pub fn check_program(program: &Program, policy: &Policy) -> Result<PolicyReport, Vec<PolicyViolation>> {
let mut violations = Vec::new();
let mut total_budget: u64 = 0;
let mut declared_effects: IndexMap<String, Vec<DeclaredEffect>> = IndexMap::new();
for f in &program.functions {
for e in &f.effects {
declared_effects.entry(f.name.clone()).or_default().push(e.clone());
if !is_effect_allowed(&policy.allow_effects, e) {
violations.push(PolicyViolation::effect_not_allowed(
&declared_effect_pretty(e), &f.name));
continue;
}
if e.kind == "fs_read" || e.kind == "fs_write" {
if let Some(EffectArg::Str(path)) = &e.arg {
let allowlist = if e.kind == "fs_read" {
&policy.allow_fs_read
} else {
&policy.allow_fs_write
};
if !path_under_any(path, allowlist) {
violations.push(PolicyViolation::fs_path_not_allowed(&e.kind, path, &f.name));
}
}
}
if e.kind == "budget" {
if let Some(EffectArg::Int(n)) = &e.arg {
if *n >= 0 { total_budget = total_budget.saturating_add(*n as u64); }
}
}
}
}
if let Some(ceiling) = policy.budget {
if total_budget > ceiling {
violations.push(PolicyViolation::budget_exceeded(total_budget, ceiling));
}
}
if violations.is_empty() {
Ok(PolicyReport { declared_effects, total_budget })
} else {
Err(violations)
}
}
#[derive(Debug, Clone)]
pub struct PolicyReport {
pub declared_effects: IndexMap<String, Vec<DeclaredEffect>>,
pub total_budget: u64,
}
fn path_under_any(p: &str, list: &[PathBuf]) -> bool {
let candidate = Path::new(p);
list.iter().any(|allowed| candidate.starts_with(allowed))
}
fn declared_effect_pretty(e: &DeclaredEffect) -> String {
match &e.arg {
None => e.kind.clone(),
Some(EffectArg::Str(s)) => format!("{}(\"{}\")", e.kind, s),
Some(EffectArg::Int(n)) => format!("{}({})", e.kind, n),
Some(EffectArg::Ident(s)) => format!("{}({})", e.kind, s),
}
}
pub fn is_effect_allowed(grants: &BTreeSet<String>, e: &DeclaredEffect) -> bool {
grants.iter().any(|g| grant_subsumes(g, e))
}
fn grant_subsumes(grant: &str, e: &DeclaredEffect) -> bool {
let (g_name, g_arg) = parse_grant(grant);
if g_name != e.kind { return false; }
match (g_arg, &e.arg) {
(None, _) => true, (Some(_), None) => false, (Some(g), Some(EffectArg::Str(d))) => g == d,
(Some(_), Some(_)) => false,
}
}
fn parse_grant(s: &str) -> (&str, Option<&str>) {
if let Some((name, rest)) = s.split_once('(') {
if let Some(arg) = rest.strip_suffix(')') {
return (name, Some(arg.trim_matches('"')));
}
}
if let Some((name, arg)) = s.split_once(':') {
return (name, Some(arg));
}
(s, None)
}