execpolicy helpers (#7032)

this PR 
- adds a helper function to amend `.codexpolicy` files with new prefix
rules
- adds a utility to `Policy` allowing prefix rules to be added to
existing `Policy` structs

both additions will be helpful as we thread codexpolicy into the TUI
workflow
This commit is contained in:
zhao-oai
2025-12-02 15:05:27 -05:00
committed by GitHub
Unverified
parent 127e307f89
commit 1d09ac89a1
6 changed files with 336 additions and 35 deletions
+226
View File
@@ -0,0 +1,226 @@
use std::fs::OpenOptions;
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use serde_json;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum AmendError {
#[error("prefix rule requires at least one token")]
EmptyPrefix,
#[error("policy path has no parent: {path}")]
MissingParent { path: PathBuf },
#[error("failed to create policy directory {dir}: {source}")]
CreatePolicyDir {
dir: PathBuf,
source: std::io::Error,
},
#[error("failed to format prefix tokens: {source}")]
SerializePrefix { source: serde_json::Error },
#[error("failed to open policy file {path}: {source}")]
OpenPolicyFile {
path: PathBuf,
source: std::io::Error,
},
#[error("failed to write to policy file {path}: {source}")]
WritePolicyFile {
path: PathBuf,
source: std::io::Error,
},
#[error("failed to lock policy file {path}: {source}")]
LockPolicyFile {
path: PathBuf,
source: std::io::Error,
},
#[error("failed to seek policy file {path}: {source}")]
SeekPolicyFile {
path: PathBuf,
source: std::io::Error,
},
#[error("failed to read policy file {path}: {source}")]
ReadPolicyFile {
path: PathBuf,
source: std::io::Error,
},
#[error("failed to read metadata for policy file {path}: {source}")]
PolicyMetadata {
path: PathBuf,
source: std::io::Error,
},
}
/// Note this thread uses advisory file locking and performs blocking I/O, so it should be used with
/// [`tokio::task::spawn_blocking`] when called from an async context.
pub fn blocking_append_allow_prefix_rule(
policy_path: &Path,
prefix: &[String],
) -> Result<(), AmendError> {
if prefix.is_empty() {
return Err(AmendError::EmptyPrefix);
}
let tokens = prefix
.iter()
.map(serde_json::to_string)
.collect::<Result<Vec<_>, _>>()
.map_err(|source| AmendError::SerializePrefix { source })?;
let pattern = format!("[{}]", tokens.join(", "));
let rule = format!(r#"prefix_rule(pattern={pattern}, decision="allow")"#);
let dir = policy_path
.parent()
.ok_or_else(|| AmendError::MissingParent {
path: policy_path.to_path_buf(),
})?;
match std::fs::create_dir(dir) {
Ok(()) => {}
Err(ref source) if source.kind() == std::io::ErrorKind::AlreadyExists => {}
Err(source) => {
return Err(AmendError::CreatePolicyDir {
dir: dir.to_path_buf(),
source,
});
}
}
append_locked_line(policy_path, &rule)
}
fn append_locked_line(policy_path: &Path, line: &str) -> Result<(), AmendError> {
let mut file = OpenOptions::new()
.create(true)
.read(true)
.append(true)
.open(policy_path)
.map_err(|source| AmendError::OpenPolicyFile {
path: policy_path.to_path_buf(),
source,
})?;
file.lock().map_err(|source| AmendError::LockPolicyFile {
path: policy_path.to_path_buf(),
source,
})?;
let len = file
.metadata()
.map_err(|source| AmendError::PolicyMetadata {
path: policy_path.to_path_buf(),
source,
})?
.len();
// Ensure file ends in a newline before appending.
if len > 0 {
file.seek(SeekFrom::End(-1))
.map_err(|source| AmendError::SeekPolicyFile {
path: policy_path.to_path_buf(),
source,
})?;
let mut last = [0; 1];
file.read_exact(&mut last)
.map_err(|source| AmendError::ReadPolicyFile {
path: policy_path.to_path_buf(),
source,
})?;
if last[0] != b'\n' {
file.write_all(b"\n")
.map_err(|source| AmendError::WritePolicyFile {
path: policy_path.to_path_buf(),
source,
})?;
}
}
file.write_all(format!("{line}\n").as_bytes())
.map_err(|source| AmendError::WritePolicyFile {
path: policy_path.to_path_buf(),
source,
})?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
use tempfile::tempdir;
#[test]
fn appends_rule_and_creates_directories() {
let tmp = tempdir().expect("create temp dir");
let policy_path = tmp.path().join("policy").join("default.codexpolicy");
blocking_append_allow_prefix_rule(
&policy_path,
&[String::from("echo"), String::from("Hello, world!")],
)
.expect("append rule");
let contents =
std::fs::read_to_string(&policy_path).expect("default.codexpolicy should exist");
assert_eq!(
contents,
r#"prefix_rule(pattern=["echo", "Hello, world!"], decision="allow")
"#
);
}
#[test]
fn appends_rule_without_duplicate_newline() {
let tmp = tempdir().expect("create temp dir");
let policy_path = tmp.path().join("policy").join("default.codexpolicy");
std::fs::create_dir_all(policy_path.parent().unwrap()).expect("create policy dir");
std::fs::write(
&policy_path,
r#"prefix_rule(pattern=["ls"], decision="allow")
"#,
)
.expect("write seed rule");
blocking_append_allow_prefix_rule(
&policy_path,
&[String::from("echo"), String::from("Hello, world!")],
)
.expect("append rule");
let contents = std::fs::read_to_string(&policy_path).expect("read policy");
assert_eq!(
contents,
r#"prefix_rule(pattern=["ls"], decision="allow")
prefix_rule(pattern=["echo", "Hello, world!"], decision="allow")
"#
);
}
#[test]
fn inserts_newline_when_missing_before_append() {
let tmp = tempdir().expect("create temp dir");
let policy_path = tmp.path().join("policy").join("default.codexpolicy");
std::fs::create_dir_all(policy_path.parent().unwrap()).expect("create policy dir");
std::fs::write(
&policy_path,
r#"prefix_rule(pattern=["ls"], decision="allow")"#,
)
.expect("write seed rule without newline");
blocking_append_allow_prefix_rule(
&policy_path,
&[String::from("echo"), String::from("Hello, world!")],
)
.expect("append rule");
let contents = std::fs::read_to_string(&policy_path).expect("read policy");
assert_eq!(
contents,
r#"prefix_rule(pattern=["ls"], decision="allow")
prefix_rule(pattern=["echo", "Hello, world!"], decision="allow")
"#
);
}
}
+3
View File
@@ -1,3 +1,4 @@
pub mod amend;
pub mod decision;
pub mod error;
pub mod execpolicycheck;
@@ -5,6 +6,8 @@ pub mod parser;
pub mod policy;
pub mod rule;
pub use amend::AmendError;
pub use amend::blocking_append_allow_prefix_rule;
pub use decision::Decision;
pub use error::Error;
pub use error::Result;
+27
View File
@@ -1,9 +1,15 @@
use crate::decision::Decision;
use crate::error::Error;
use crate::error::Result;
use crate::rule::PatternToken;
use crate::rule::PrefixPattern;
use crate::rule::PrefixRule;
use crate::rule::RuleMatch;
use crate::rule::RuleRef;
use multimap::MultiMap;
use serde::Deserialize;
use serde::Serialize;
use std::sync::Arc;
#[derive(Clone, Debug)]
pub struct Policy {
@@ -23,6 +29,27 @@ impl Policy {
&self.rules_by_program
}
pub fn add_prefix_rule(&mut self, prefix: &[String], decision: Decision) -> Result<()> {
let (first_token, rest) = prefix
.split_first()
.ok_or_else(|| Error::InvalidPattern("prefix cannot be empty".to_string()))?;
let rule: RuleRef = Arc::new(PrefixRule {
pattern: PrefixPattern {
first: Arc::from(first_token.as_str()),
rest: rest
.iter()
.map(|token| PatternToken::Single(token.clone()))
.collect::<Vec<_>>()
.into(),
},
decision,
});
self.rules_by_program.insert(first_token.clone(), rule);
Ok(())
}
pub fn check(&self, cmd: &[String]) -> Evaluation {
let rules = match cmd.first() {
Some(first) => match self.rules_by_program.get_vec(first) {