feat: add exectrack cli

Author: not-matthias

Cargo.lock

@@ -80,7 +80,7 @@ rstest_reuse = "0.7.0"
 shell-quote = "0.7.2"
 
 [workspace]
-members = ["crates/runner-shared", "crates/memtrack", "crates/codspeed-bpf"]
+members = ["crates/runner-shared", "crates/memtrack", "crates/codspeed-bpf", "crates/exectrack"]
 
 [workspace.dependencies]
 anyhow = "1.0"

Cargo.toml

@@ -40,7 +40,7 @@ pub fn build_bpf(program_name: &str, source_file: &str) {
         .expect("CARGO_CFG_TARGET_ARCH must be set in build script");
 
     let output =
-        PathBuf::from(env::var("OUT_DIR").unwrap()).join(format!("{}.skel.rs", program_name));
+        PathBuf::from(env::var("OUT_DIR").unwrap()).join(format!("{program_name}.skel.rs"));
 
     // Get the path to codspeed-bpf's C headers
     let codspeed_bpf_include = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap())
@@ -57,7 +57,7 @@ pub fn build_bpf(program_name: &str, source_file: &str) {
             &codspeed_bpf_include.to_string_lossy(),
         ])
         .build_and_generate(&output)
-        .expect(&format!("Failed to build {}.bpf.c", program_name));
+        .unwrap_or_else(|_| panic!("Failed to build {program_name}.bpf.c"));
 }
 
 /// Generate Rust bindings for a C header file

crates/codspeed-bpf/src/build.rs

@@ -0,0 +1,29 @@
+[package]
+name = "exectrack"
+version = "0.1.0"
+edition = "2024"
+repository = "https://github.com/CodSpeedHQ/runner"
+
+[lib]
+name = "exectrack"
+path = "src/lib.rs"
+
+[[bin]]
+name = "codspeed-exectrack"
+path = "src/main.rs"
+
+[dependencies]
+anyhow = { workspace = true }
+clap = { workspace = true }
+libc = { workspace = true }
+log = { workspace = true }
+env_logger = { workspace = true }
+serde_json = { workspace = true }
+serde = { workspace = true }
+static_assertions = "1.1"
+libbpf-rs = { version = "0.25.0", features = ["vendored"] }
+codspeed-bpf = { path = "../codspeed-bpf" }
+runner-shared = { path = "../runner-shared" }
+
+[build-dependencies]
+codspeed-bpf = { path = "../codspeed-bpf", features = ["build"] }

crates/exectrack/Cargo.toml

@@ -0,0 +1,4 @@
+fn main() {
+    codspeed_bpf::build::build_bpf("exectrack", "src/ebpf/c/exectrack.bpf.c");
+    codspeed_bpf::build::generate_bindings("wrapper.h");
+}

crates/exectrack/build.rs

@@ -0,0 +1,18 @@
+#ifndef __EXECTRACK_EVENT_H__
+#define __EXECTRACK_EVENT_H__
+
+#define EVENT_TYPE_FORK 1
+#define EVENT_TYPE_EXEC 2
+#define EVENT_TYPE_EXIT 3
+
+/* Event structure - shared between BPF and userspace */
+struct event {
+    uint8_t event_type; /* See EVENT_TYPE_* constants above */
+    uint64_t timestamp; /* monotonic time in nanoseconds (CLOCK_MONOTONIC) */
+    uint32_t pid;       /* Process ID */
+    uint32_t tid;       /* Thread ID */
+    uint32_t ppid;      /* Parent Process ID (for fork events) */
+    char comm[16];      /* Command name (null-terminated) */
+};
+
+#endif /* __EXECTRACK_EVENT_H__ */

crates/exectrack/src/ebpf/c/event.h

