use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Scope {
All,
Only(BTreeSet<String>),
}
impl Scope {
pub fn only<I, S>(names: I) -> Scope
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Scope::Only(names.into_iter().map(Into::into).collect())
}
pub fn allows(&self, name: &str) -> bool {
match self {
Scope::All => true,
Scope::Only(set) => set.contains(name),
}
}
pub fn narrow(&self, requested: &Scope) -> Scope {
match (self, requested) {
(Scope::All, r) => r.clone(),
(p @ Scope::Only(_), Scope::All) => p.clone(),
(Scope::Only(p), Scope::Only(r)) => Scope::Only(p.intersection(r).cloned().collect()),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolScope {
pub servers: Scope,
pub tools: Scope,
}
impl ToolScope {
pub fn all() -> ToolScope {
ToolScope {
servers: Scope::All,
tools: Scope::All,
}
}
pub fn allows_server(&self, server: &str) -> bool {
self.servers.allows(server)
}
pub fn allows(&self, server: &str, tool: &str) -> bool {
self.servers.allows(server) && self.tools.allows(tool)
}
pub fn narrow(&self, requested: &ToolScope) -> ToolScope {
ToolScope {
servers: self.servers.narrow(&requested.servers),
tools: self.tools.narrow(&requested.tools),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Trifecta {
pub untrusted_input: bool,
pub sensitive: bool,
pub egress: bool,
}
impl Trifecta {
pub fn legs(self) -> u8 {
self.untrusted_input as u8 + self.sensitive as u8 + self.egress as u8
}
pub fn merge(self, other: Trifecta) -> Trifecta {
Trifecta {
untrusted_input: self.untrusted_input || other.untrusted_input,
sensitive: self.sensitive || other.sensitive,
egress: self.egress || other.egress,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RuleOfTwo {
Ok,
Warn,
Refuse,
}
pub fn evaluate(tags: Trifecta, allow_trifecta: bool) -> RuleOfTwo {
if tags.legs() < 3 {
RuleOfTwo::Ok
} else if allow_trifecta {
RuleOfTwo::Warn
} else {
RuleOfTwo::Refuse
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TrifectaTag {
UntrustedInput,
Sensitive,
Egress,
}
impl TrifectaTag {
pub fn parse(s: &str) -> Option<TrifectaTag> {
match s {
"untrusted_input" => Some(TrifectaTag::UntrustedInput),
"sensitive" => Some(TrifectaTag::Sensitive),
"egress" => Some(TrifectaTag::Egress),
_ => None,
}
}
pub fn as_trifecta(self) -> Trifecta {
match self {
TrifectaTag::UntrustedInput => Trifecta {
untrusted_input: true,
..Default::default()
},
TrifectaTag::Sensitive => Trifecta {
sensitive: true,
..Default::default()
},
TrifectaTag::Egress => Trifecta {
egress: true,
..Default::default()
},
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TrifectaVerdict {
Ok,
RefusedTrifecta,
AllowedWithWarning,
}
impl TrifectaVerdict {
pub fn is_refused(self) -> bool {
matches!(self, TrifectaVerdict::RefusedTrifecta)
}
}
pub fn check_trifecta<I>(tags: I, allow_trifecta: bool) -> TrifectaVerdict
where
I: IntoIterator<Item = TrifectaTag>,
{
let budget = tags
.into_iter()
.fold(Trifecta::default(), |acc, t| acc.merge(t.as_trifecta()));
match evaluate(budget, allow_trifecta) {
RuleOfTwo::Ok => TrifectaVerdict::Ok,
RuleOfTwo::Warn => TrifectaVerdict::AllowedWithWarning,
RuleOfTwo::Refuse => TrifectaVerdict::RefusedTrifecta,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn scope_all_allows_everything() {
assert!(Scope::All.allows("anything"));
}
#[test]
fn scope_only_is_a_whitelist() {
let s = Scope::only(["read_file", "list_dir"]);
assert!(s.allows("read_file"));
assert!(!s.allows("write_file"));
}
#[test]
fn narrow_never_widens() {
let parent = Scope::only(["a", "b"]);
assert_eq!(parent.narrow(&Scope::All), parent);
let child = Scope::only(["a", "c"]);
assert_eq!(parent.narrow(&child), Scope::only(["a"]));
assert_eq!(Scope::All.narrow(&child), child);
}
#[test]
fn tool_scope_requires_both_dimensions() {
let scope = ToolScope {
servers: Scope::only(["fs"]),
tools: Scope::only(["read_file"]),
};
assert!(scope.allows("fs", "read_file"));
assert!(!scope.allows("github", "read_file")); assert!(!scope.allows("fs", "write_file")); }
#[test]
fn tool_scope_narrow_intersects_both() {
let parent = ToolScope {
servers: Scope::only(["fs", "db"]),
tools: Scope::All,
};
let child = ToolScope {
servers: Scope::only(["fs", "net"]),
tools: Scope::only(["read"]),
};
let n = parent.narrow(&child);
assert_eq!(n.servers, Scope::only(["fs"]));
assert_eq!(n.tools, Scope::only(["read"]));
}
#[test]
fn rule_of_two() {
let two = Trifecta {
untrusted_input: true,
sensitive: true,
egress: false,
};
assert_eq!(evaluate(two, false), RuleOfTwo::Ok);
let three = Trifecta {
untrusted_input: true,
sensitive: true,
egress: true,
};
assert_eq!(evaluate(three, false), RuleOfTwo::Refuse);
assert_eq!(evaluate(three, true), RuleOfTwo::Warn);
assert_eq!(three.legs(), 3);
}
#[test]
fn trifecta_merge_accumulates() {
let a = Trifecta {
untrusted_input: true,
..Default::default()
};
let b = Trifecta {
egress: true,
..Default::default()
};
assert_eq!(a.merge(b).legs(), 2);
}
use TrifectaTag::{Egress, Sensitive, UntrustedInput};
#[test]
fn tag_maps_to_single_leg() {
assert_eq!(UntrustedInput.as_trifecta().legs(), 1);
assert_eq!(Sensitive.as_trifecta().legs(), 1);
assert_eq!(Egress.as_trifecta().legs(), 1);
assert!(UntrustedInput.as_trifecta().untrusted_input);
assert!(Sensitive.as_trifecta().sensitive);
assert!(Egress.as_trifecta().egress);
}
#[test]
fn empty_grant_is_ok() {
assert_eq!(check_trifecta([], false), TrifectaVerdict::Ok);
}
#[test]
fn each_single_leg_is_ok() {
for tag in [UntrustedInput, Sensitive, Egress] {
assert_eq!(check_trifecta([tag], false), TrifectaVerdict::Ok);
}
}
#[test]
fn every_pair_is_allowed() {
let pairs = [
[UntrustedInput, Sensitive],
[UntrustedInput, Egress],
[Sensitive, Egress],
];
for pair in pairs {
assert_eq!(
check_trifecta(pair, false),
TrifectaVerdict::Ok,
"pair {pair:?} should be allowed"
);
assert_eq!(check_trifecta(pair, true), TrifectaVerdict::Ok);
}
}
#[test]
fn all_three_refused_without_override() {
assert_eq!(
check_trifecta([UntrustedInput, Sensitive, Egress], false),
TrifectaVerdict::RefusedTrifecta
);
}
#[test]
fn all_three_warns_with_override() {
assert_eq!(
check_trifecta([UntrustedInput, Sensitive, Egress], true),
TrifectaVerdict::AllowedWithWarning
);
}
#[test]
fn duplicate_tags_do_not_inflate_legs() {
assert_eq!(
check_trifecta([Egress, Egress, Egress], false),
TrifectaVerdict::Ok
);
assert_eq!(
check_trifecta([Sensitive, Sensitive, Egress, Egress], false),
TrifectaVerdict::Ok
);
}
#[test]
fn only_refused_blocks_the_chokepoint() {
assert!(TrifectaVerdict::RefusedTrifecta.is_refused());
assert!(!TrifectaVerdict::Ok.is_refused());
assert!(!TrifectaVerdict::AllowedWithWarning.is_refused());
}
#[test]
fn tag_serde_roundtrips_snake_case() {
let json = serde_json::to_string(&UntrustedInput).unwrap();
assert_eq!(json, "\"untrusted_input\"");
let back: TrifectaTag = serde_json::from_str("\"egress\"").unwrap();
assert_eq!(back, Egress);
}
}