agent: simplify ai agent enforcement config

Author: kylewanginchina

agent/crates/enterprise-utils/src/lib.rs

@@ -502,45 +502,61 @@ pub mod kernel_version {
 }
 
 pub mod ai_agent_enforcement {
-    use std::sync::Arc;
+    use std::sync::{Arc, OnceLock, RwLock};
+
+    pub const BPF_PATH_LEN: usize = 256;
+    pub const BPF_RULE_ID_LEN: usize = 64;
+    pub const BPF_EXEC_MAX_RULES: usize = 256;
+    pub const BPF_SYSCALL_NAME_LEN: usize = 32;
+    pub const BPF_SYSCALL_MAX_RULES: usize = 32;
+    pub const BPF_CMDLINE_PATTERN_MAX_PARTS: usize = 4;
+    pub const BPF_CMDLINE_PATTERN_PART_LEN: usize = 64;
 
     #[derive(Clone, Copy, Debug, PartialEq, Eq)]
-    pub enum EnforcementMode {
-        AuditOnly,
-        Block,
+    pub enum EnforcementMechanism {
+        KprobeOverride,
+        Lsm,
     }
 
-    #[derive(Clone, Debug, PartialEq, Eq)]
-    pub struct ExecRuleInput {
-        pub id: String,
-        pub mode: EnforcementMode,
-        pub exact: Vec<String>,
-        pub prefix: Vec<String>,
-        pub suffix: Vec<String>,
-        pub argv_matches: Vec<ExecArgvMatchInput>,
-        pub argv_contains_any: Vec<String>,
+    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+    pub enum EnforcementAction {
+        Audit,
+        Block,
     }
 
     #[derive(Clone, Copy, Debug, PartialEq, Eq)]
-    pub enum ExecArgvMatchOp {
-        Exact,
-        Prefix,
-        Suffix,
+    pub enum EnforcementMode {
+        AuditOnly,
+        Block,
     }
 
     #[derive(Clone, Debug, PartialEq, Eq)]
-    pub struct ExecArgvMatchInput {
-        pub index: u8,
-        pub op: ExecArgvMatchOp,
-        pub value: String,
+    pub struct AiAgentEnforcementInput {
+        pub enabled: bool,
+        pub mechanism: EnforcementMechanism,
+        pub rules: Vec<EnforcementRuleInput>,
     }
 
     #[derive(Clone, Debug, PartialEq, Eq)]
-    pub struct SyscallRuleInput {
+    pub struct EnforcementRuleInput {
         pub id: String,
-        pub mode: EnforcementMode,
-        pub names: Vec<String>,
-        pub symbols: Vec<String>,
+        pub action: EnforcementAction,
+        pub target: EnforcementTargetInput,
+    }
+
+    #[derive(Clone, Debug, PartialEq, Eq)]
+    pub enum EnforcementTargetInput {
+        ExecCmdline {
+            patterns: Vec<String>,
+        },
+        ExecPath {
+            exact: Vec<String>,
+            prefix: Vec<String>,
+            suffix: Vec<String>,
+        },
+        Syscall {
+            names: Vec<String>,
+        },
     }
 
     #[derive(Clone, Debug, PartialEq, Eq)]
@@ -558,92 +574,202 @@ pub mod ai_agent_enforcement {
         KprobeOverride,
     }
 
+    #[repr(C)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+    pub struct BpfKprobeExecRuleRecord {
+        pub rule_index: u32,
+        pub mode: u8,
+        pub pattern_flags: u8,
+        pub part_count: u8,
+        pub reserved: u8,
+        pub rule_id: [u8; BPF_RULE_ID_LEN],
+        pub part_lens: [u16; BPF_CMDLINE_PATTERN_MAX_PARTS],
+        pub parts: [[u8; BPF_CMDLINE_PATTERN_PART_LEN]; BPF_CMDLINE_PATTERN_MAX_PARTS],
+    }
+
+    #[repr(C)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+    pub struct BpfLsmExecRuleRecord {
+        pub rule_index: u32,
+        pub mode: u8,
+        pub match_type: u8,
+        pub pattern_len: u16,
+        pub rule_id: [u8; BPF_RULE_ID_LEN],
+        pub pattern: [u8; BPF_PATH_LEN],
+    }
+
+    #[repr(C)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+    pub struct BpfSyscallRuleRecord {
+        pub rule_index: u32,
+        pub mode: u8,
+        pub syscall_key: u8,
+        pub reserved: u16,
+        pub syscall_id: u32,
+        pub errno_code: i32,
+        pub rule_id: [u8; BPF_RULE_ID_LEN],
+        pub syscall_name: [u8; BPF_SYSCALL_NAME_LEN],
+    }
+
+    impl Default for BpfKprobeExecRuleRecord {
+        fn default() -> Self {
+            Self {
+                rule_index: 0,
+                mode: 0,
+                pattern_flags: 0,
+                part_count: 0,
+                reserved: 0,
+                rule_id: [0; BPF_RULE_ID_LEN],
+                part_lens: [0; BPF_CMDLINE_PATTERN_MAX_PARTS],
+                parts: [[0; BPF_CMDLINE_PATTERN_PART_LEN]; BPF_CMDLINE_PATTERN_MAX_PARTS],
+            }
+        }
+    }
+
+    impl Default for BpfLsmExecRuleRecord {
+        fn default() -> Self {
+            Self {
+                rule_index: 0,
+                mode: 0,
+                match_type: 0,
+                pattern_len: 0,
+                rule_id: [0; BPF_RULE_ID_LEN],
+                pattern: [0; BPF_PATH_LEN],
+            }
+        }
+    }
+
+    impl Default for BpfSyscallRuleRecord {
+        fn default() -> Self {
+            Self {
+                rule_index: 0,
+                mode: 0,
+                syscall_key: 0,
+                reserved: 0,
+                syscall_id: 0,
+                errno_code: 0,
+                rule_id: [0; BPF_RULE_ID_LEN],
+                syscall_name: [0; BPF_SYSCALL_NAME_LEN],
+            }
+        }
+    }
+
     #[derive(Clone, Debug, PartialEq, Eq)]
-    pub struct CompiledExecPolicy {
+    pub struct CompiledKprobePolicy {
         pub epoch: u64,
+        pub exec_rules: Vec<BpfKprobeExecRuleRecord>,
+        pub syscall_rules: Vec<BpfSyscallRuleRecord>,
     }
 
     #[derive(Clone, Debug, PartialEq, Eq)]
-    pub struct CompiledSyscallPolicy {
+    pub struct CompiledLsmPolicy {
         pub epoch: u64,
+        pub exec_rules: Vec<BpfLsmExecRuleRecord>,
     }
 
-    impl CompiledExecPolicy {
-        pub fn match_exec(&self, _exec_path: &str, _cmdline: &str) -> Option<PolicyHit> {
+    static GLOBAL_KPROBE_POLICY: OnceLock<RwLock<Option<Arc<CompiledKprobePolicy>>>> =
+        OnceLock::new();
+    static GLOBAL_LSM_POLICY: OnceLock<RwLock<Option<Arc<CompiledLsmPolicy>>>> = OnceLock::new();
+
+    impl CompiledKprobePolicy {
+        pub fn exec_records(&self) -> &[BpfKprobeExecRuleRecord] {
+            &self.exec_rules
+        }
+
+        pub fn syscall_records(&self) -> &[BpfSyscallRuleRecord] {
+            &self.syscall_rules
+        }
+
+        pub fn match_cmdline(&self, _cmdline: &str) -> Option<PolicyHit> {
             None
         }
 
         pub fn sync_to_bpf_maps(
             &self,
-            _exec_rules_fd: i32,
             _policy_epoch_fd: i32,
-            _max_records: usize,
+            _kprobe_exec_rules_fd: i32,
+            _syscall_rules_fd: i32,
         ) -> Result<(), String> {
             Ok(())
         }
     }
 
-    impl CompiledSyscallPolicy {
-        pub fn to_bpf_records(&self) -> Vec<BpfSyscallRuleRecord> {
-            vec![]
+    impl CompiledLsmPolicy {
+        pub fn exec_records(&self) -> &[BpfLsmExecRuleRecord] {
+            &self.exec_rules
+        }
+
+        pub fn match_exec_path(&self, _exec_path: &str) -> Option<PolicyHit> {
+            None
         }
 
         pub fn sync_to_bpf_maps(
             &self,
-            _syscall_rules_fd: i32,
             _policy_epoch_fd: i32,
-            _max_records: usize,
+            _lsm_exec_rules_fd: i32,
         ) -> Result<(), String> {
             Ok(())
         }
     }
 
-    #[repr(C)]
-    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
-    pub struct BpfSyscallRuleRecord {
-        pub rule_index: u32,
-        pub mode: u8,
-        pub syscall_key: u8,
-        pub reserved: u16,
-        pub syscall_id: u32,
-        pub errno_code: i32,
-        pub rule_id: [u8; 64],
-        pub syscall_name: [u8; 32],
+    pub fn clear_kprobe_bpf_maps(
+        _kprobe_exec_rules_fd: i32,
+        _syscall_rules_fd: i32,
+    ) -> Result<(), String> {
+        Ok(())
     }
 
-    impl Default for BpfSyscallRuleRecord {
-        fn default() -> Self {
-            Self {
-                rule_index: 0,
-                mode: 0,
-                syscall_key: 0,
-                reserved: 0,
-                syscall_id: 0,
-                errno_code: 0,
-                rule_id: [0; 64],
-                syscall_name: [0; 32],
-            }
-        }
+    pub fn clear_lsm_bpf_maps(_lsm_exec_rules_fd: i32) -> Result<(), String> {
+        Ok(())
+    }
+
+    pub fn compile_kprobe_policy(
+        _input: &AiAgentEnforcementInput,
+    ) -> Result<CompiledKprobePolicy, String> {
+        Ok(CompiledKprobePolicy {
+            epoch: 0,
+            exec_rules: Vec::new(),
+            syscall_rules: Vec::new(),
+        })
     }
 
-    pub fn compile_exec_rules(_rules: &[ExecRuleInput]) -> Result<CompiledExecPolicy, String> {
-        Ok(CompiledExecPolicy { epoch: 0 })
+    pub fn compile_lsm_policy(
+        _input: &AiAgentEnforcementInput,
+    ) -> Result<CompiledLsmPolicy, String> {
+        Ok(CompiledLsmPolicy {
+            epoch: 0,
+            exec_rules: Vec::new(),
+        })
     }
 
-    pub fn compile_syscall_rules(
-        _rules: &[SyscallRuleInput],
-    ) -> Result<CompiledSyscallPolicy, String> {
-        Ok(CompiledSyscallPolicy { epoch: 0 })
+    pub fn set_global_kprobe_policy(policy: Option<CompiledKprobePolicy>) {
+        *GLOBAL_KPROBE_POLICY
+            .get_or_init(|| RwLock::new(None))
+            .write()
+            .unwrap() = policy.map(Arc::new);
     }
 
-    pub fn syscall_override_symbols(_syscall_key: u8) -> &'static [&'static str] {
-        &[]
+    pub fn global_kprobe_policy() -> Option<Arc<CompiledKprobePolicy>> {
+        GLOBAL_KPROBE_POLICY
+            .get_or_init(|| RwLock::new(None))
+            .read()
+            .unwrap()
+            .clone()
     }
 
-    pub fn set_global_exec_policy(_policy: Option<CompiledExecPolicy>) {}
+    pub fn set_global_lsm_policy(policy: Option<CompiledLsmPolicy>) {
+        *GLOBAL_LSM_POLICY
+            .get_or_init(|| RwLock::new(None))
+            .write()
+            .unwrap() = policy.map(Arc::new);
+    }
 
-    pub fn global_exec_policy() -> Option<Arc<CompiledExecPolicy>> {
-        None
+    pub fn global_lsm_policy() -> Option<Arc<CompiledLsmPolicy>> {
+        GLOBAL_LSM_POLICY
+            .get_or_init(|| RwLock::new(None))
+            .read()
+            .unwrap()
+            .clone()
     }
 }