execpolicy helpers

codex-rs/execpolicy/Cargo.toml

@@ -28,3 +28,4 @@ thiserror = { workspace = true }
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }
+tempfile = { workspace = true }

codex-rs/execpolicy/src/amend.rs

@@ -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")
+"#
+        );
+    }
+}

codex-rs/execpolicy/src/lib.rs

@@ -1,10 +1,13 @@
+pub mod amend;
 pub mod decision;
 pub mod error;
 pub mod execpolicycheck;
 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;

codex-rs/execpolicy/src/policy.rs

@@ -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) {