From f65ea23746c5ef7d09ce3f581bc5093d9ef566a9 Mon Sep 17 00:00:00 2001 From: kevin zhao Date: Mon, 1 Dec 2025 18:36:05 -0500 Subject: [PATCH] feat: integrating heuristics-based fallback in execpolicy --- codex-rs/execpolicy/src/main.rs | 55 ++++++++++-- codex-rs/execpolicy/src/policy.rs | 93 +++++++++++-------- codex-rs/execpolicy/src/rule.rs | 5 ++ codex-rs/execpolicy/tests/basic.rs | 140 +++++++++++++++++++++++------ 4 files changed, 222 insertions(+), 71 deletions(-) diff --git a/codex-rs/execpolicy/src/main.rs b/codex-rs/execpolicy/src/main.rs index e1373b6d16..5b5296890a 100644 --- a/codex-rs/execpolicy/src/main.rs +++ b/codex-rs/execpolicy/src/main.rs @@ -1,22 +1,67 @@ +use std::fs; +use std::path::PathBuf; + +use anyhow::Context; use anyhow::Result; use clap::Parser; -use codex_execpolicy::ExecPolicyCheckCommand; +use codex_execpolicy::Decision; +use codex_execpolicy::PolicyParser; /// CLI for evaluating exec policies #[derive(Parser)] #[command(name = "codex-execpolicy")] enum Cli { /// Evaluate a command against a policy. - Check(ExecPolicyCheckCommand), + Check { + #[arg(short, long = "policy", value_name = "PATH", required = true)] + policies: Vec, + + /// Pretty-print the JSON output. + #[arg(long)] + pretty: bool, + + /// Command tokens to check. + #[arg( + value_name = "COMMAND", + required = true, + trailing_var_arg = true, + allow_hyphen_values = true + )] + command: Vec, + }, } fn main() -> Result<()> { let cli = Cli::parse(); match cli { - Cli::Check(cmd) => cmd_check(cmd), + Cli::Check { + policies, + command, + pretty, + } => cmd_check(policies, command, pretty), } } -fn cmd_check(cmd: ExecPolicyCheckCommand) -> Result<()> { - cmd.run() +fn cmd_check(policy_paths: Vec, args: Vec, pretty: bool) -> Result<()> { + let policy = load_policies(&policy_paths)?; + + let eval = policy.check(&args, &|_| Decision::Allow); + let json = if pretty { + serde_json::to_string_pretty(&eval)? + } else { + serde_json::to_string(&eval)? + }; + println!("{json}"); + Ok(()) +} + +fn load_policies(policy_paths: &[PathBuf]) -> Result { + let mut parser = PolicyParser::new(); + for policy_path in policy_paths { + let policy_file_contents = fs::read_to_string(policy_path) + .with_context(|| format!("failed to read policy at {}", policy_path.display()))?; + let policy_identifier = policy_path.to_string_lossy().to_string(); + parser.parse(&policy_identifier, &policy_file_contents)?; + } + Ok(parser.build()) } diff --git a/codex-rs/execpolicy/src/policy.rs b/codex-rs/execpolicy/src/policy.rs index 10858c9fad..6da5857c37 100644 --- a/codex-rs/execpolicy/src/policy.rs +++ b/codex-rs/execpolicy/src/policy.rs @@ -50,62 +50,81 @@ impl Policy { Ok(()) } - pub fn check(&self, cmd: &[String]) -> Evaluation { - let rules = match cmd.first() { - Some(first) => match self.rules_by_program.get_vec(first) { - Some(rules) => rules, - None => return Evaluation::NoMatch {}, - }, - None => return Evaluation::NoMatch {}, - }; - - let matched_rules: Vec = - rules.iter().filter_map(|rule| rule.matches(cmd)).collect(); - match matched_rules.iter().map(RuleMatch::decision).max() { - Some(decision) => Evaluation::Match { - decision, - matched_rules, - }, - None => Evaluation::NoMatch {}, - } + pub fn check(&self, cmd: &[String], heuristics_fallback: &F) -> Evaluation + where + F: Fn(&[String]) -> Decision, + { + let matched_rules = self.matches_for_command(cmd, heuristics_fallback); + Evaluation::from_matches(matched_rules) } - pub fn check_multiple(&self, commands: Commands) -> Evaluation + pub fn check_multiple( + &self, + commands: Commands, + heuristics_fallback: &F, + ) -> Evaluation where Commands: IntoIterator, Commands::Item: AsRef<[String]>, + F: Fn(&[String]) -> Decision, { let matched_rules: Vec = commands .into_iter() - .flat_map(|command| match self.check(command.as_ref()) { - Evaluation::Match { matched_rules, .. } => matched_rules, - Evaluation::NoMatch { .. } => Vec::new(), - }) + .flat_map(|command| self.matches_for_command(command.as_ref(), heuristics_fallback)) .collect(); - match matched_rules.iter().map(RuleMatch::decision).max() { - Some(decision) => Evaluation::Match { - decision, - matched_rules, - }, - None => Evaluation::NoMatch {}, + Evaluation::from_matches(matched_rules) + } + + fn matches_for_command(&self, cmd: &[String], heuristics_fallback: &F) -> Vec + where + F: Fn(&[String]) -> Decision, + { + let mut matched_rules: Vec = match cmd.first() { + Some(first) => self + .rules_by_program + .get_vec(first) + .map(|rules| rules.iter().filter_map(|rule| rule.matches(cmd)).collect()) + .unwrap_or_default(), + None => Vec::new(), + }; + + if matched_rules.is_empty() { + matched_rules.push(RuleMatch::HeuristicsRuleMatch { + command: cmd.to_vec(), + decision: heuristics_fallback(cmd), + }); } + + matched_rules } } #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub enum Evaluation { - NoMatch {}, - Match { - decision: Decision, - #[serde(rename = "matchedRules")] - matched_rules: Vec, - }, +pub struct Evaluation { + pub decision: Decision, + #[serde(rename = "matchedRules")] + pub matched_rules: Vec, } impl Evaluation { pub fn is_match(&self) -> bool { - matches!(self, Self::Match { .. }) + self.matched_rules + .iter() + .any(|rule_match| !matches!(rule_match, RuleMatch::HeuristicsRuleMatch { .. })) + } + + fn from_matches(matched_rules: Vec) -> Self { + let decision = matched_rules + .iter() + .map(RuleMatch::decision) + .max() + .unwrap_or(Decision::Allow); + + Self { + decision, + matched_rules, + } } } diff --git a/codex-rs/execpolicy/src/rule.rs b/codex-rs/execpolicy/src/rule.rs index 20e23fe6a2..cd0756bbb3 100644 --- a/codex-rs/execpolicy/src/rule.rs +++ b/codex-rs/execpolicy/src/rule.rs @@ -64,12 +64,17 @@ pub enum RuleMatch { matched_prefix: Vec, decision: Decision, }, + HeuristicsRuleMatch { + command: Vec, + decision: Decision, + }, } impl RuleMatch { pub fn decision(&self) -> Decision { match self { Self::PrefixRuleMatch { decision, .. } => *decision, + Self::HeuristicsRuleMatch { decision, .. } => *decision, } } } diff --git a/codex-rs/execpolicy/tests/basic.rs b/codex-rs/execpolicy/tests/basic.rs index 9a7ec58b1e..654291cc98 100644 --- a/codex-rs/execpolicy/tests/basic.rs +++ b/codex-rs/execpolicy/tests/basic.rs @@ -1,5 +1,6 @@ use std::any::Any; use std::sync::Arc; +use std::sync::Mutex; use anyhow::Context; use anyhow::Result; @@ -19,6 +20,14 @@ fn tokens(cmd: &[&str]) -> Vec { cmd.iter().map(std::string::ToString::to_string).collect() } +fn allow_all(_: &[String]) -> Decision { + Decision::Allow +} + +fn prompt_all(_: &[String]) -> Decision { + Decision::Prompt +} + #[derive(Clone, Debug, Eq, PartialEq)] enum RuleSnapshot { Prefix(PrefixRule), @@ -49,9 +58,9 @@ prefix_rule( parser.parse("test.codexpolicy", policy_src)?; let policy = parser.build(); let cmd = tokens(&["git", "status"]); - let evaluation = policy.check(&cmd); + let evaluation = policy.check(&cmd, &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Allow, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["git", "status"]), @@ -80,9 +89,9 @@ fn add_prefix_rule_extends_policy() -> Result<()> { rules ); - let evaluation = policy.check(&tokens(&["ls", "-l", "/tmp"])); + let evaluation = policy.check(&tokens(&["ls", "-l", "/tmp"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Prompt, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["ls", "-l"]), @@ -146,9 +155,9 @@ prefix_rule( git_rules ); - let status_eval = policy.check(&tokens(&["git", "status"])); + let status_eval = policy.check(&tokens(&["git", "status"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Prompt, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["git"]), @@ -158,9 +167,9 @@ prefix_rule( status_eval ); - let commit_eval = policy.check(&tokens(&["git", "commit", "-m", "hi"])); + let commit_eval = policy.check(&tokens(&["git", "commit", "-m", "hi"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Forbidden, matched_rules: vec![ RuleMatch::PrefixRuleMatch { @@ -217,9 +226,9 @@ prefix_rule( sh_rules ); - let bash_eval = policy.check(&tokens(&["bash", "-c", "echo", "hi"])); + let bash_eval = policy.check(&tokens(&["bash", "-c", "echo", "hi"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Allow, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["bash", "-c"]), @@ -229,9 +238,9 @@ prefix_rule( bash_eval ); - let sh_eval = policy.check(&tokens(&["sh", "-l", "echo", "hi"])); + let sh_eval = policy.check(&tokens(&["sh", "-l", "echo", "hi"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Allow, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["sh", "-l"]), @@ -273,9 +282,9 @@ prefix_rule( rules ); - let npm_i = policy.check(&tokens(&["npm", "i", "--legacy-peer-deps"])); + let npm_i = policy.check(&tokens(&["npm", "i", "--legacy-peer-deps"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Allow, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["npm", "i", "--legacy-peer-deps"]), @@ -285,9 +294,12 @@ prefix_rule( npm_i ); - let npm_install = policy.check(&tokens(&["npm", "install", "--no-save", "leftpad"])); + let npm_install = policy.check( + &tokens(&["npm", "install", "--no-save", "leftpad"]), + &allow_all, + ); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Allow, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["npm", "install", "--no-save"]), @@ -314,9 +326,9 @@ prefix_rule( let mut parser = PolicyParser::new(); parser.parse("test.codexpolicy", policy_src)?; let policy = parser.build(); - let match_eval = policy.check(&tokens(&["git", "status"])); + let match_eval = policy.check(&tokens(&["git", "status"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Allow, matched_rules: vec![RuleMatch::PrefixRuleMatch { matched_prefix: tokens(&["git", "status"]), @@ -326,13 +338,20 @@ prefix_rule( match_eval ); - let no_match_eval = policy.check(&tokens(&[ - "git", - "--config", - "color.status=always", - "status", - ])); - assert_eq!(Evaluation::NoMatch {}, no_match_eval); + let no_match_eval = policy.check( + &tokens(&["git", "--config", "color.status=always", "status"]), + &allow_all, + ); + assert_eq!( + Evaluation { + decision: Decision::Allow, + matched_rules: vec![RuleMatch::HeuristicsRuleMatch { + command: tokens(&["git", "--config", "color.status=always", "status",]), + decision: Decision::Allow, + }], + }, + no_match_eval + ); Ok(()) } @@ -352,9 +371,9 @@ prefix_rule( parser.parse("test.codexpolicy", policy_src)?; let policy = parser.build(); - let commit = policy.check(&tokens(&["git", "commit", "-m", "hi"])); + let commit = policy.check(&tokens(&["git", "commit", "-m", "hi"]), &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Forbidden, matched_rules: vec![ RuleMatch::PrefixRuleMatch { @@ -393,9 +412,9 @@ prefix_rule( tokens(&["git", "commit", "-m", "hi"]), ]; - let evaluation = policy.check_multiple(&commands); + let evaluation = policy.check_multiple(&commands, &allow_all); assert_eq!( - Evaluation::Match { + Evaluation { decision: Decision::Forbidden, matched_rules: vec![ RuleMatch::PrefixRuleMatch { @@ -416,3 +435,66 @@ prefix_rule( ); Ok(()) } + +#[test] +fn heuristics_match_is_returned_when_no_policy_matches() { + let policy = Policy::empty(); + let command = tokens(&["python"]); + + let evaluation = policy.check(&command, &prompt_all); + assert_eq!( + Evaluation { + decision: Decision::Prompt, + matched_rules: vec![RuleMatch::HeuristicsRuleMatch { + command, + decision: Decision::Prompt, + }], + }, + evaluation + ); +} + +#[test] +fn heuristics_only_runs_for_commands_without_policy_matches() { + let policy_src = r#" +prefix_rule( + pattern = ["git"], + decision = "allow", +) + "#; + let mut parser = PolicyParser::new(); + parser + .parse("policy.codexpolicy", policy_src) + .expect("parse policy"); + let policy = parser.build(); + + let commands = vec![tokens(&["git", "status"]), tokens(&["python"])]; + let heuristics_calls = Arc::new(Mutex::new(Vec::new())); + let heuristics_call_log = Arc::clone(&heuristics_calls); + let heuristics = move |cmd: &[String]| { + heuristics_call_log + .lock() + .expect("lock heuristics call log") + .push(cmd.to_vec()); + Decision::Prompt + }; + + let evaluation = policy.check_multiple(&commands, &heuristics); + + assert_eq!(Decision::Prompt, evaluation.decision); + assert!(evaluation.matched_rules.iter().any(|rule_match| { + matches!( + rule_match, + RuleMatch::HeuristicsRuleMatch { + command, + decision: Decision::Prompt + } if command == &tokens(&["python"]) + ) + })); + assert_eq!( + vec![tokens(&["python"])], + *heuristics_calls + .lock() + .expect("lock heuristics call log after evaluation") + ); +}