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

crates/memtrack/build.rs

@@ -12,18 +12,33 @@ fn build_ebpf() {
 
     println!("cargo:rerun-if-changed=src/ebpf/c");
 
-    // Build the BPF program
+    // The same source (main.bpf.c) is compiled into two skeletons through thin
+    // wrappers that differ only in whether MEMTRACK_BPF_VARIANT_TOKEN is set:
+    // the token variant attaches uprobes as uprobe_multi links, which a
+    // delegated BPF token can authorize, and the legacy variant uses
+    // perf_event_open for kernels predating uprobe_multi. See
+    // src/ebpf/c/utils/variant.h. The runtime picks one based on token
+    // availability.
     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_legacy.bpf.c",
+            "memtrack_legacy.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/bpf_token.rs

@@ -0,0 +1,9 @@
+/// Whether libbpf has a delegated BPF token to attach to its `bpf()` calls. It
+/// reads one from the bpffs mount named by `LIBBPF_BPF_TOKEN_PATH`, which is how
+/// eBPF gets loaded without host privileges.
+///
+/// <https://docs.ebpf.io/linux/concepts/token/>
+pub fn has_delegated_bpf_token() -> bool {
+    std::env::var_os("LIBBPF_BPF_TOKEN_PATH")
+        .is_some_and(|p| !p.is_empty() && std::path::Path::new(&p).is_dir())
+}

crates/memtrack/src/ebpf/attach_worker.rs

@@ -232,8 +232,38 @@ impl Worker {
             ),
         };
 
