use std::collections::HashSet;
use std::sync::Mutex;
use super::rules::Decision;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApprovalOutcome {
Deny,
Allow,
AllowForSession,
}
#[derive(Debug, Clone)]
pub struct ApprovalRequest<'a> {
pub tool: &'a str,
pub subject: Option<&'a str>,
pub raw_args: &'a serde_json::Value,
}
pub trait PermissionsApprovalHandler: Send + Sync {
fn ask(&self, req: &ApprovalRequest) -> ApprovalOutcome;
}
#[derive(Debug, Default)]
pub struct ApprovalCache {
granted: Mutex<HashSet<String>>,
}
impl ApprovalCache {
pub fn new() -> Self {
ApprovalCache::default()
}
pub fn key(tool: &str, subject: Option<&str>) -> String {
match subject {
Some(s) => length_prefixed(&[tool, s]),
None => length_prefixed(&[tool]),
}
}
pub fn key_for_request(req: &ApprovalRequest) -> String {
match req.subject {
Some(s) => Self::key(req.tool, Some(s)),
None => length_prefixed(&[req.tool, "args", &canonical_json_string(req.raw_args)]),
}
}
pub fn is_approved(&self, key: &str) -> bool {
self.granted
.lock()
.map(|g| g.contains(key))
.unwrap_or(false)
}
pub fn approve(&self, key: &str) {
if let Ok(mut g) = self.granted.lock() {
g.insert(key.to_string());
}
}
}
fn length_prefixed(parts: &[&str]) -> String {
let mut out = String::new();
for p in parts {
out.push_str(&p.len().to_string());
out.push(':');
out.push_str(p);
}
out
}
fn canonical_json_string(value: &serde_json::Value) -> String {
fn sorted(value: &serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::Object(map) => {
let ordered: std::collections::BTreeMap<&String, &serde_json::Value> =
map.iter().collect();
serde_json::Value::Object(
ordered
.into_iter()
.map(|(k, v)| (k.clone(), sorted(v)))
.collect(),
)
}
serde_json::Value::Array(items) => {
serde_json::Value::Array(items.iter().map(sorted).collect())
}
other => other.clone(),
}
}
sorted(value).to_string()
}
pub fn resolve_ask(
cache: &ApprovalCache,
handler: Option<&dyn PermissionsApprovalHandler>,
req: &ApprovalRequest,
) -> bool {
let key = ApprovalCache::key_for_request(req);
if cache.is_approved(&key) {
return true;
}
match handler {
None => false,
Some(h) => match h.ask(req) {
ApprovalOutcome::Deny => false,
ApprovalOutcome::Allow => true,
ApprovalOutcome::AllowForSession => {
cache.approve(&key);
true
}
},
}
}
pub fn decision_to_approved(decision: Decision, ask_fn: impl FnOnce() -> bool) -> bool {
match decision {
Decision::Deny => false,
Decision::Allow => true,
Decision::Ask => ask_fn(),
}
}
#[cfg(test)]
mod tests {
use super::*;
struct AlwaysAllow;
impl PermissionsApprovalHandler for AlwaysAllow {
fn ask(&self, _req: &ApprovalRequest) -> ApprovalOutcome {
ApprovalOutcome::Allow
}
}
struct AlwaysDeny;
impl PermissionsApprovalHandler for AlwaysDeny {
fn ask(&self, _req: &ApprovalRequest) -> ApprovalOutcome {
ApprovalOutcome::Deny
}
}
struct AlwaysAllowForSession;
impl PermissionsApprovalHandler for AlwaysAllowForSession {
fn ask(&self, _req: &ApprovalRequest) -> ApprovalOutcome {
ApprovalOutcome::AllowForSession
}
}
fn req<'a>(
tool: &'a str,
subject: Option<&'a str>,
args: &'a serde_json::Value,
) -> ApprovalRequest<'a> {
ApprovalRequest {
tool,
subject,
raw_args: args,
}
}
#[test]
fn no_handler_denies_fail_closed() {
let cache = ApprovalCache::new();
let args = serde_json::json!({});
assert!(!resolve_ask(
&cache,
None,
&req("bash", Some("rm -rf /"), &args)
));
}
#[test]
fn handler_deny_is_denied_and_not_cached() {
let cache = ApprovalCache::new();
let args = serde_json::json!({});
let h = AlwaysDeny;
assert!(!resolve_ask(
&cache,
Some(&h),
&req("bash", Some("ls"), &args)
));
assert!(!cache.is_approved(&ApprovalCache::key("bash", Some("ls"))));
}
#[test]
fn handler_allow_once_is_not_cached() {
let cache = ApprovalCache::new();
let args = serde_json::json!({});
let h = AlwaysAllow;
assert!(resolve_ask(
&cache,
Some(&h),
&req("bash", Some("ls"), &args)
));
assert!(!cache.is_approved(&ApprovalCache::key("bash", Some("ls"))));
}
#[test]
fn allow_for_session_is_cached_and_skips_handler_next_time() {
let cache = ApprovalCache::new();
let args = serde_json::json!({});
let h = AlwaysAllowForSession;
assert!(resolve_ask(
&cache,
Some(&h),
&req("bash", Some("ls -la"), &args)
));
assert!(cache.is_approved(&ApprovalCache::key("bash", Some("ls -la"))));
let deny = AlwaysDeny;
assert!(resolve_ask(
&cache,
Some(&deny),
&req("bash", Some("ls -la"), &args)
));
}
#[test]
fn cache_key_distinguishes_subjects() {
let cache = ApprovalCache::new();
cache.approve(&ApprovalCache::key("bash", Some("ls -la")));
assert!(cache.is_approved(&ApprovalCache::key("bash", Some("ls -la"))));
assert!(!cache.is_approved(&ApprovalCache::key("bash", Some("rm -rf /"))));
assert!(!cache.is_approved(&ApprovalCache::key("write_file", Some("ls -la"))));
}
#[test]
fn decision_to_approved_short_circuits_deny_and_allow() {
assert!(!decision_to_approved(Decision::Deny, || panic!(
"must not call"
)));
assert!(decision_to_approved(Decision::Allow, || panic!(
"must not call"
)));
assert!(decision_to_approved(Decision::Ask, || true));
assert!(!decision_to_approved(Decision::Ask, || false));
}
#[test]
fn allow_for_session_on_a_subject_less_tool_does_not_leak_to_different_args() {
let cache = ApprovalCache::new();
let handler = AlwaysAllowForSession;
let safe_args = serde_json::json!({"sql": "SELECT 1"});
let safe = req("mcp_db_query", None, &safe_args);
assert!(resolve_ask(&cache, Some(&handler), &safe));
assert!(resolve_ask(&cache, Some(&handler), &safe));
let deny = AlwaysDeny;
let dangerous_args = serde_json::json!({"sql": "DROP TABLE users"});
let dangerous = req("mcp_db_query", None, &dangerous_args);
assert!(
!resolve_ask(&cache, Some(&deny), &dangerous),
"a session grant for one arg set must not leak to a different one"
);
}
#[test]
fn key_for_request_is_per_args_when_subject_is_none() {
let a = serde_json::json!({"sql": "SELECT 1"});
let b = serde_json::json!({"sql": "DROP TABLE users"});
let key_a = ApprovalCache::key_for_request(&req("mcp_db_query", None, &a));
let key_b = ApprovalCache::key_for_request(&req("mcp_db_query", None, &b));
assert_ne!(key_a, key_b);
let a_reordered = serde_json::json!({
"extra": "same",
"sql": "SELECT 1",
});
let a2 = serde_json::json!({
"sql": "SELECT 1",
"extra": "same",
});
let key_a_reordered =
ApprovalCache::key_for_request(&req("mcp_db_query", None, &a_reordered));
let key_a2 = ApprovalCache::key_for_request(&req("mcp_db_query", None, &a2));
assert_eq!(key_a_reordered, key_a2);
}
#[test]
fn key_for_request_is_unchanged_for_bash_and_file_tools() {
let args = serde_json::json!({"command": "git status"});
let r = req("bash", Some("git status"), &args);
assert_eq!(
ApprovalCache::key_for_request(&r),
ApprovalCache::key("bash", Some("git status"))
);
}
#[test]
fn bash_allow_for_session_still_does_not_leak_to_a_different_command() {
let cache = ApprovalCache::new();
let handler = AlwaysAllowForSession;
let status_args = serde_json::json!({"command": "git status"});
assert!(resolve_ask(
&cache,
Some(&handler),
&req("bash", Some("git status"), &status_args)
));
let deny = AlwaysDeny;
let push_args = serde_json::json!({"command": "git push"});
assert!(!resolve_ask(
&cache,
Some(&deny),
&req("bash", Some("git push"), &push_args)
));
}
#[test]
fn canonical_json_string_sorts_nested_objects_and_arrays() {
let a = serde_json::json!({"z": 1, "a": {"y": 2, "b": 3}, "list": [{"n": 2, "m": 1}]});
let b = serde_json::json!({"a": {"b": 3, "y": 2}, "z": 1, "list": [{"m": 1, "n": 2}]});
assert_eq!(canonical_json_string(&a), canonical_json_string(&b));
}
#[test]
fn key_collision_probe_bash_subject_vs_colon_named_tool_now_distinct() {
let empty_args = serde_json::json!({});
let subject_key = ApprovalCache::key("bash", Some("x\0args:{}"));
let colon_named_tool_req = req("bash:x", None, &empty_args);
let args_digest_key = ApprovalCache::key_for_request(&colon_named_tool_req);
assert_ne!(
subject_key, args_digest_key,
"a subject embedding the old separator bytes must not collide with an \
unrelated colon-named tool's subject-less key"
);
}
#[test]
fn key_collision_probe_does_not_let_one_grant_cover_the_other() {
let cache = ApprovalCache::new();
let empty_args = serde_json::json!({});
let bash_with_tricky_subject = req("bash", Some("x\0args:{}"), &empty_args);
cache.approve(&ApprovalCache::key_for_request(&bash_with_tricky_subject));
let colon_named_tool_req = req("bash:x", None, &empty_args);
assert!(
!cache.is_approved(&ApprovalCache::key_for_request(&colon_named_tool_req)),
"granting the tricky-subject bash call must not also grant the \
unrelated colon-named, subject-less tool call"
);
}
#[test]
fn length_prefixed_distinguishes_differing_arity_and_embedded_separators() {
assert_ne!(
length_prefixed(&["bash", "x\0args:{}"]),
length_prefixed(&["bash:x", "args", "{}"])
);
assert_ne!(
length_prefixed(&["a:b", "c"]),
length_prefixed(&["a", "b:c"])
);
}
}