@@ -0,0 +1,90 @@
+// clang-format off
+#include "vmlinux.h"
+// clang-format on
+
+#include <bpf/bpf_core_read.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+#include "event.h"
+
+// Include shared BPF utilities from codspeed-bpf
+#include "codspeed/common.h"
+#include "codspeed/process_tracking.h"
+
+char LICENSE[] SEC("license") = "GPL";
+
+// Define standard process tracking maps using shared macro
+PROCESS_TRACKING_MAPS();
+
+// Define ring buffer for events
+BPF_RINGBUF(events, 256 * 1024);
+
+/* Helper to submit an event to the ring buffer */
+static __always_inline int submit_event(__u8 event_type, __u32 pid, __u32 ppid) {
+    if (!is_tracked(pid) && !is_tracked(ppid)) {
+        return 0;
+    }
+
+    struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
+    if (!e) {
+        return 0;
+    }
+
+    __u64 tid_full = bpf_get_current_pid_tgid();
+
+    e->timestamp = bpf_ktime_get_ns();
+    e->pid = pid;
+    e->tid = tid_full & 0xFFFFFFFF;
+    e->ppid = ppid;
+    e->event_type = event_type;
+
+    // Get current command name
+    bpf_get_current_comm(e->comm, sizeof(e->comm));
+
+    bpf_ringbuf_submit(e, 0);
+    return 0;
+}
+
+/* Track process creation via fork/clone */
+SEC("tracepoint/sched/sched_process_fork")
+int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) {
+    __u32 parent_pid = ctx->parent_pid;
+    __u32 child_pid = ctx->child_pid;
+
+    // Use shared fork handler to track child
+    if (handle_fork(parent_pid, child_pid)) {
+        // Submit fork event with parent/child relationship
+        submit_event(EVENT_TYPE_FORK, child_pid, parent_pid);
+    }
+
+    return 0;
+}
+
+/* Track process execution via execve */
+SEC("tracepoint/sched/sched_process_exec")
+int tracepoint_sched_exec(struct trace_event_raw_sched_process_exec* ctx) {
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
+
+    if (is_tracked(pid)) {
+        submit_event(EVENT_TYPE_EXEC, pid, 0);
+    }
+
+    return 0;
+}
+
+/* Track process termination via exit */
+SEC("tracepoint/sched/sched_process_exit")
+int tracepoint_sched_exit(struct trace_event_raw_sched_process_template* ctx) {
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
+
+    if (is_tracked(pid)) {
+        submit_event(EVENT_TYPE_EXIT, pid, 0);
+        // Use shared exit handler to clean up maps
+        handle_exit(pid);
+    }
+
+    return 0;
+}

crates/exectrack/src/ebpf/c/exectrack.bpf.c