-        let Ok(lib) = AllocatorLib::from_path_static(&mapping.attach_path) else {
-            debug!("on-demand: not an allocator: {}", mapping.display);
+        let candidates = mapping.candidate_paths();
+        let mut lib = None;
+        let mut unreadable = Vec::new();
+        for path in &candidates {
+            match AllocatorLib::from_path_static(path) {
+                Ok(found) => {
+                    lib = Some(found);
+                    break;
+                }
+                // Classification reads the whole file, so a failed open is
+                // indistinguishable from a readable non-allocator. Probe
+                // openability to tell them apart.
+                Err(_) => {
+                    if let Err(e) = std::fs::File::open(path) {
+                        unreadable.push(format!("{}: {e}", path.display()));
+                    }
+                }
+            }
+        }
+
+        let Some(lib) = lib else {
+            // A mapping no path can reach costs coverage of that one library,
+            // which is not worth killing the run over.
+            if unreadable.len() == candidates.len() {
+                warn!(
+                    "no readable path for mapping {} in pid {}, allocator coverage may be incomplete",
+                    mapping.display, req.pid
+                );
+                debug!("unreadable paths: {}", unreadable.join(", "));
+            } else {
+                debug!("on-demand: not an allocator: {}", mapping.display);
+            }
             return Ok(true);
         };
 

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,12 @@
         __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;                                                \
+        struct task_ids ids = current_task_ids();                             \
+        __u64 tid = ids.tid;                                                  \
                                                                               \
-        if (!is_tracked(pid)) {                                               \
+        if (!is_tracked(ids.tgid)) {                                          \
             return 0;                                                         \
         }                                                                     \
                                                                               \
@@ -53,9 +53,9 @@
         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();                               \
+        __u64 tid = current_tid();                                            \
         struct name##_args_t* args = bpf_map_lookup_elem(&name##_args, &tid); \
                                                                               \
         if (!args) {                                                          \
@@ -105,11 +105,11 @@ struct posix_memalign_args_t {
 };
 BPF_HASH_MAP(posix_memalign_args, __u64, struct posix_memalign_args_t, 10000);
 
-SEC("uprobe")
+SEC(UPROBE_SEC)
 int uprobe_posix_memalign(struct pt_regs* ctx) {
-    __u64 tid = bpf_get_current_pid_tgid();
-    __u32 pid = tid >> 32;
-    if (!is_tracked(pid)) {
+    struct task_ids ids = current_task_ids();
+    __u64 tid = ids.tid;
+    if (!is_tracked(ids.tgid)) {
         return 0;
     }
 
@@ -118,9 +118,9 @@ int uprobe_posix_memalign(struct pt_regs* ctx) {
     return 0;
 }
 
-SEC("uretprobe")
+SEC(URETPROBE_SEC)
 int uretprobe_posix_memalign(struct pt_regs* ctx) {
-    __u64 tid = bpf_get_current_pid_tgid();
+    __u64 tid = current_tid();
     struct posix_memalign_args_t* args = bpf_map_lookup_elem(&posix_memalign_args, &tid);
     if (!args) {
         return 0;
@@ -149,9 +149,9 @@ struct mmap_args {
 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)) {
+    struct task_ids ids = current_task_ids();
+    __u64 tid = ids.tid;
+    if (is_tracked(ids.tgid)) {
         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_tgid();
     if (!is_tracked(tgid)) {
         return 0;
     }

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

@@ -0,0 +1,5 @@
+/* Legacy variant: uprobes attach through perf_event_open(), which works on
+ * kernels predating uprobe_multi (< 6.6) but needs CAP_PERFMON in the init user
+ * namespace and so cannot be delegated via a BPF token. See utils/variant.h.
+ * Program bodies live in main.bpf.c. */
+#include "main.bpf.c"

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

@@ -0,0 +1,7 @@
+/* Token variant: uprobes attach as uprobe_multi links, so every attach goes
+ * through bpf() and a delegated BPF token can authorize it without host
+ * privileges. Requires kernel >= 6.6. See utils/variant.h for why classic
+ * uprobes cannot be delegated. Program bodies live in main.bpf.c; only the
+ * SEC() annotations differ, keyed on this define. */
+#define MEMTRACK_BPF_VARIANT_TOKEN 1
+#include "main.bpf.c"

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

@@ -21,16 +21,18 @@ static __always_inline long wake_flags(void) {
 }
 
 static __always_inline int store_param(void* map, __u64 value) {
-    __u64 tid = bpf_get_current_pid_tgid();
-    __u32 pid = tid >> 32;
-    if (is_tracked(pid)) {
+    /* Key by the tid: unique per thread, so it survives the entry/exit pair even
+     * when several threads are inside the same allocator call. */
+    struct task_ids ids = current_task_ids();
+    __u64 tid = ids.tid;
+    if (is_tracked(ids.tgid)) {
         bpf_map_update_elem(map, &tid, &value, BPF_ANY);
     }
     return 0;
 }
 
 static __always_inline __u64* take_param(void* map) {
-    __u64 tid = bpf_get_current_pid_tgid();
+    __u64 tid = current_tid();
     __u64* value = bpf_map_lookup_elem(map, &tid);
     if (value) {
         bpf_map_delete_elem(map, &tid);
@@ -40,10 +42,9 @@ 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;                                          \
+        struct task_ids ids = current_task_ids();                       \
                                                                         \
-        if (!is_tracked(pid) || !is_enabled()) {                        \
+        if (!is_tracked(ids.tgid) || !is_enabled()) {                   \
             return 0;                                                   \
         }                                                               \
                                                                         \
@@ -58,8 +59,8 @@ static __always_inline __u64* take_param(void* map) {
         }                                                               \
                                                                         \
         e->header.timestamp = bpf_ktime_get_ns();                       \
-        e->header.pid = pid;                                            \
-        e->header.tid = tid & 0xFFFFFFFF;                               \
+        e->header.pid = ids.tgid;                                       \
+        e->header.tid = ids.tid;                                        \
         e->header.event_type = evt_type;                                \
                                                                         \
         fill_data;                                                      \

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

@@ -2,6 +2,7 @@
 #define __PROCESS_TRACKING_H__
 
 #include "map_helpers.h"
+#include "variant.h"
 
 BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000);
 BPF_HASH_MAP(pids_ppid, __u32, __u32, 10000);
@@ -43,15 +44,26 @@ 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);
     }
+}
 
+/* tp_btf rather than a classic tracepoint on two counts: it attaches with
+ * BPF_RAW_TRACEPOINT_OPEN, so a token can delegate it (perf_event_open() cannot
+ * be), and its arguments are task_struct pointers rather than a flattened trace
+ * event, so PIDs resolve through the tracker's namespace.
+ *
+ * https://docs.ebpf.io/linux/program-type/BPF_PROG_TYPE_TRACING/
+ */
+SEC("tp_btf/sched_process_fork")
+int BPF_PROG(tracepoint_sched_fork, struct task_struct* parent, struct task_struct* child) {
+    follow_fork(task_ns_tgid(parent), task_ns_tgid(child));
     return 0;
 }