fix(agent): retry block event delivery and backfill best effort

Author: kylewanginchina

agent/src/common/proc_event/linux.rs

@@ -319,11 +319,21 @@ pub struct ProcLifecycleExecInfo<'a> {
     pub parent_pid: u32,
     pub uid: u32,
     pub gid: u32,
+    pub timestamp: u64,
     pub comm: &'a [u8],
     pub cmdline: &'a [u8],
     pub exec_path: &'a [u8],
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct ProcBlockInfo<'a> {
+    pub pid: u32,
+    pub action: u8,
+    pub timestamp: u64,
+    pub rule_id: &'a str,
+    pub exec_path: &'a [u8],
+}
+
 impl TryFrom<&[u8]> for ProcLifecycleEventData {
     type Error = Error;
 
@@ -406,7 +416,9 @@ const PROC_BLOCK_CMDLINE_LEN: usize = 256;
 const PROC_BLOCK_EXEC_PATH_LEN: usize = 256;
 const ENFORCEMENT_TARGET_EXEC: u8 = 1;
 const ENFORCEMENT_ACTION_AUDIT: u8 = 1;
+const ENFORCEMENT_ACTION_DENY: u8 = 2;
 const ENFORCEMENT_MECHANISM_USER_SPACE_AUDIT: u8 = 5;
+const ENFORCEMENT_GUARANTEE_BEST_EFFORT: u8 = 2;
 const ENFORCEMENT_GUARANTEE_AUDIT_ONLY: u8 = 3;
 
 struct ProcBlockEventData {
@@ -725,6 +737,7 @@ impl ProcEvent {
                     parent_pid: data.parent_pid,
                     uid: data.uid,
                     gid: data.gid,
+                    timestamp: data.timestamp,
                     comm: &data.comm,
                     cmdline: &data.cmdline,
                     exec_path: &data.exec_path,
@@ -734,6 +747,19 @@ impl ProcEvent {
         }
     }
 
+    pub fn proc_block_info(&self) -> Option<ProcBlockInfo<'_>> {
+        match &self.event_data {
+            EventData::ProcBlockEvent(data) => Some(ProcBlockInfo {
+                pid: data.pid,
+                action: data.action,
+                timestamp: data.timestamp,
+                rule_id: data.rule_id.as_str(),
+                exec_path: &data.exec_path,
+            }),
+            _ => None,
+        }
+    }
+
     pub fn new_proc_block_event_for_audit(
         &self,
         rule_id: &str,
@@ -786,6 +812,59 @@ impl ProcEvent {
             event_data: EventData::ProcBlockEvent(block_event),
         })))
     }
