mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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:
committed by
GitHub
Unverified
parent
127e307f89
commit
1d09ac89a1
@@ -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")
|
||||
"#
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user