feat(agent): support ai agent syscall override enforcement

Author: kylewanginchina

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

@@ -520,6 +520,14 @@ pub mod ai_agent_enforcement {
         pub argv_contains_any: Vec<String>,
     }
 
+    #[derive(Clone, Debug, PartialEq, Eq)]
+    pub struct SyscallRuleInput {
+        pub id: String,
+        pub mode: EnforcementMode,
+        pub names: Vec<String>,
+        pub symbols: Vec<String>,
+    }
+
     #[derive(Clone, Debug, PartialEq, Eq)]
     pub struct PolicyHit {
         pub rule_index: u32,
@@ -532,6 +540,11 @@ pub mod ai_agent_enforcement {
         pub epoch: u64,
     }
 
+    #[derive(Clone, Debug, PartialEq, Eq)]
+    pub struct CompiledSyscallPolicy {
+        pub epoch: u64,
+    }
+
     impl CompiledExecPolicy {
         pub fn match_exec(&self, _exec_path: &str, _cmdline: &str) -> Option<PolicyHit> {
             None
@@ -547,10 +560,63 @@ pub mod ai_agent_enforcement {
         }
     }
 
+    impl CompiledSyscallPolicy {
+        pub fn to_bpf_records(&self) -> Vec<BpfSyscallRuleRecord> {
+            vec![]
+        }
+
+        pub fn sync_to_bpf_maps(
+            &self,
+            _syscall_rules_fd: i32,
+            _policy_epoch_fd: i32,
+            _max_records: usize,
+        ) -> 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],
+    }
+
+    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 compile_exec_rules(_rules: &[ExecRuleInput]) -> Result<CompiledExecPolicy, String> {
         Ok(CompiledExecPolicy { epoch: 0 })
     }
 
+    pub fn compile_syscall_rules(
+        _rules: &[SyscallRuleInput],
+    ) -> Result<CompiledSyscallPolicy, String> {
+        Ok(CompiledSyscallPolicy { epoch: 0 })
+    }
+
+    pub fn syscall_override_symbols(_syscall_key: u8) -> &'static [&'static str] {
+        &[]
+    }
+
     pub fn set_global_exec_policy(_policy: Option<CompiledExecPolicy>) {}
 
     pub fn global_exec_policy() -> Option<Arc<CompiledExecPolicy>> {

agent/src/common/kernel_capability.rs

@@ -11,6 +11,8 @@ pub struct KernelCapability {
     pub bpf_lsm_configured: bool,
     pub bpf_lsm_active: bool,
     pub bpf_kprobe_override_configured: bool,
+    pub bpf_kprobe_override_available: bool,
+    pub bpf_kprobe_override_symbols: Vec<String>,
     pub seccomp_filter_configured: bool,
     pub btf_vmlinux_available: bool,
 }
@@ -29,14 +31,18 @@ impl KernelCapability {
         let lsm_text = fs::read_to_string(sys_root.join("kernel/security/lsm")).unwrap_or_default();
         let config_text = read_kernel_config_from_roots(proc_root, boot_root).unwrap_or_default();
         let bpf_lsm_active = lsm_has_bpf(&lsm_text);
+        let bpf_kprobe_override_symbols = read_kprobe_override_symbols(sys_root);
+        let bpf_kprobe_override_configured =
+            config_enabled(&config_text, "CONFIG_BPF_KPROBE_OVERRIDE")
+                || !bpf_kprobe_override_symbols.is_empty();
 
         Self {
             bpf_lsm_configured: config_enabled(&config_text, "CONFIG_BPF_LSM") || bpf_lsm_active,
             bpf_lsm_active,
-            bpf_kprobe_override_configured: config_enabled(
-                &config_text,
-                "CONFIG_BPF_KPROBE_OVERRIDE",
-            ),
+            bpf_kprobe_override_configured,
+            bpf_kprobe_override_available: bpf_kprobe_override_configured
+                && !bpf_kprobe_override_symbols.is_empty(),
+            bpf_kprobe_override_symbols,
             seccomp_filter_configured: config_enabled(&config_text, "CONFIG_SECCOMP_FILTER"),
             btf_vmlinux_available: sys_root.join("kernel/btf/vmlinux").exists(),
         }
@@ -45,6 +51,14 @@ impl KernelCapability {
     pub fn supports_exec_lsm_enforcement(&self) -> bool {
         self.bpf_lsm_configured && self.bpf_lsm_active
     }
+
+    pub fn supports_kprobe_override_symbol(&self, symbol: &str) -> bool {
+        self.bpf_kprobe_override_available
+            && self
+                .bpf_kprobe_override_symbols
+                .iter()
+                .any(|allowed| allowed == symbol)
+    }
 }
 
 fn path_from_env(name: &str, default: &str) -> PathBuf {
@@ -97,6 +111,37 @@ fn read_proc_kernel_config(proc_root: &Path) -> Option<String> {
     decode_gzip(&compressed).ok()
 }
 
+fn read_kprobe_override_symbols(sys_root: &Path) -> Vec<String> {
+    const REL_PATHS: [&str; 2] = [
+        "kernel/debug/error_injection/list",
+        "kernel/debug/fail_function/injectable",
+    ];
+
+    let mut symbols = Vec::new();
+    for rel_path in REL_PATHS {
+        let Ok(text) = fs::read_to_string(sys_root.join(rel_path)) else {
+            continue;
+        };
+        symbols.extend(parse_error_injection_symbols(&text));
+    }
+    symbols.sort();
+    symbols.dedup();
+    symbols
+}
+
+fn parse_error_injection_symbols(text: &str) -> Vec<String> {
+    text.lines()
+        .filter_map(|line| {
+            let token = line.split_whitespace().next().unwrap_or_default().trim();
+            if token.is_empty() || token.starts_with('#') {
+                None
+            } else {
+                Some(token.to_string())
+            }
+        })
+        .collect()
+}
+
 fn decode_gzip(bytes: &[u8]) -> Result<String, std::io::Error> {
     let mut decoder = GzDecoder::new(Cursor::new(bytes));
     let mut output = String::new();
@@ -123,6 +168,14 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn parse_error_injection_list_takes_first_column() {
+        assert_eq!(
+            parse_error_injection_symbols("__x64_sys_reboot\tEI_ETYPE_ERRNO\n# ignored\n"),
+            vec!["__x64_sys_reboot".to_string()]
+        );
+    }
+
     #[test]
     fn support_exec_lsm_requires_config_and_active_lsm() {
         assert!(KernelCapability {
@@ -166,6 +219,31 @@ mod tests {
         let _ = fs::remove_dir_all(root);
     }
 
+    #[test]
+    fn detect_from_roots_uses_error_injection_allowlist_for_kprobe_override() {
+        let root = make_temp_root("kprobe-override-allowlist");
+        let proc_root = root.join("host-proc");
+        let sys_root = root.join("host-sys");
+        let boot_root = root.join("boot");
+        fs::create_dir_all(proc_root.join("sys/kernel")).unwrap();
+        fs::create_dir_all(sys_root.join("kernel/debug/error_injection")).unwrap();
+        fs::write(proc_root.join("sys/kernel/osrelease"), "4.18.0-test\n").unwrap();
+        fs::write(
+            sys_root.join("kernel/debug/error_injection/list"),
+            "__x64_sys_reboot\n__x64_sys_init_module\n",
+        )
+        .unwrap();
+
+        let capability = KernelCapability::detect_from_roots(&proc_root, &sys_root, &boot_root);
+
+        assert!(capability.bpf_kprobe_override_configured);
+        assert!(capability.bpf_kprobe_override_available);
+        assert!(capability.supports_kprobe_override_symbol("__x64_sys_reboot"));
+        assert!(!capability.supports_kprobe_override_symbol("__x64_sys_mount"));
+
+        let _ = fs::remove_dir_all(root);
+    }
+
     fn make_temp_root(name: &str) -> std::path::PathBuf {
         let root = std::env::temp_dir().join(format!(
             "deepflow-kernel-capability-{name}-{}",

agent/src/common/proc_event/linux.rs

@@ -952,6 +952,37 @@ mod tests {
         assert_eq!(pb.mechanism, "lsm");
     }
 
+    #[test]
+    fn test_proc_block_event_into_metric_carries_syscall_override() {
+        let raw = make_proc_block_raw(
+            2,
+            2,
+            2,
+            1,
+            1,
+            13,
+            100,
+            10,
+            1000,
+            1000,
+            169,
+            42,
+            b"block-direct-reboot",
+            b"reboot",
+            b"",
+            b"reboot",
+        );
+        let event = ProcBlockEventData::try_from(raw.as_slice()).unwrap();
+        let pb: metric::ProcBlockEventData = event.into();
+        assert_eq!(
+            pb.target_type,
+            metric::EnforcementTargetType::EnforcementTargetSyscall as i32
+        );
+        assert_eq!(pb.mechanism, "kprobe_override");
+        assert_eq!(pb.syscall_name, "reboot");
+        assert_eq!(pb.syscall_id, 169);
+    }
+
     #[test]
     fn test_new_proc_block_event_for_audit_encodes_proc_block_event() {
         let proc_event = ProcEvent {

agent/src/config/config.rs

@@ -668,6 +668,7 @@ pub struct AiAgentEnforcementRule {
     pub action: AiAgentEnforcementAction,
     pub audit: bool,
     pub exec: AiAgentExecMatch,
+    pub syscall: AiAgentSyscallMatch,
 }
 
 impl Default for AiAgentEnforcementRule {
@@ -680,6 +681,7 @@ impl Default for AiAgentEnforcementRule {
             action: AiAgentEnforcementAction::default(),
             audit: true,
             exec: AiAgentExecMatch::default(),
+            syscall: AiAgentSyscallMatch::default(),
         }
     }
 }
@@ -710,6 +712,13 @@ pub struct AiAgentExecMatch {
     pub argv_contains_any: Vec<String>,
 }
 
+#[derive(Clone, Debug, Default, Deserialize, PartialEq, Eq)]
+#[serde(default)]
+pub struct AiAgentSyscallMatch {
+    pub names: Vec<String>,
+    pub symbols: Vec<String>,
+}
+
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
 #[serde(default)]
 pub struct Proc {

agent/src/ebpf/kernel/include/bpf_base.h

@@ -41,6 +41,10 @@ struct task_struct;
 // Helper ID for bpf_get_current_task_btf (introduced in Linux 5.11).
 #define BPF_FUNC_get_current_task_btf 158
 #endif
+#ifndef BPF_FUNC_override_return
+// Helper ID for bpf_override_return (kprobe override / error injection).
+#define BPF_FUNC_override_return 58
+#endif
 
 /*
  * bpf helpers
@@ -148,6 +152,11 @@ static int
 static int
     __attribute__ ((__unused__)) (*bpf_get_stack) (void *ctx, void *buf, __u32 size,
 						     int flags) = (void *)67;
+static long
+    __attribute__ ((__unused__)) (*bpf_override_return) (struct pt_regs *regs,
+							 __u64 rc) =
+    (void *)BPF_FUNC_override_return;
+#define DF_BPF_OVERRIDE_RETURN_HELPER_DECLARED 1
 
 
 // Linux 4.14: Added support for BPF_MAP_TYPE_CPUMAP, allowing packets to be redirected to specific CPUs.

agent/src/ebpf/test/test_ai_agent_source_contracts.py

@@ -163,6 +163,10 @@ def read_source(path: Path) -> str:
     "lsm_programs_handle" in tracer_c_text,
     "tracer.c must include LSM programs in the attach lifecycle",
 )
+require(
+    "optional_kprobe_programs_handle" in tracer_c_text,
+    "tracer.c must include optional AI Agent kprobe programs in the attach lifecycle",
+)
 require(
     "new_prog->type == BPF_PROG_TYPE_LSM" in load_text
     and "Skip optional BPF LSM program" in load_text,
@@ -225,5 +229,33 @@ def read_source(path: Path) -> str:
         "ai_agent_exec_enforce.bpf.c" in support_text,
         "support_extended_observability must include ai_agent_exec_enforce.bpf.c",
     )
+    syscall_override_bpf = ENTERPRISE_BPF / "ai_agent_syscall_override.bpf.c"
+    require(
+        syscall_override_bpf.exists(),
+        f"missing enterprise AI Agent syscall override BPF: {syscall_override_bpf}",
+    )
+    syscall_override_text = read_source(syscall_override_bpf)
+    require(
+        "bpf_override_return(ctx," in syscall_override_text,
+        "AI Agent syscall enforcement must use bpf_override_return for blocking",
+    )
+    require(
+        'SEC("kprobe/__x64_sys_reboot")' in syscall_override_text,
+        "AI Agent syscall enforcement must hook direct reboot syscall with kprobe override",
+    )
+    require(
+        "df_K_ai_agent_syscall_override_" in tracer_c_text
+        or "optional kprobe: 'kprobe/__x64_sys_reboot'" in tracer_c_text,
+        "tracer.c must explicitly attach AI Agent syscall override kprobes",
+    )
+    require(
+        "ai_agent_syscall_override.bpf.c" in support_text,
+        "support_extended_observability must include ai_agent_syscall_override.bpf.c",
+    )
+    require(
+        "df_K_ai_agent_syscall_override_" in load_text
+        and "Skip optional AI Agent kprobe override program" in load_text,
+        "load.c must keep unsupported AI Agent kprobe override programs non-fatal",
+    )
 
 print("[OK]")

agent/src/ebpf/user/load.c

@@ -665,6 +665,13 @@ static enum bpf_prog_type get_prog_type(struct sec_desc *desc)
 	return prog_type;
 }
 
+static bool is_optional_ai_agent_kprobe_override_prog(struct ebpf_prog *prog)
+{
+	return prog != NULL && prog->type == BPF_PROG_TYPE_KPROBE &&
+	    prog->name != NULL &&
+	    strstr(prog->name, "df_K_ai_agent_syscall_override_") == prog->name;
+}
+
 static int load_obj__progs(struct ebpf_object *obj)
 {
 	int i;
@@ -1024,6 +1031,13 @@ static int load_obj__progs(struct ebpf_object *obj)
 				continue;
 			}
 
+			if (is_optional_ai_agent_kprobe_override_prog(new_prog)) {
+				ebpf_warning
+				    ("Skip optional AI Agent kprobe override program '%s'; syscall enforcement disabled for this hook.\n",
+				     new_prog->name);
+				continue;
+			}
+
 			if (memcmp(desc->name, "uprobe/", 7) &&
 			    memcmp(desc->name, "uretprobe/", 10)) {
 				return ETR_INVAL;