feat: run CPU isolation and the memory instrument inside the macro-agent sandbox

.github/workflows/ci.yml

@@ -95,7 +95,15 @@ jobs:
         # Each memtrack integration test binary runs its cases serially
         # (eBPF tracker can't overlap with itself in one process), so we
         # shard at the test-binary level to parallelize across jobs.
-        test: [c_tests, cpp_tests, rust_tests, spawn_tests, dlopen_tests]
+        test:
+          [
+            c_tests,
+            cpp_tests,
+            rust_tests,
+            spawn_tests,
+            dlopen_tests,
+            flavor_equivalence_tests,
+          ]
     steps:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
         with:

crates/memtrack/build.rs

@@ -12,18 +12,29 @@ fn build_ebpf() {
 
     println!("cargo:rerun-if-changed=src/ebpf/c");
 
-    // Build the BPF program
+    // Build the BPF program. The same shared source (main.bpf.c) is compiled
+    // into two skeletons via thin wrappers: memtrack_token.bpf.c defines
+    // MEMTRACK_BPF_LINKS to attach through bpf() links (uprobe_multi/tp_btf) so a
+    // delegated token authorizes them, while memtrack_perf.bpf.c uses the classic
+    // perf-based paths for kernels without uprobe_multi. The runtime picks the
+    // matching skeleton based on whether a BPF token is available.
     let arch = env::var("CARGO_CFG_TARGET_ARCH")
         .expect("CARGO_CFG_TARGET_ARCH must be set in build script");
-    let memtrack_out = PathBuf::from(env::var("OUT_DIR").unwrap()).join("memtrack.skel.rs");
-    SkeletonBuilder::new()
-        .source("src/ebpf/c/main.bpf.c")
-        .clang_args([
-            "-I",
-            &vmlinux::include_path_root().join(arch).to_string_lossy(),
-        ])
-        .build_and_generate(&memtrack_out)
-        .unwrap();
+    let vmlinux_inc = vmlinux::include_path_root()
+        .join(arch)
+        .to_string_lossy()
+        .into_owned();
+    let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
+    for (source, skel) in [
+        ("src/ebpf/c/memtrack_token.bpf.c", "memtrack_token.skel.rs"),
+        ("src/ebpf/c/memtrack_perf.bpf.c", "memtrack_perf.skel.rs"),
+    ] {
+        SkeletonBuilder::new()
+            .source(source)
+            .clang_args(["-I", &vmlinux_inc])
+            .build_and_generate(out_dir.join(skel))
+            .unwrap();
+    }
 
     // Generate bindings for event.h
     let bindings = bindgen::Builder::default()

crates/memtrack/src/ebpf/c/allocator.h

@@ -7,9 +7,9 @@
 
 #define UPROBE_ARG_RET(name, arg_expr, submit_block)                                      \
     BPF_HASH_MAP(name##_arg, __u64, __u64, 10000);                                        \
-    SEC("uprobe")                                                                         \
+    SEC(UPROBE_SEC)                                                                       \
     int uprobe_##name(struct pt_regs* ctx) { return store_param(&name##_arg, arg_expr); } \
-    SEC("uretprobe")                                                                      \
+    SEC(URETPROBE_SEC)                                                                    \
     int uretprobe_##name(struct pt_regs* ctx) {                                           \
         __u64* arg_ptr = take_param(&name##_arg);                                         \
         if (!arg_ptr) {                                                                   \
@@ -24,7 +24,7 @@
     }
 
 #define UPROBE_RET(name, arg_expr, submit_block) \
-    SEC("uprobe")                                \
+    SEC(UPROBE_SEC)                              \
     int uprobe_##name(struct pt_regs* ctx) {     \
         __u64 arg0 = arg_expr;                   \
         if (arg0 == 0) {                         \
@@ -39,12 +39,11 @@
         __u64 arg1;                                                           \
     };                                                                        \
     BPF_HASH_MAP(name##_args, __u64, struct name##_args_t, 10000);            \
-    SEC("uprobe")                                                             \
+    SEC(UPROBE_SEC)                                                           \
     int uprobe_##name(struct pt_regs* ctx) {                                  \
         __u64 tid = bpf_get_current_pid_tgid();                               \
-        __u32 pid = tid >> 32;                                                \
                                                                               \
-        if (!is_tracked(pid)) {                                               \
+        if (!is_tracked(current_pid())) {                                     \
             return 0;                                                         \
         }                                                                     \
                                                                               \
@@ -53,7 +52,7 @@
         bpf_map_update_elem(&name##_args, &tid, &args, BPF_ANY);              \
         return 0;                                                             \
     }                                                                         \
-    SEC("uretprobe")                                                          \
+    SEC(URETPROBE_SEC)                                                        \
     int uretprobe_##name(struct pt_regs* ctx) {                               \
         __u64 tid = bpf_get_current_pid_tgid();                               \
         struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \
@@ -99,8 +98,7 @@ BPF_HASH_MAP(mmap_temp, __u64, struct mmap_args, 10000);
 
 static __always_inline void store_mmap_args(__u64 addr, __u64 len) {
     __u64 tid = bpf_get_current_pid_tgid();
-    __u32 pid = tid >> 32;
-    if (is_tracked(pid)) {
+    if (is_tracked(current_pid())) {
         struct mmap_args args = {.addr = addr, .len = len};
         bpf_map_update_elem(&mmap_temp, &tid, &args, BPF_ANY);
     }

crates/memtrack/src/ebpf/c/attach.h

@@ -31,8 +31,7 @@ int BPF_PROG(watch_exec_mmap, struct file* file, unsigned long prot, unsigned lo
         return 0;
     }
 
-    __u64 pid_tgid = bpf_get_current_pid_tgid();
-    __u32 tgid = pid_tgid >> 32;
+    __u32 tgid = current_pid();
     if (!is_tracked(tgid)) {
         return 0;
     }

crates/memtrack/src/ebpf/c/memtrack_perf.bpf.c

@@ -0,0 +1,5 @@
+/* Classic perf attach variant: perf-based uprobe/uretprobe + a perf
+ * sched_process_fork tracepoint. Works on kernels predating uprobe_multi but
+ * needs CAP_PERFMON in the init user namespace, so it cannot be delegated into
+ * the sandbox. The program bodies live in the shared main.bpf.c. */
+#include "main.bpf.c"

crates/memtrack/src/ebpf/c/memtrack_token.bpf.c

@@ -0,0 +1,6 @@
+/* BPF-token attach variant: uprobe_multi links + tp_btf fork hook, authorized
+ * through bpf() so a delegated token can load them inside the sandbox. Requires
+ * a kernel with uprobe_multi (>= 6.6). The program bodies live in the shared
+ * main.bpf.c; only the SEC() annotations differ, keyed on this define. */
+#define MEMTRACK_BPF_LINKS 1
+#include "main.bpf.c"

crates/memtrack/src/ebpf/c/utils/event_helpers.h

@@ -21,9 +21,10 @@ static __always_inline long wake_flags(void) {
 }
 
 static __always_inline int store_param(void* map, __u64 value) {
+    /* Key by the global tid (stable, unique per thread across the entry/exit
+     * pair); gate on the namespace-relative PID that the tracker registered. */
     __u64 tid = bpf_get_current_pid_tgid();
-    __u32 pid = tid >> 32;
-    if (is_tracked(pid)) {
+    if (is_tracked(current_pid())) {
         bpf_map_update_elem(map, &tid, &value, BPF_ANY);
     }
     return 0;
@@ -41,7 +42,7 @@ static __always_inline __u64* take_param(void* map) {
 #define SUBMIT_EVENT(evt_type, fill_data)                               \
     {                                                                   \
         __u64 tid = bpf_get_current_pid_tgid();                         \
-        __u32 pid = tid >> 32;                                          \
+        __u32 pid = current_pid();                                      \
                                                                         \
         if (!is_tracked(pid) || !is_enabled()) {                        \
             return 0;                                                   \

crates/memtrack/src/ebpf/c/utils/flavor.h

@@ -0,0 +1,61 @@
+#ifndef __FLAVOR_H__
+#define __FLAVOR_H__
+
+/* Attach mechanism, selected at build time. With a delegated BPF token
+ * (MEMTRACK_BPF_LINKS) the uprobes attach as uprobe_multi links and the fork
+ * hook as tp_btf, both authorized through bpf(). Without it they use the classic
+ * perf-based uprobe/tracepoint paths, which work on kernels predating
+ * uprobe_multi (< 6.6) but cannot be delegated into the sandbox. */
+#ifdef MEMTRACK_BPF_LINKS
+#define UPROBE_SEC "uprobe.multi"
+#define URETPROBE_SEC "uretprobe.multi"
+#else
+#define UPROBE_SEC "uprobe"
+#define URETPROBE_SEC "uretprobe"
+#endif
+
+/* PID namespace the userspace tracker resolves PIDs in. When set (ino != 0),
+ * PIDs are read relative to this namespace so tracking works when the tracker
+ * runs inside a PID namespace (e.g. the macro-agent sandbox) while eBPF sees
+ * global PIDs. Zero means "use global PIDs" — the default outside a sandbox. */
+const volatile __u64 target_pidns_dev = 0;
+const volatile __u64 target_pidns_ino = 0;
+
+/* Current task's PID in the configured namespace (or global when unset). */
+static __always_inline __u32 current_pid(void) {
+    if (target_pidns_ino == 0) {
+        return bpf_get_current_pid_tgid() >> 32;
+    }
+    struct bpf_pidns_info nsinfo = {};
+    if (bpf_get_ns_current_pid_tgid(target_pidns_dev, target_pidns_ino, &nsinfo, sizeof(nsinfo)) !=
+        0) {
+        return 0;
+    }
+    return nsinfo.tgid;
+}
+
+/* A task's PID as seen in the configured namespace (or global when unset).
+ * Walks the task's pid struct: numbers[level].nr holds the PID at that namespace
+ * depth, and numbers[level].ns identifies which namespace it is. */
+static __always_inline __u32 task_ns_pid(struct task_struct* task) {
+    if (target_pidns_ino == 0) {
+        return BPF_CORE_READ(task, pid);
+    }
+    struct pid* thread_pid = BPF_CORE_READ(task, thread_pid);
+    if (!thread_pid) {
+        return 0;
+    }
+    unsigned int level = BPF_CORE_READ(thread_pid, level);
+    /* Bound the level to satisfy the verifier and match the pid array. */
+    if (level >= 4) {
+        return 0;
+    }
+    struct upid* up = &thread_pid->numbers[level];
+    __u64 ino = BPF_CORE_READ(up, ns, ns.inum);
+    if (ino != target_pidns_ino) {
+        return 0;
+    }
+    return BPF_CORE_READ(up, nr);
+}
+
+#endif /* __FLAVOR_H__ */

crates/memtrack/src/ebpf/c/utils/process_tracking.h

@@ -1,6 +1,7 @@
 #ifndef __PROCESS_TRACKING_H__
 #define __PROCESS_TRACKING_H__
 
+#include "flavor.h"
 #include "map_helpers.h"
 
 BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000);
@@ -43,16 +44,34 @@ static __always_inline void track_child(__u32 child_pid, __u32 parent_pid) {
     bpf_map_update_elem(&pids_ppid, &child_pid, &parent_pid, BPF_ANY);
 }
 
-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;
-
+/* Record a parent→child fork so the child inherits the parent's tracked state. */
+static __always_inline void follow_fork(__u32 parent_pid, __u32 child_pid) {
+    if (parent_pid == 0 || child_pid == 0) {
+        return;
+    }
     if (is_tracked(parent_pid)) {
         track_child(child_pid, parent_pid);
     }
+}
 
+#ifdef MEMTRACK_BPF_LINKS
+/* tp_btf attaches via bpf(), so a delegated BPF token can authorize it inside
+ * the sandbox. Reads task_struct pointers, resolving PIDs in the tracker's
+ * namespace. */
+SEC("tp_btf/sched_process_fork")
+int BPF_PROG(tracepoint_sched_fork, struct task_struct* parent, struct task_struct* child) {
+    follow_fork(task_ns_pid(parent), task_ns_pid(child));
+    return 0;
+}
+#else
+/* Classic perf tracepoint for kernels without (or runs without) a BPF token.
+ * The raw tracepoint context carries global PIDs directly; without a token we
+ * never run in a PID namespace, so global PIDs are correct. */
+SEC("tracepoint/sched/sched_process_fork")
+int tracepoint_sched_fork(struct trace_event_raw_sched_process_fork* ctx) {
+    follow_fork(ctx->parent_pid, ctx->child_pid);
     return 0;
 }
+#endif
 
 #endif /* __PROCESS_TRACKING_H__ */

crates/memtrack/src/ebpf/memtrack/allocator.rs

@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use libbpf_rs::UprobeOpts;
+use libbpf_rs::{UprobeMultiOpts, UprobeOpts};
 use paste::paste;
 use std::path::Path;