diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8aee923822..8e609eb6ac 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1193,6 +1193,7 @@ name = "codex-execpolicy2" version = "0.0.0" dependencies = [ "anyhow", + "multimap", "serde", "serde_json", "starlark", diff --git a/codex-rs/execpolicy2/Cargo.toml b/codex-rs/execpolicy2/Cargo.toml index ceae2bd989..3c659575a6 100644 --- a/codex-rs/execpolicy2/Cargo.toml +++ b/codex-rs/execpolicy2/Cargo.toml @@ -19,3 +19,4 @@ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } starlark = { workspace = true } thiserror = { workspace = true } +multimap = { workspace = true } diff --git a/codex-rs/execpolicy2/src/parser.rs b/codex-rs/execpolicy2/src/parser.rs index acf5e4d7ba..b65e5db386 100644 --- a/codex-rs/execpolicy2/src/parser.rs +++ b/codex-rs/execpolicy2/src/parser.rs @@ -1,5 +1,6 @@ use std::cell::RefCell; +use multimap::MultiMap; use starlark::any::ProvidesStaticType; use starlark::environment::GlobalsBuilder; use starlark::environment::Module; @@ -15,6 +16,8 @@ use starlark::values::none::NoneType; use crate::decision::Decision; use crate::error::Error; use crate::error::Result; +use crate::rule::PatternToken; +use crate::rule::PrefixPattern; use crate::rule::Rule; pub struct PolicyParser { @@ -51,14 +54,14 @@ impl PolicyParser { #[derive(Debug, ProvidesStaticType)] struct PolicyBuilder { - rules: RefCell>, + rules_by_program: RefCell>, next_auto_id: RefCell, } impl PolicyBuilder { fn new() -> Self { Self { - rules: RefCell::new(Vec::new()), + rules_by_program: RefCell::new(MultiMap::new()), next_auto_id: RefCell::new(0), } } @@ -71,68 +74,81 @@ impl PolicyBuilder { } fn add_rule(&self, rule: Rule) { - self.rules.borrow_mut().push(rule); + self.rules_by_program + .borrow_mut() + .insert(rule.pattern.first.clone(), rule); } fn build(&self) -> crate::policy::Policy { - crate::policy::Policy::new(self.rules.borrow().clone()) + crate::policy::Policy::new(self.rules_by_program.borrow().clone()) } } #[derive(Debug)] -enum PatternPart { - Single(String), - Alts(Vec), +struct ParsedPattern { + heads: Vec, + tail: Vec, } -fn expand_pattern(parts: &[PatternPart]) -> Vec> { - let mut acc: Vec> = vec![Vec::new()]; - for part in parts { - let alts: Vec = match part { - PatternPart::Single(s) => vec![s.clone()], - PatternPart::Alts(v) => v.clone(), - }; - let mut next = Vec::new(); - for prefix in &acc { - for alt in &alts { - let mut combined = prefix.clone(); - combined.push(alt.clone()); - next.push(combined); - } - } - acc = next; +fn parse_pattern<'v>(pattern: UnpackList>) -> Result { + let mut items = pattern.items.into_iter(); + let first = items + .next() + .ok_or_else(|| Error::InvalidPattern("pattern cannot be empty".to_string()))?; + let heads = parse_first_token(first)?; + let mut tail = Vec::new(); + for item in items { + tail.push(parse_tail_token(item)?); } - acc + Ok(ParsedPattern { heads, tail }) } -fn parse_pattern<'v>(pattern: UnpackList>) -> Result>> { - let mut parts = Vec::new(); - for item in pattern.items { - if let Some(s) = item.unpack_str() { - parts.push(PatternPart::Single(s.to_string())); - continue; - } +fn parse_first_token<'v>(value: Value<'v>) -> Result> { + if let Some(s) = value.unpack_str() { + return Ok(vec![s.to_string()]); + } + if let Some(list) = ListRef::from_value(value) { let mut alts = Vec::new(); - if let Some(list) = ListRef::from_value(item) { - for value in list.content() { - let s = value.unpack_str().ok_or_else(|| { - Error::InvalidPattern("pattern alternative must be a string".to_string()) - })?; - alts.push(s.to_string()); - } - } else { - return Err(Error::InvalidPattern( - "pattern element must be a string or list of strings".to_string(), - )); + for value in list.content() { + let s = value.unpack_str().ok_or_else(|| { + Error::InvalidPattern("pattern alternative must be a string".to_string()) + })?; + alts.push(s.to_string()); } if alts.is_empty() { return Err(Error::InvalidPattern( "pattern alternatives cannot be empty".to_string(), )); } - parts.push(PatternPart::Alts(alts)); + return Ok(alts); } - Ok(expand_pattern(&parts)) + Err(Error::InvalidPattern( + "pattern element must be a string or list of strings".to_string(), + )) +} + +fn parse_tail_token<'v>(value: Value<'v>) -> Result { + if let Some(s) = value.unpack_str() { + return Ok(PatternToken::Single(s.to_string())); + } + if let Some(list) = ListRef::from_value(value) { + let mut alts = Vec::new(); + for value in list.content() { + let s = value.unpack_str().ok_or_else(|| { + Error::InvalidPattern("pattern alternative must be a string".to_string()) + })?; + alts.push(s.to_string()); + } + if alts.is_empty() { + return Err(Error::InvalidPattern( + "pattern alternatives cannot be empty".to_string(), + )); + } + return Ok(PatternToken::Alts(alts)); + } + Err(Error::InvalidPattern( + "pattern element must be a string or list of strings".to_string(), + )) } fn parse_examples<'v>(examples: UnpackList>) -> Result>> { @@ -173,7 +189,7 @@ fn policy_builtins(builder: &mut GlobalsBuilder) { None => Decision::Allow, }; - let prefixes = parse_pattern(pattern)?; + let parsed_pattern = parse_pattern(pattern)?; let positive_examples: Vec> = r#match.map(parse_examples).transpose()?.unwrap_or_default(); @@ -193,13 +209,6 @@ fn policy_builtins(builder: &mut GlobalsBuilder) { builder.alloc_id() }); - let rule = Rule { - id: id.clone(), - prefixes, - decision, - }; - rule.validate_examples(&positive_examples, &negative_examples)?; - #[expect(clippy::unwrap_used)] let builder = eval .extra @@ -207,7 +216,19 @@ fn policy_builtins(builder: &mut GlobalsBuilder) { .unwrap() .downcast_ref::() .unwrap(); - builder.add_rule(rule); + + for head in &parsed_pattern.heads { + let rule = Rule { + id: id.clone(), + pattern: PrefixPattern { + first: head.clone(), + tail: parsed_pattern.tail.clone(), + }, + decision, + }; + rule.validate_examples(&positive_examples, &negative_examples)?; + builder.add_rule(rule); + } Ok(NoneType) } } diff --git a/codex-rs/execpolicy2/src/policy.rs b/codex-rs/execpolicy2/src/policy.rs index bc0765f5e9..f94600fa01 100644 --- a/codex-rs/execpolicy2/src/policy.rs +++ b/codex-rs/execpolicy2/src/policy.rs @@ -1,11 +1,12 @@ use crate::decision::Decision; use crate::rule::Rule; +use multimap::MultiMap; use serde::Deserialize; use serde::Serialize; #[derive(Clone, Debug)] pub struct Policy { - rules: Vec, + rules_by_program: MultiMap, } #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] @@ -15,18 +16,22 @@ pub struct Evaluation { } impl Policy { - pub fn new(rules: Vec) -> Self { - Self { rules } + pub fn new(rules_by_program: MultiMap) -> Self { + Self { rules_by_program } } - pub fn rules(&self) -> &[Rule] { - &self.rules + pub fn rules(&self) -> &MultiMap { + &self.rules_by_program } pub fn evaluate(&self, cmd: &[String]) -> Option { + let first = cmd.first()?; + let Some(rules) = self.rules_by_program.get_vec(first) else { + return None; + }; let mut matched_rules: Vec = Vec::new(); let mut best_decision: Option = None; - for rule in &self.rules { + for rule in rules { if let Some(matched) = rule.matches(cmd) { let decision = match best_decision { None => matched.decision, diff --git a/codex-rs/execpolicy2/src/rule.rs b/codex-rs/execpolicy2/src/rule.rs index 50b2601863..5117c5d24a 100644 --- a/codex-rs/execpolicy2/src/rule.rs +++ b/codex-rs/execpolicy2/src/rule.rs @@ -4,10 +4,51 @@ use crate::error::Result; use serde::Deserialize; use serde::Serialize; +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum PatternToken { + Single(String), + Alts(Vec), +} + +impl PatternToken { + fn matches(&self, token: &str) -> bool { + match self { + Self::Single(expected) => expected == token, + Self::Alts(alternatives) => alternatives.iter().any(|alt| alt == token), + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct PrefixPattern { + pub first: String, + pub tail: Vec, +} + +impl PrefixPattern { + pub fn len(&self) -> usize { + self.tail.len() + 1 + } + + pub fn matches_prefix(&self, cmd: &[String]) -> Option> { + if cmd.len() < self.len() || cmd[0] != self.first { + return None; + } + + for (pattern_token, cmd_token) in self.tail.iter().zip(&cmd[1..self.len()]) { + if !pattern_token.matches(cmd_token) { + return None; + } + } + + Some(cmd[..self.len()].to_vec()) + } +} + #[derive(Clone, Debug)] pub struct Rule { pub id: String, - pub prefixes: Vec>, + pub pattern: PrefixPattern, pub decision: Decision, } @@ -20,23 +61,13 @@ pub struct RuleMatch { impl Rule { pub fn matches(&self, cmd: &[String]) -> Option { - for prefix in &self.prefixes { - if prefix.len() > cmd.len() { - continue; - } - if cmd - .iter() - .zip(prefix) - .all(|(cmd_tok, prefix_tok)| cmd_tok == prefix_tok) - { - return Some(RuleMatch { - rule_id: self.id.clone(), - matched_prefix: prefix.clone(), - decision: self.decision, - }); - } - } - None + self.pattern + .matches_prefix(cmd) + .map(|matched_prefix| RuleMatch { + rule_id: self.id.clone(), + matched_prefix, + decision: self.decision, + }) } pub fn validate_examples( diff --git a/codex-rs/execpolicy2/tests/basic.rs b/codex-rs/execpolicy2/tests/basic.rs index 5273409e76..448bd2aacd 100644 --- a/codex-rs/execpolicy2/tests/basic.rs +++ b/codex-rs/execpolicy2/tests/basic.rs @@ -1,6 +1,7 @@ use codex_execpolicy2::Decision; use codex_execpolicy2::PolicyParser; use codex_execpolicy2::RuleMatch; +use codex_execpolicy2::rule::PatternToken; fn tokens(cmd: &[&str]) -> Vec { cmd.iter().map(|token| token.to_string()).collect() @@ -52,6 +53,67 @@ prefix_rule( assert!(policy.evaluate(&no_match).is_none()); } +#[test] +fn only_first_token_alias_expands_to_multiple_rules() { + let policy_src = r#" +prefix_rule( + id = "shell", + pattern = [["bash", "sh"], ["-c", "-l"]], +) + "#; + let parser = PolicyParser::new("test.policy", policy_src); + let policy = parser.parse().expect("parse policy"); + + let bash_rules = policy.rules().get_vec("bash").expect("bash rules"); + let sh_rules = policy.rules().get_vec("sh").expect("sh rules"); + assert_eq!(bash_rules.len(), 1); + assert_eq!(sh_rules.len(), 1); + + for (cmd, prefix) in [ + ( + tokens(&["bash", "-c", "echo", "hi"]), + tokens(&["bash", "-c"]), + ), + (tokens(&["sh", "-l", "echo", "hi"]), tokens(&["sh", "-l"])), + ] { + let eval = policy.evaluate(&cmd).expect("match"); + assert_eq!(eval.matched_rules[0].matched_prefix, prefix); + } +} + +#[test] +fn tail_aliases_are_not_cartesian_expanded() { + let policy_src = r#" +prefix_rule( + id = "npm_install_variants", + pattern = ["npm", ["i", "install"], ["--legacy-peer-deps", "--no-save"]], +) + "#; + let parser = PolicyParser::new("test.policy", policy_src); + let policy = parser.parse().expect("parse policy"); + + let rules = policy.rules().get_vec("npm").expect("npm rules"); + assert_eq!(rules.len(), 1); + let rule = &rules[0]; + assert_eq!( + rule.pattern.tail, + vec![ + PatternToken::Alts(vec!["i".to_string(), "install".to_string()]), + PatternToken::Alts(vec![ + "--legacy-peer-deps".to_string(), + "--no-save".to_string() + ]), + ], + ); + + for cmd in [ + tokens(&["npm", "i", "--legacy-peer-deps"]), + tokens(&["npm", "install", "--no-save", "leftpad"]), + ] { + assert!(policy.evaluate(&cmd).is_some()); + } +} + #[test] fn match_and_not_match_examples_are_enforced() { let policy_src = r#"