feat: integrating heuristics-based fallback in execpolicy

This commit is contained in:
kevin zhao
2025-12-01 18:36:05 -05:00
parent 134b48d95a
commit f65ea23746
4 changed files with 222 additions and 71 deletions

View File

@@ -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<PathBuf>,
/// 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<String>,
},
}
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<PathBuf>, args: Vec<String>, 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<codex_execpolicy::Policy> {
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())
}

View File

@@ -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<RuleMatch> =
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<F>(&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<Commands>(&self, commands: Commands) -> Evaluation
pub fn check_multiple<Commands, F>(
&self,
commands: Commands,
heuristics_fallback: &F,
) -> Evaluation
where
Commands: IntoIterator,
Commands::Item: AsRef<[String]>,
F: Fn(&[String]) -> Decision,
{
let matched_rules: Vec<RuleMatch> = 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<F>(&self, cmd: &[String], heuristics_fallback: &F) -> Vec<RuleMatch>
where
F: Fn(&[String]) -> Decision,
{
let mut matched_rules: Vec<RuleMatch> = 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<RuleMatch>,
},
pub struct Evaluation {
pub decision: Decision,
#[serde(rename = "matchedRules")]
pub matched_rules: Vec<RuleMatch>,
}
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<RuleMatch>) -> Self {
let decision = matched_rules
.iter()
.map(RuleMatch::decision)
.max()
.unwrap_or(Decision::Allow);
Self {
decision,
matched_rules,
}
}
}

View File

@@ -64,12 +64,17 @@ pub enum RuleMatch {
matched_prefix: Vec<String>,
decision: Decision,
},
HeuristicsRuleMatch {
command: Vec<String>,
decision: Decision,
},
}
impl RuleMatch {
pub fn decision(&self) -> Decision {
match self {
Self::PrefixRuleMatch { decision, .. } => *decision,
Self::HeuristicsRuleMatch { decision, .. } => *decision,
}
}
}

View File

@@ -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<String> {
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")
);
}