@@ -0,0 +1,77 @@
+use serde::{Deserialize, Serialize};
+
+// Include the bindings for event.h
+pub mod bindings {
+    #![allow(non_upper_case_globals)]
+    #![allow(non_camel_case_types)]
+    #![allow(non_snake_case)]
+    #![allow(dead_code)]
+
+    include!(concat!(env!("OUT_DIR"), "/event.rs"));
+}
+use bindings::*;
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
+pub enum EventType {
+    Fork = EVENT_TYPE_FORK as u8,
+    Exec = EVENT_TYPE_EXEC as u8,
+    Exit = EVENT_TYPE_EXIT as u8,
+}
+
+impl From<u8> for EventType {
+    fn from(val: u8) -> Self {
+        match val as u32 {
+            bindings::EVENT_TYPE_FORK => EventType::Fork,
+            bindings::EVENT_TYPE_EXEC => EventType::Exec,
+            bindings::EVENT_TYPE_EXIT => EventType::Exit,
+            _ => panic!("Unknown event type: {val}"),
+        }
+    }
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct Event {
+    pub event_type: u8,
+    pub timestamp: u64,
+    pub pid: u32,
+    pub tid: u32,
+    pub ppid: u32,
+    pub comm: [u8; 16],
+}
+
+impl Event {
+    /// Get the event type as an enum
+    pub fn event_type(&self) -> EventType {
+        EventType::from(self.event_type)
+    }
+
+    /// Get the command name as a string
+    pub fn comm_str(&self) -> &str {
+        let len = self.comm.iter().position(|&c| c == 0).unwrap_or(16);
+        std::str::from_utf8(&self.comm[..len]).unwrap_or("<invalid>")
+    }
+}
+
+// Static assertions for C/Rust ABI safety
+mod assertions {
+    use super::*;
+    use static_assertions::{assert_eq_align, assert_eq_size, const_assert_eq};
+    use std::mem::offset_of;
+
+    assert_eq_size!(Event, bindings::event);
+    assert_eq_align!(Event, bindings::event);
+
+    const_assert_eq!(
+        offset_of!(Event, timestamp),
+        offset_of!(bindings::event, timestamp)
+    );
+    const_assert_eq!(offset_of!(Event, pid), offset_of!(bindings::event, pid));
+    const_assert_eq!(offset_of!(Event, tid), offset_of!(bindings::event, tid));
+    const_assert_eq!(
+        offset_of!(Event, event_type),
+        offset_of!(bindings::event, event_type)
+    );
+    const_assert_eq!(offset_of!(Event, comm), offset_of!(bindings::event, comm));
+}

crates/exectrack/src/ebpf/events.rs

@@ -0,0 +1,66 @@
+use anyhow::{Context, Result};
+use codspeed_bpf::ProcessTracking;
+use codspeed_bpf::{RingBufferPoller, attach_tracepoint};
+use libbpf_rs::Link;
+use libbpf_rs::skel::{OpenSkel, SkelBuilder};
+use std::mem::MaybeUninit;
+
+pub mod exectrack_skel {
+    include!(concat!(env!("OUT_DIR"), "/exectrack.skel.rs"));
+}
+pub use exectrack_skel::*;
+
+pub struct ExectrackBpf {
+    skel: Box<ExectrackSkel<'static>>,
+    probes: Vec<Link>,
+}
+
+impl ExectrackBpf {
+    pub fn new() -> Result<Self> {
+        let builder = ExectrackSkelBuilder::default();
+        let open_object = Box::leak(Box::new(MaybeUninit::uninit()));
+        let open_skel = builder
+            .open(open_object)
+            .context("Failed to open exectrack BPF skeleton")?;
+
+        let skel = Box::new(
+            open_skel
+                .load()
+                .context("Failed to load exectrack BPF skeleton")?,
+        );
+
+        Ok(Self {
+            skel,
+            probes: Vec::new(),
+        })
+    }
+
+    // Use the shared macro from codspeed-bpf for attaching tracepoints
+    attach_tracepoint!(attach_sched_fork, tracepoint_sched_fork);
+    attach_tracepoint!(attach_sched_exec, tracepoint_sched_exec);
+    attach_tracepoint!(attach_sched_exit, tracepoint_sched_exit);
+
+    pub fn attach_tracepoints(&mut self) -> Result<()> {
+        self.attach_sched_fork()?;
+        self.attach_sched_exec()?;
+        self.attach_sched_exit()?;
+        Ok(())
+    }
+
+    /// Start polling with an mpsc channel for events
+    pub fn start_polling_with_channel(
+        &self,
+        poll_timeout_ms: u64,
+    ) -> Result<(
+        RingBufferPoller,
+        std::sync::mpsc::Receiver<super::events::Event>,
+    )> {
+        RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms)
+    }
+}
+
+impl ProcessTracking for ExectrackBpf {
+    fn tracked_pids_map(&self) -> &impl libbpf_rs::MapCore {
+        &self.skel.maps.tracked_pids
+    }
+}

crates/exectrack/src/ebpf/exectrack.rs

@@ -0,0 +1,7 @@
+pub mod events;
+mod exectrack;
+mod tracker;
+
+pub use events::{Event, EventType};
+pub use exectrack::ExectrackBpf;
+pub use tracker::Tracker;