+
+    pub fn new_proc_block_event_for_best_effort(
+        &self,
+        rule_id: &str,
+        policy_epoch: u64,
+    ) -> Option<BoxedProcEvents> {
+        let data = match &self.event_data {
+            EventData::ProcLifecycleEvent(data) if data.lifecycle_type == PROC_LIFECYCLE_EXEC => {
+                data
+            }
+            _ => return None,
+        };
+        let process_kname = if data.comm.is_empty() {
+            self.process_kname.clone()
+        } else {
+            data.comm.clone()
+        };
+        let root_pid = self.ai_agent_root_pid;
+        let block_event = ProcBlockEventData {
+            rule_id: rule_id.chars().take(PROC_BLOCK_RULE_ID_LEN).collect(),
+            target_type: ENFORCEMENT_TARGET_EXEC,
+            action: ENFORCEMENT_ACTION_DENY,
+            mechanism: enforcement_mechanism_name(ENFORCEMENT_MECHANISM_USER_SPACE_AUDIT)
+                .to_string(),
+            guarantee: enforcement_guarantee_name(ENFORCEMENT_GUARANTEE_BEST_EFFORT).to_string(),
+            errno: 1,
+            pid: data.pid,
+            parent_pid: data.parent_pid,
+            ai_agent_root_pid: root_pid,
+            uid: data.uid,
+            gid: data.gid,
+            comm: process_kname.clone(),
+            cmdline: truncate_bytes(&data.cmdline, PROC_BLOCK_CMDLINE_LEN),
+            exec_path: truncate_bytes(&data.exec_path, PROC_BLOCK_EXEC_PATH_LEN),
+            syscall_name: String::new(),
+            syscall_id: 0,
+            timestamp: data.timestamp,
+            policy_epoch,
+        };
+
+        Some(BoxedProcEvents(Box::new(ProcEvent {
+            pid: data.pid,
+            pod_id: self.pod_id,
+            ai_agent_root_pid: root_pid,
+            thread_id: self.thread_id,
+            coroutine_id: self.coroutine_id,
+            process_kname,
+            start_time: self.start_time,
+            end_time: self.end_time,
+            event_type: EventType::ProcBlockEvent,
+            event_data: EventData::ProcBlockEvent(block_event),
+        })))
+    }
 }
 
 fn truncate_bytes(bytes: &[u8], limit: usize) -> Vec<u8> {
@@ -1025,6 +1104,49 @@ mod tests {
         assert_eq!(block.exec_path, b"/sbin/reboot");
     }
 
+    #[test]
+    fn test_new_proc_block_event_for_best_effort_encodes_proc_block_event() {
+        let proc_event = ProcEvent {
+            pid: 13,
+            pod_id: 7,
+            ai_agent_root_pid: 100,
+            thread_id: 13,
+            coroutine_id: 0,
+            process_kname: b"reboot".to_vec(),
+            start_time: 42,
+            end_time: 43,
+            event_type: EventType::ProcLifecycleEvent,
+            event_data: EventData::ProcLifecycleEvent(ProcLifecycleEventData {
+                lifecycle_type: PROC_LIFECYCLE_EXEC,
+                pid: 13,
+                parent_pid: 10,
+                uid: 1000,
+                gid: 1000,
+                timestamp: 42,
+                comm: b"reboot".to_vec(),
+                cmdline: b"reboot now".to_vec(),
+                exec_path: b"/sbin/reboot".to_vec(),
+            }),
+        };
+
+        let boxed = proc_event
+            .new_proc_block_event_for_best_effort("block-reboot", 99)
+            .unwrap();
+        let mut buf = Vec::new();
+        boxed.encode(&mut buf).unwrap();
+        let pb = metric::ProcEvent::decode(buf.as_slice()).unwrap();
+        let block = pb.proc_block_event_data.unwrap();
+
+        assert_eq!(pb.event_type, metric::EventType::ProcBlockEvent as i32);
+        assert_eq!(block.rule_id, "block-reboot");
+        assert_eq!(block.action, metric::EnforcementAction::Deny as i32);
+        assert_eq!(block.mechanism, "user_space_audit");
+        assert_eq!(block.guarantee, "best_effort");
+        assert_eq!(block.errno, 1);
+        assert_eq!(block.ai_agent_root_pid, 100);
+        assert_eq!(block.exec_path, b"/sbin/reboot");
+    }
+
     fn make_proc_block_raw(
         target_type: u8,
         action: u8,

agent/src/ebpf/test/test_ai_agent_source_contracts.py

@@ -405,6 +405,20 @@ def read_source(path: Path) -> str:
         "buf->path,\n\t\t\t\t       AI_AGENT_EXEC_PATTERN_LEN" in exec_override_text,
         "AI Agent exec override must report exec_path as cmdline placeholder instead of a partial argv slot",
     )
+    file_io_bpf = ENTERPRISE_BPF / "ai_agent_file_io.bpf.c"
+    file_io_text = read_source(file_io_bpf)
+    require(
+        "for (__u32 attempt = 0; attempt < 3; attempt++)" in exec_override_standalone_text
+        and "ret = bpf_perf_event_output(" in exec_override_standalone_text
+        and "if (ret >= 0)" in exec_override_standalone_text,
+        "standalone exec override helper must retry perf event output up to 3 attempts",
+    )
+    require(
+        "for (__u32 attempt = 0; attempt < 3; attempt++)" in file_io_text
+        and "ret = bpf_perf_event_output(" in file_io_text
+        and "if (ret >= 0)" in file_io_text,
+        "shared AI Agent event helper must retry perf event output up to 3 attempts",
+    )
     syscall_override_bpf = ENTERPRISE_BPF / "ai_agent_syscall_override.bpf.c"
     require(
         syscall_override_bpf.exists(),

agent/src/ebpf_dispatcher.rs

@@ -40,13 +40,17 @@ fn main() {
 #[cfg(feature = "extended_observability")]
 pub mod memory_profile;
 
+#[cfg(feature = "enterprise")]
+use std::collections::HashMap;
 use std::ffi::{CStr, CString};
 use std::ptr::{self, null_mut};
 use std::slice;
 #[cfg(feature = "enterprise")]
 use std::sync::atomic::AtomicI32;
 use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
 use std::sync::Arc;
+#[cfg(feature = "enterprise")]
+use std::sync::{Mutex, OnceLock};
 use std::thread::{self, JoinHandle};
 use std::time::Duration;
 
@@ -241,6 +245,98 @@ fn emit_ai_agent_enforcement_audit_event(event: &BoxedProcEvents) {
     }
 }
 
+#[cfg(feature = "enterprise")]
+fn kernel_block_event_cache() -> &'static Mutex<HashMap<KernelBlockMarkerKey, u64>> {
+    RECENT_KERNEL_BLOCK_EVENTS.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+#[cfg(feature = "enterprise")]
+fn prune_kernel_block_event_cache(cache: &mut HashMap<KernelBlockMarkerKey, u64>, now: u64) {
+    cache.retain(|_, ts| now.saturating_sub(*ts) <= KERNEL_BLOCK_EVENT_CACHE_WINDOW_NS);
+}
+
+#[cfg(feature = "enterprise")]
+fn record_kernel_block_event(event: &BoxedProcEvents) {
+    let Some(info) = event.0.proc_block_info() else {
+        return;
+    };
+    if info.action != metric::EnforcementAction::Deny as u8 || info.exec_path.is_empty() {
+        return;
+    }
+    let mut cache = kernel_block_event_cache().lock().unwrap();
+    prune_kernel_block_event_cache(&mut cache, info.timestamp);
+    cache.insert(
+        KernelBlockMarkerKey {
+            pid: info.pid,
+            rule_id: info.rule_id.to_string(),
+            exec_path: info.exec_path.to_vec(),
+        },
+        info.timestamp,
+    );
+}
+
+#[cfg(feature = "enterprise")]
+fn consume_recent_kernel_block_event(pid: u32, rule_id: &str, exec_path: &[u8], now: u64) -> bool {
+    let mut cache = kernel_block_event_cache().lock().unwrap();
+    prune_kernel_block_event_cache(&mut cache, now);
+    cache
+        .remove(&KernelBlockMarkerKey {
+            pid,
+            rule_id: rule_id.to_string(),
+            exec_path: exec_path.to_vec(),
+        })
+        .map(|ts| now.saturating_sub(ts) <= KERNEL_BLOCK_EVENT_CACHE_WINDOW_NS)
+        .unwrap_or(false)
+}
+
+#[cfg(feature = "enterprise")]
+#[allow(static_mut_refs)]
+fn emit_ai_agent_enforcement_best_effort_event(event: &BoxedProcEvents) {
+    use enterprise_utils::ai_agent_enforcement::EnforcementMode;
+
+    if event.0.ai_agent_root_pid == 0 {
+        return;
+    }
+    let Some(exec_info) = event.0.proc_lifecycle_exec_info() else {
+        return;
+    };
+    if exec_info.exec_path.is_empty() {
+        return;
+    }
+    let Some(policy) = enterprise_utils::ai_agent_enforcement::global_exec_policy() else {
+        return;
+    };
+    let exec_path = String::from_utf8_lossy(exec_info.exec_path);
+    let cmdline = String::from_utf8_lossy(exec_info.cmdline);
+    let Some(hit) = policy.match_exec(&exec_path, &cmdline) else {
+        return;
+    };
+    if hit.mode != EnforcementMode::Block {
+        return;
+    }
+    if consume_recent_kernel_block_event(
+        exec_info.pid,
+        &hit.rule_id,
+        exec_info.exec_path,
+        exec_info.timestamp,
+    ) {
+        return;
+    }
+    let Some(best_effort_event) = event
+        .0
+        .new_proc_block_event_for_best_effort(&hit.rule_id, policy.epoch)
+    else {
+        return;
+    };
+    unsafe {
+        if let Some(sender) = PROC_EVENT_SENDER.as_mut() {
+            if let Err(e) = sender.send(best_effort_event) {
+                warn!("ai agent enforcement best_effort event send error: {:?}", e);
+            }
+        }
+    }
+}
+
 impl OwnedCountable for SyncEbpfCounter {
     fn get_counters(&self) -> Vec<Counter> {
         let rx = self.counter.rx.swap(0, Ordering::Relaxed);
@@ -705,9 +801,22 @@ static AI_AGENT_EXEC_LSM_EVENTS_ACTIVE: AtomicBool = AtomicBool::new(false);
 #[cfg(feature = "enterprise")]
 static AI_AGENT_EXEC_KPROBE_EVENTS_ACTIVE: AtomicBool = AtomicBool::new(false);
 #[cfg(feature = "enterprise")]
+static RECENT_KERNEL_BLOCK_EVENTS: OnceLock<Mutex<HashMap<KernelBlockMarkerKey, u64>>> =
+    OnceLock::new();
+#[cfg(feature = "enterprise")]
 const AI_AGENT_EXEC_RULES_BPF_MAX: usize = 256;
 #[cfg(feature = "enterprise")]
 const AI_AGENT_SYSCALL_RULES_BPF_MAX: usize = 32;
+#[cfg(feature = "enterprise")]
+const KERNEL_BLOCK_EVENT_CACHE_WINDOW_NS: u64 = 5_000_000_000;
+
+#[cfg(feature = "enterprise")]
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+struct KernelBlockMarkerKey {
+    pid: u32,
+    rule_id: String,
+    exec_path: Vec<u8>,
+}
 
 #[cfg(feature = "enterprise")]
 fn ai_agent_enforcement_mode_eq(value: &str, expected: &str) -> bool {
@@ -922,7 +1031,11 @@ impl EbpfCollector {
                 #[cfg(feature = "enterprise")]
                 fill_ai_agent_root_pid(&mut event);
                 #[cfg(feature = "enterprise")]
+                record_kernel_block_event(&event);
+                #[cfg(feature = "enterprise")]
                 emit_ai_agent_enforcement_audit_event(&event);
+                #[cfg(feature = "enterprise")]
+                emit_ai_agent_enforcement_best_effort_event(&event);
                 if let Err(e) = PROC_EVENT_SENDER.as_mut().unwrap().send(event) {
                     warn!("event send ebpf error: {:?}", e);
                 }