mirror of
https://github.com/openai/codex.git
synced 2026-09-15 12:08:01 +00:00
feat: integrating heuristics-based fallback in execpolicy
This commit is contained in:
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user