From 81c4b15016c4f8f93fc47a0dd9765c01e5281ae4 Mon Sep 17 00:00:00 2001 From: Guillaume Lagrange Date: Wed, 8 Jul 2026 15:44:19 +0200 Subject: [PATCH] feat(memtrack): allow memtrack usage with a bpf token --- .github/workflows/ci.yml | 9 +- crates/memtrack/build.rs | 31 +- crates/memtrack/src/ebpf/c/memtrack.bpf.c | 109 +++++- .../memtrack/src/ebpf/c/memtrack_perf.bpf.c | 5 + .../memtrack/src/ebpf/c/memtrack_token.bpf.c | 6 + crates/memtrack/src/ebpf/memtrack.rs | 368 +++++++++++------- crates/memtrack/src/ebpf/mod.rs | 2 +- crates/memtrack/src/ebpf/tracker.rs | 18 +- crates/memtrack/testdata/flavor_equivalence.c | 50 +++ .../tests/flavor_equivalence_tests.rs | 123 ++++++ crates/memtrack/tests/shared.rs | 77 ++-- ...nce_tests__flavor_equivalence_summary.snap | 26 ++ src/executor/memory/executor.rs | 26 +- 13 files changed, 638 insertions(+), 212 deletions(-) create mode 100644 crates/memtrack/src/ebpf/c/memtrack_perf.bpf.c create mode 100644 crates/memtrack/src/ebpf/c/memtrack_token.bpf.c create mode 100644 crates/memtrack/testdata/flavor_equivalence.c create mode 100644 crates/memtrack/tests/flavor_equivalence_tests.rs create mode 100644 crates/memtrack/tests/snapshots/flavor_equivalence_tests__flavor_equivalence_summary.snap diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b76aa609..23de02f9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,7 +96,14 @@ 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] + test: + [ + c_tests, + cpp_tests, + rust_tests, + spawn_tests, + flavor_equivalence_tests, + ] steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: diff --git a/crates/memtrack/build.rs b/crates/memtrack/build.rs index a9227d13..8e6e4542 100644 --- a/crates/memtrack/build.rs +++ b/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 (memtrack.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/memtrack.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() diff --git a/crates/memtrack/src/ebpf/c/memtrack.bpf.c b/crates/memtrack/src/ebpf/c/memtrack.bpf.c index 06601e27..36704825 100644 --- a/crates/memtrack/src/ebpf/c/memtrack.bpf.c +++ b/crates/memtrack/src/ebpf/c/memtrack.bpf.c @@ -12,6 +12,19 @@ char LICENSE[] SEC("license") = "GPL"; +/* 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 + /* Macros for common map definitions */ #define BPF_HASH_MAP(name, key_type, value_type, max_ents) \ struct { \ @@ -35,6 +48,50 @@ char LICENSE[] SEC("license") = "GPL"; __uint(max_entries, size); \ } name SEC(".maps") +/* 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); +} + /* Map to store PIDs we're tracking */ BPF_HASH_MAP(tracked_pids, __u32, __u8, 10000); /* Map to store parent-child relationships to detect hierarchy */ @@ -71,26 +128,37 @@ static __always_inline int is_tracked(__u32 pid) { return 0; } -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; - - /* Print process fork with PIDs */ - // bpf_printk("sched_fork: parent_pid=%u child_pid=%u", parent_pid, child_pid); - - /* Check if parent is being tracked */ +/* 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)) { - /* Auto-track this child */ __u8 marker = 1; bpf_map_update_elem(&tracked_pids, &child_pid, &marker, BPF_ANY); bpf_map_update_elem(&pids_ppid, &child_pid, &parent_pid, BPF_ANY); - - // bpf_printk("auto-tracking child process: child_pid=%u", child_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 /* == Helper functions for the allocation tracking == */ @@ -110,9 +178,10 @@ static __always_inline int is_enabled(void) { /* Helper to store parameter value in map for tracking between entry and return */ 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; @@ -134,7 +203,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; \ @@ -210,9 +279,9 @@ static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_ /* Macro to generate uprobe/uretprobe pairs for allocation functions with 1 argument */ #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) { \ @@ -228,7 +297,7 @@ static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_ /* Macro for simple return value only functions like free */ #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) { \ @@ -244,7 +313,7 @@ static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_ __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; \ @@ -258,7 +327,7 @@ static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_ 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); \ diff --git a/crates/memtrack/src/ebpf/c/memtrack_perf.bpf.c b/crates/memtrack/src/ebpf/c/memtrack_perf.bpf.c new file mode 100644 index 00000000..b3307edd --- /dev/null +++ b/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 memtrack.bpf.c. */ +#include "memtrack.bpf.c" diff --git a/crates/memtrack/src/ebpf/c/memtrack_token.bpf.c b/crates/memtrack/src/ebpf/c/memtrack_token.bpf.c new file mode 100644 index 00000000..f942e48d --- /dev/null +++ b/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 + * memtrack.bpf.c; only the SEC() annotations differ, keyed on this define. */ +#define MEMTRACK_BPF_LINKS 1 +#include "memtrack.bpf.c" diff --git a/crates/memtrack/src/ebpf/memtrack.rs b/crates/memtrack/src/ebpf/memtrack.rs index 3a1bfa26..f299752b 100644 --- a/crates/memtrack/src/ebpf/memtrack.rs +++ b/crates/memtrack/src/ebpf/memtrack.rs @@ -2,7 +2,7 @@ use crate::prelude::*; use libbpf_rs::Link; use libbpf_rs::skel::OpenSkel; use libbpf_rs::skel::SkelBuilder; -use libbpf_rs::{MapCore, UprobeOpts}; +use libbpf_rs::{MapCore, UprobeMultiOpts, UprobeOpts}; use paste::paste; use std::mem::MaybeUninit; use std::path::Path; @@ -10,20 +10,85 @@ use std::path::Path; use crate::allocators::AllocatorKind; use crate::ebpf::poller::RingBufferPoller; -pub mod memtrack_skel { - include!(concat!(env!("OUT_DIR"), "/memtrack.skel.rs")); +/// Attach via `bpf()` links (uprobe_multi + tp_btf), authorized by a delegated +/// BPF token. Requires a kernel with uprobe_multi (>= 6.6) and is the only path +/// that works inside the macro-agent sandbox. +mod token { + include!(concat!(env!("OUT_DIR"), "/memtrack_token.skel.rs")); +} +/// Classic perf-based attach (uprobe/uretprobe + perf tracepoint). Works on +/// kernels predating uprobe_multi but needs CAP_PERFMON in the init user +/// namespace, so it cannot be delegated into the sandbox. +mod perf { + include!(concat!(env!("OUT_DIR"), "/memtrack_perf.skel.rs")); +} + +/// Whether a delegated BPF token is available to this process. libbpf reads the +/// token from the directory named by `LIBBPF_BPF_TOKEN_PATH` and attaches it to +/// its `bpf()` calls; its presence is also what tells us to use the bpf()-link +/// attach paths instead of the perf-based ones. +fn has_delegated_bpf_token() -> bool { + std::env::var_os("LIBBPF_BPF_TOKEN_PATH").is_some_and(|p| !p.is_empty()) +} + +/// Which set of attach mechanisms a loaded skeleton uses. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Flavor { + /// `bpf()`-link attach (uprobe_multi + tp_btf); delegatable via a BPF token, + /// requires uprobe_multi (kernel >= 6.6). + Token, + /// Classic perf-based attach (uprobe/uretprobe + perf tracepoint); works on + /// older kernels but needs CAP_PERFMON in the init user namespace. + Perf, +} + +/// The loaded skeleton, in whichever attach flavor [`MemtrackBpf::new`] selected. +/// Both flavors are generated from the same BPF source and share identical maps; +/// only the programs' attach mechanism differs. +enum Skel { + Token(Box>), + Perf(Box>), } -pub use memtrack_skel::*; -/// Resolve symbol offset from .symtab to ensure that libbpf can find it. Otherwise -/// it will print a warning at runtime. -fn ensure_symbol_exists(lib_path: &Path, symbol_name: &str) -> Result<()> { +/// Run `$body` against the loaded skeleton, binding `$skel` to the concrete +/// skeleton of whichever flavor is active. Used where both flavors expose the +/// same field/method names (all maps, and the auto-attached programs). +macro_rules! with_skel { + ($self:expr, $skel:ident => $body:expr) => { + match &$self.skel { + Skel::Token($skel) => $body, + Skel::Perf($skel) => $body, + } + }; + (mut $self:expr, $skel:ident => $body:expr) => { + match &mut $self.skel { + Skel::Token($skel) => $body, + Skel::Perf($skel) => $body, + } + }; +} + +/// Device and inode of our PID namespace, as `bpf_get_ns_current_pid_tgid` +/// expects them (`stat` of `/proc/self/ns/pid`). Returns `None` if it can't be +/// read, in which case tracking falls back to global PIDs. +fn current_pidns_ids() -> Option<(u64, u64)> { + use std::os::unix::fs::MetadataExt; + let meta = std::fs::metadata("/proc/self/ns/pid").ok()?; + Some((meta.dev(), meta.ino())) +} + +/// Resolve a symbol to its file offset in the target library. +/// +/// uprobe_multi attaches by offset rather than symbol name: libbpf's own +/// name-based resolution misses some libc symbols (returning ENOENT), so we +/// resolve from the ELF symbol tables ourselves, exactly as the offset the probe +/// needs. Checks both `.symtab` and `.dynsym`. +fn resolve_symbol_offset(lib_path: &Path, symbol_name: &str) -> Result { use object::{Object, ObjectSymbol}; let data = std::fs::read(lib_path)?; let file = object::File::parse(&*data)?; - // Check both regular and dynamic symbols for symbol in file.symbols().chain(file.dynamic_symbols()) { if !symbol.is_definition() { continue; @@ -36,7 +101,7 @@ fn ensure_symbol_exists(lib_path: &Path, symbol_name: &str) -> Result<()> { if name == symbol_name { let addr = symbol.address(); if addr != 0 { - return Ok(()); + return Ok(addr as usize); } } } @@ -44,58 +109,62 @@ fn ensure_symbol_exists(lib_path: &Path, symbol_name: &str) -> Result<()> { bail!("Symbol {symbol_name} not found in {}", lib_path.display()) } +/// Attach a single program (entry or return) to `symbol` in `lib_path`, using +/// the attach mechanism matching the loaded skeleton flavor. +/// +/// `$prog` is the program field name, identical across both flavors; `$skel` +/// binds the concrete skeleton so field access type-checks in each arm. +macro_rules! attach_one { + ($self:expr, $prog:ident, $lib_path:expr, $offset:expr, $retprobe:expr) => {{ + let lib_path = $lib_path; + let offset = $offset; + let retprobe = $retprobe; + match &mut $self.skel { + Skel::Token(skel) => skel.progs.$prog.attach_uprobe_multi_with_opts( + -1, + lib_path, + "", + UprobeMultiOpts { + offsets: vec![offset], + retprobe, + ..Default::default() + }, + ), + Skel::Perf(skel) => skel.progs.$prog.attach_uprobe_with_opts( + -1, + lib_path, + offset, + UprobeOpts { + retprobe, + ..Default::default() + }, + ), + } + }}; +} + /// Macro to attach a function with both entry and return probes. /// Also generates a `try_attach_*` variant that logs errors instead of returning them. /// -/// Uses offset-based attachment by resolving symbols from .symtab. -/// Fails if the symbol is not found. +/// Resolves the symbol to a file offset and attaches through whichever +/// mechanism the loaded skeleton uses. Fails if the symbol is not found. macro_rules! attach_uprobe_uretprobe { ($name:ident, $prog_entry:ident, $prog_return:ident) => { fn $name(&mut self, lib_path: &Path, symbol: &str) -> Result<()> { - ensure_symbol_exists(lib_path, symbol)?; - - // Attach entry probe at function entry via func_name - let link = self - .skel - .progs - .$prog_entry - .attach_uprobe_with_opts( - -1, - lib_path, - 0, - UprobeOpts { - retprobe: false, - func_name: Some(symbol.to_owned()), - ..Default::default() - }, - ) - .context(format!( - "Failed to attach {} uprobe in {}", - symbol, - lib_path.display() - ))?; + let offset = resolve_symbol_offset(lib_path, symbol)?; + + let link = attach_one!(self, $prog_entry, lib_path, offset, false).context(format!( + "Failed to attach {} uprobe in {}", + symbol, + lib_path.display() + ))?; self.probes.push(link); - // Attach return probe at function entry via func_name - let link = self - .skel - .progs - .$prog_return - .attach_uprobe_with_opts( - -1, - lib_path, - 0, - UprobeOpts { - retprobe: true, - func_name: Some(symbol.to_owned()), - ..Default::default() - }, - ) - .context(format!( - "Failed to attach {} uretprobe in {}", - symbol, - lib_path.display() - ))?; + let link = attach_one!(self, $prog_return, lib_path, offset, true).context(format!( + "Failed to attach {} uretprobe in {}", + symbol, + lib_path.display() + ))?; self.probes.push(link); Ok(()) @@ -113,32 +182,18 @@ macro_rules! attach_uprobe_uretprobe { /// Macro to attach a function with only an entry probe (no return probe). /// Also generates a `try_attach_*` variant that logs errors instead of returning them. /// -/// Uses offset-based attachment by resolving symbols from .symtab. -/// Fails if the symbol is not found. +/// Resolves the symbol to a file offset and attaches through whichever +/// mechanism the loaded skeleton uses. Fails if the symbol is not found. macro_rules! attach_uprobe { ($name:ident, $prog:ident) => { fn $name(&mut self, lib_path: &Path, symbol: &str) -> Result<()> { - ensure_symbol_exists(lib_path, symbol)?; - - let link = self - .skel - .progs - .$prog - .attach_uprobe_with_opts( - -1, - lib_path, - 0, - UprobeOpts { - retprobe: false, - func_name: Some(symbol.to_owned()), - ..Default::default() - }, - ) - .context(format!( - "Failed to attach {} uprobe in {}", - symbol, - lib_path.display() - ))?; + let offset = resolve_symbol_offset(lib_path, symbol)?; + + let link = attach_one!(self, $prog, lib_path, offset, false).context(format!( + "Failed to attach {} uprobe in {}", + symbol, + lib_path.display() + ))?; self.probes.push(link); Ok(()) } @@ -152,45 +207,76 @@ macro_rules! attach_uprobe { }; } -macro_rules! attach_tracepoint { - ($func:ident, $prog:ident) => { - fn $func(&mut self) -> Result<()> { - let link = self - .skel - .progs - .$prog - .attach() - .context(format!("Failed to attach {} tracepoint", stringify!($prog)))?; - self.probes.push(link); - Ok(()) - } - }; - ($name:ident) => { - paste! { - attach_tracepoint!([], []); - } - }; -} - pub struct MemtrackBpf { - skel: Box>, + skel: Skel, probes: Vec, } +/// Set the PID-namespace rodata on an open skeleton, so both flavors resolve +/// PIDs the same way. `$open` is the open skeleton, of either flavor. +macro_rules! set_pidns_rodata { + ($open:expr, $ids:expr) => { + if let (Some((dev, ino)), Some(rodata)) = ($ids, $open.maps.rodata_data.as_deref_mut()) { + rodata.target_pidns_dev = dev; + rodata.target_pidns_ino = ino; + } + }; +} + impl MemtrackBpf { + /// Load the skeleton, selecting the attach flavor from the environment: a + /// delegated BPF token means we are the sandboxed workload and must use the + /// token path; otherwise the perf path, which also supports kernels + /// predating uprobe_multi. pub fn new() -> Result { - // Build and open the syscalls BPF program - let builder = MemtrackSkelBuilder::default(); - let open_object = Box::leak(Box::new(MaybeUninit::uninit())); - let open_skel = builder - .open(open_object) - .context("Failed to open syscalls BPF skeleton")?; - - let skel = Box::new( - open_skel - .load() - .context("Failed to load syscalls BPF skeleton")?, - ); + let flavor = if has_delegated_bpf_token() { + Flavor::Token + } else { + Flavor::Perf + }; + Self::with_flavor(flavor) + } + + /// Load the skeleton for a specific attach flavor, bypassing environment + /// detection. Used by tests to exercise either path directly; both attach + /// fine on a privileged kernel without a token (the token only authorizes + /// `bpf()` from inside the unprivileged sandbox). + pub fn with_flavor(flavor: Flavor) -> Result { + // Resolve PIDs relative to our own PID namespace. When we run inside a + // namespace (e.g. the macro-agent sandbox), the PIDs we register are + // namespace-local while eBPF sees global PIDs; telling the programs which + // namespace to resolve in keeps tracking correct. In the init namespace + // this resolves to global PIDs, matching the previous behavior. + let pidns_ids = current_pidns_ids(); + + let skel = match flavor { + Flavor::Token => { + let builder = token::MemtrackTokenSkelBuilder::default(); + let open_object = Box::leak(Box::new(MaybeUninit::uninit())); + let mut open_skel = builder + .open(open_object) + .context("Failed to open memtrack BPF skeleton")?; + set_pidns_rodata!(open_skel, pidns_ids); + Skel::Token(Box::new( + open_skel + .load() + .context("Failed to load memtrack BPF skeleton")?, + )) + } + Flavor::Perf => { + let builder = perf::MemtrackPerfSkelBuilder::default(); + let open_object = Box::leak(Box::new(MaybeUninit::uninit())); + let mut open_skel = builder + .open(open_object) + .context("Failed to open memtrack BPF skeleton")?; + set_pidns_rodata!(open_skel, pidns_ids); + Skel::Perf(Box::new( + open_skel + .load() + .context("Failed to load memtrack BPF skeleton")?, + )) + } + }; Ok(Self { skel, @@ -199,15 +285,12 @@ impl MemtrackBpf { } pub fn add_tracked_pid(&mut self, pid: i32) -> Result<()> { - self.skel - .maps - .tracked_pids - .update( - &pid.to_le_bytes(), - &1u8.to_le_bytes(), - libbpf_rs::MapFlags::ANY, - ) - .context("Failed to add PID to uprobes tracked set")?; + with_skel!(self, skel => skel.maps.tracked_pids.update( + &pid.to_le_bytes(), + &1u8.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + )) + .context("Failed to add PID to uprobes tracked set")?; Ok(()) } @@ -216,28 +299,24 @@ impl MemtrackBpf { pub fn enable_tracking(&mut self) -> Result<()> { let key = 0u32; let value = true as u8; - self.skel - .maps - .tracking_enabled - .update( - &key.to_le_bytes(), - &value.to_le_bytes(), - libbpf_rs::MapFlags::ANY, - ) - .context("Failed to enable tracking")?; + with_skel!(self, skel => skel.maps.tracking_enabled.update( + &key.to_le_bytes(), + &value.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + )) + .context("Failed to enable tracking")?; Ok(()) } /// Read the count of events dropped because the ring buffer was full. pub fn dropped_events_count(&self) -> Result { let key = 0u32; - let value = self - .skel + let value = with_skel!(self, skel => skel .maps .dropped_events - .lookup(&key.to_le_bytes(), libbpf_rs::MapFlags::ANY) - .context("Failed to read dropped_events counter")? - .ok_or_else(|| anyhow!("dropped_events slot 0 missing"))?; + .lookup(&key.to_le_bytes(), libbpf_rs::MapFlags::ANY)) + .context("Failed to read dropped_events counter")? + .ok_or_else(|| anyhow!("dropped_events slot 0 missing"))?; let bytes: [u8; 8] = value .as_slice() @@ -250,15 +329,12 @@ impl MemtrackBpf { pub fn disable_tracking(&mut self) -> Result<()> { let key = 0u32; let value = false as u8; - self.skel - .maps - .tracking_enabled - .update( - &key.to_le_bytes(), - &value.to_le_bytes(), - libbpf_rs::MapFlags::ANY, - ) - .context("Failed to disable tracking")?; + with_skel!(self, skel => skel.maps.tracking_enabled.update( + &key.to_le_bytes(), + &value.to_le_bytes(), + libbpf_rs::MapFlags::ANY, + )) + .context("Failed to disable tracking")?; Ok(()) } @@ -446,10 +522,13 @@ impl MemtrackBpf { Ok(()) } - attach_tracepoint!(sched_fork); pub fn attach_tracepoints(&mut self) -> Result<()> { - self.attach_sched_fork()?; + // The fork hook auto-attaches by its section (tp_btf or classic + // tracepoint, depending on the flavor); both go through `attach()`. + let link = with_skel!(mut self, skel => skel.progs.tracepoint_sched_fork.attach()) + .context("Failed to attach sched_process_fork tracepoint")?; + self.probes.push(link); Ok(()) } @@ -461,8 +540,7 @@ impl MemtrackBpf { RingBufferPoller, std::sync::mpsc::Receiver, )> { - // Use the syscalls skeleton's ring buffer (both programs share the same one) - RingBufferPoller::with_channel(&self.skel.maps.events, poll_timeout_ms) + with_skel!(self, skel => RingBufferPoller::with_channel(&skel.maps.events, poll_timeout_ms)) } } diff --git a/crates/memtrack/src/ebpf/mod.rs b/crates/memtrack/src/ebpf/mod.rs index 0badf542..e7a6be2e 100644 --- a/crates/memtrack/src/ebpf/mod.rs +++ b/crates/memtrack/src/ebpf/mod.rs @@ -3,5 +3,5 @@ mod memtrack; mod poller; mod tracker; -pub use memtrack::MemtrackBpf; +pub use memtrack::{Flavor, MemtrackBpf}; pub use tracker::Tracker; diff --git a/crates/memtrack/src/ebpf/tracker.rs b/crates/memtrack/src/ebpf/tracker.rs index 9f1af587..e4872593 100644 --- a/crates/memtrack/src/ebpf/tracker.rs +++ b/crates/memtrack/src/ebpf/tracker.rs @@ -1,3 +1,4 @@ +use crate::ebpf::Flavor; use crate::prelude::*; use crate::{AllocatorLib, ebpf::MemtrackBpf}; use runner_shared::artifacts::MemtrackEvent as Event; @@ -16,7 +17,17 @@ impl Tracker { /// - Attach uprobes to all libc instances /// - Attach tracepoints for fork tracking pub fn new() -> Result { - let mut instance = Self::new_without_allocators()?; + Self::with_bpf(MemtrackBpf::new()?) + } + + /// Like [`Tracker::new`], but pinned to a specific attach flavor instead of + /// the environment-selected one. Used by tests to exercise either path. + pub fn with_flavor(flavor: Flavor) -> Result { + Self::with_bpf(MemtrackBpf::with_flavor(flavor)?) + } + + fn with_bpf(bpf: MemtrackBpf) -> Result { + let mut instance = Self::from_bpf_without_allocators(bpf)?; let allocators = AllocatorLib::find_all()?; debug!("Found {} allocator instance(s)", allocators.len()); @@ -26,10 +37,13 @@ impl Tracker { } pub fn new_without_allocators() -> Result { + Self::from_bpf_without_allocators(MemtrackBpf::new()?) + } + + fn from_bpf_without_allocators(mut bpf: MemtrackBpf) -> Result { // Bump memlock limits Self::bump_memlock_rlimit()?; - let mut bpf = MemtrackBpf::new()?; bpf.attach_tracepoints()?; Ok(Self { bpf }) diff --git a/crates/memtrack/testdata/flavor_equivalence.c b/crates/memtrack/testdata/flavor_equivalence.c new file mode 100644 index 00000000..ed81d780 --- /dev/null +++ b/crates/memtrack/testdata/flavor_equivalence.c @@ -0,0 +1,50 @@ +/* Deterministic allocation workload for the perf-vs-token flavor equivalence + * test. Brackets a fixed sequence of allocations with 0xC0D59EED marker + * mallocs so the test can isolate exactly these events from libc/runtime noise + * (see assert_events_with_marker! in tests/shared.rs). + * + * The leading sleep gives the tracker time to attach its probes and register + * this PID before the marked allocations run. The fork exercises the + * fork-follow hook (tp_btf in the token flavor, a classic tracepoint in the + * perf flavor), which must behave identically. */ +#include +#include +#include +#include + +static void marker(void) { + void* m = malloc(0xC0D59EED); + free(m); +} + +int main(void) { + sleep(1); + + marker(); + + /* A fixed, varied set of allocations the test asserts on. */ + void* a = malloc(4096); + memset(a, 1, 4096); + void* c = calloc(16, 256); + void* r = malloc(1024); + r = realloc(r, 8192); + void* al = aligned_alloc(64, 2048); + + /* Fork a child that allocates and frees; the parent must inherit-track it. */ + pid_t pid = fork(); + if (pid == 0) { + void* child = malloc(333); + free(child); + _exit(0); + } + waitpid(pid, NULL, 0); + + free(a); + free(c); + free(r); + free(al); + + marker(); + + return 0; +} diff --git a/crates/memtrack/tests/flavor_equivalence_tests.rs b/crates/memtrack/tests/flavor_equivalence_tests.rs new file mode 100644 index 00000000..915c22e5 --- /dev/null +++ b/crates/memtrack/tests/flavor_equivalence_tests.rs @@ -0,0 +1,123 @@ +//! Verifies the two eBPF attach flavors produce equivalent tracking results. +//! +//! memtrack attaches its probes one of two ways: bpf()-based links +//! (`uprobe_multi` + `tp_btf`, [`Flavor::Token`]) that a delegated BPF token can +//! authorize inside the macro-agent sandbox, or classic perf-based +//! uprobe/tracepoint attach ([`Flavor::Perf`]) that works on kernels predating +//! `uprobe_multi` but cannot be delegated. Both are generated from the same BPF +//! source; only the attach mechanism differs. This test runs the same +//! deterministic workload through each and asserts the tracked allocations +//! match, so the fallback path can't silently diverge from the sandbox path. +//! +//! Attaching either flavor requires privilege (root or CAP_BPF/CAP_PERFMON); +//! the token flavor additionally needs `uprobe_multi` (kernel >= 6.6). The test +//! skips any path it cannot exercise on the current host rather than failing. + +#[macro_use] +mod shared; + +use std::collections::BTreeMap; +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use memtrack::Flavor; +use runner_shared::artifacts::{MemtrackEvent as Event, MemtrackEventKind}; +use tempfile::TempDir; + +fn compile_c_source(source: &str, name: &str, out_dir: &Path) -> anyhow::Result { + let source_path = out_dir.join(format!("{name}.c")); + let binary_path = out_dir.join(name); + fs::write(&source_path, source)?; + + let output = Command::new("gcc") + .args(["-O0", "-o", binary_path.to_str().unwrap()]) + .arg(&source_path) + .output()?; + if !output.status.success() { + anyhow::bail!("gcc failed: {}", String::from_utf8_lossy(&output.stderr)); + } + Ok(binary_path) +} + +/// A stable summary of the tracked allocations: how many events of each kind, +/// and the total bytes they account for. Independent of addresses, timestamps, +/// and event ordering, all of which legitimately differ run to run, so it is +/// safe to compare across two separate tracking runs. +type Summary = BTreeMap; + +fn summarize(events: &[Event]) -> Summary { + let mut summary: Summary = BTreeMap::new(); + for event in events { + let (kind, size) = match event.kind { + MemtrackEventKind::Malloc { size } => ("Malloc", size), + MemtrackEventKind::Free => ("Free", 0), + MemtrackEventKind::Realloc { size, .. } => ("Realloc", size), + MemtrackEventKind::Calloc { size } => ("Calloc", size), + MemtrackEventKind::AlignedAlloc { size } => ("AlignedAlloc", size), + MemtrackEventKind::Mmap { size } => ("Mmap", size), + MemtrackEventKind::Munmap { size } => ("Munmap", size), + MemtrackEventKind::Brk { size } => ("Brk", size), + }; + let entry = summary.entry(kind.to_string()).or_insert((0, 0)); + entry.0 += 1; + entry.1 += size; + } + summary +} + +/// Track the workload under one flavor, returning the marker-isolated summary, +/// or `None` if this flavor cannot be attached on the current host (e.g. the +/// token flavor on a kernel without `uprobe_multi`). +fn summary_for_flavor(binary: &Path, flavor: Flavor) -> Option { + let command = Command::new(binary); + match shared::track_command_with_flavor(command, &[], flavor) { + Ok((events, handle)) => { + handle.join().ok(); + let filtered = shared::between_markers(&events); + assert!( + !filtered.is_empty(), + "{flavor:?}: no events captured between markers" + ); + Some(summarize(&filtered)) + } + Err(err) => { + eprintln!("skipping {flavor:?} flavor: cannot attach on this host: {err:#}"); + None + } + } +} + +#[test_with::env(GITHUB_ACTIONS)] +#[test_log::test] +fn both_flavors_track_equivalently() -> anyhow::Result<()> { + let temp_dir = TempDir::new()?; + let binary = compile_c_source( + include_str!("../testdata/flavor_equivalence.c"), + "flavor_equivalence", + temp_dir.path(), + )?; + + let Some(perf) = summary_for_flavor(&binary, Flavor::Perf) else { + // No privilege to attach at all — nothing to compare. + eprintln!("skipping: perf flavor could not attach (need root or CAP_BPF/CAP_PERFMON)"); + return Ok(()); + }; + + let Some(token) = summary_for_flavor(&binary, Flavor::Token) else { + // Perf worked but the token flavor is unavailable (kernel < 6.6). The + // perf fallback is still validated; there is nothing to compare against. + return Ok(()); + }; + + assert_eq!( + perf, token, + "perf and token flavors disagree on tracked allocations\nperf: {perf:#?}\ntoken: {token:#?}" + ); + + // Both flavors agree, so a single snapshot pins down what the workload is + // expected to produce and guards against both paths drifting together. + insta::assert_debug_snapshot!("flavor_equivalence_summary", perf); + + Ok(()) +} diff --git a/crates/memtrack/tests/shared.rs b/crates/memtrack/tests/shared.rs index 054b3b8e..d5594997 100644 --- a/crates/memtrack/tests/shared.rs +++ b/crates/memtrack/tests/shared.rs @@ -2,7 +2,7 @@ use anyhow::Context; use memtrack::prelude::*; -use memtrack::{AllocatorLib, Tracker}; +use memtrack::{AllocatorLib, Flavor, Tracker}; use runner_shared::artifacts::{MemtrackEvent as Event, MemtrackEventKind}; use std::path::Path; use std::process::Command; @@ -71,35 +71,34 @@ macro_rules! assert_events_snapshot { /// ``` macro_rules! assert_events_with_marker { ($name:expr, $events:expr) => {{ - use itertools::Itertools; - use runner_shared::artifacts::MemtrackEventKind; - use std::mem::discriminant; - - // Remove events outside our 0xC0D59EED marker allocations - let filtered_events = $events - .iter() - .sorted_by_key(|e| e.timestamp) - .dedup_by(|a, b| a.addr == b.addr && discriminant(&a.kind) == discriminant(&b.kind)) - .skip_while(|e| { - let MemtrackEventKind::Malloc { size } = e.kind else { - return true; - }; - size != 0xC0D59EED - }) - .skip(2) // Skip the marker allocation and free - .take_while(|e| { - let MemtrackEventKind::Malloc { size } = e.kind else { - return true; - }; - size != 0xC0D59EED - }) - .cloned() - .collect::>(); - - assert_events_snapshot!($name, filtered_events); + let filtered_events = shared::between_markers($events); + assert_events_snapshot!($name, &filtered_events); }}; } +/// Keep only the deterministic allocations the workload brackets with +/// `0xC0D59EED` marker mallocs, dropping libc/runtime noise. Sorts by +/// timestamp and dedups by `(addr, kind)` first, so ordering and duplicate +/// tracking don't leak into the result. +pub fn between_markers(events: &[Event]) -> Vec { + use itertools::Itertools; + use std::mem::discriminant; + + const MARKER: u64 = 0xC0D5_9EED; + let is_marker = + |e: &&Event| matches!(e.kind, MemtrackEventKind::Malloc { size } if size == MARKER); + + events + .iter() + .sorted_by_key(|e| e.timestamp) + .dedup_by(|a, b| a.addr == b.addr && discriminant(&a.kind) == discriminant(&b.kind)) + .skip_while(|e| !is_marker(e)) + .skip(2) // the opening marker malloc and its free + .take_while(|e| !is_marker(e)) + .cloned() + .collect() +} + /// Compile a Rust binary from a test crate directory. pub fn compile_rust_binary( crate_dir: &Path, @@ -129,17 +128,37 @@ pub fn compile_rust_binary( /// When `discover_system_allocators` is true, the tracker will scan for all /// allocators on the system (slower). When false, only `extra_allocators` are used. pub fn track_command( - mut command: Command, + command: Command, extra_allocators: &[AllocatorLib], discover_system_allocators: bool, ) -> TrackResult { // IMPORTANT: Always initialize the tracker BEFORE spawning the binary, as it can take some time to // attach to all the allocator libraries (especially when using NixOS). - let mut tracker = if discover_system_allocators { + let tracker = if discover_system_allocators { memtrack::Tracker::new()? } else { memtrack::Tracker::new_without_allocators()? }; + track_command_with_tracker(command, extra_allocators, tracker) +} + +/// Like [`track_command`], but pins the tracker to a specific attach flavor +/// (perf vs. bpf-token links) instead of the environment-selected default. +/// Discovers system allocators, matching [`track_binary`]. +pub fn track_command_with_flavor( + command: Command, + extra_allocators: &[AllocatorLib], + flavor: Flavor, +) -> TrackResult { + let tracker = memtrack::Tracker::with_flavor(flavor)?; + track_command_with_tracker(command, extra_allocators, tracker) +} + +fn track_command_with_tracker( + mut command: Command, + extra_allocators: &[AllocatorLib], + mut tracker: Tracker, +) -> TrackResult { tracker.attach_allocators(extra_allocators)?; let child = command.spawn().context("Failed to spawn command")?; diff --git a/crates/memtrack/tests/snapshots/flavor_equivalence_tests__flavor_equivalence_summary.snap b/crates/memtrack/tests/snapshots/flavor_equivalence_tests__flavor_equivalence_summary.snap new file mode 100644 index 00000000..d17c6606 --- /dev/null +++ b/crates/memtrack/tests/snapshots/flavor_equivalence_tests__flavor_equivalence_summary.snap @@ -0,0 +1,26 @@ +--- +source: crates/memtrack/tests/flavor_equivalence_tests.rs +expression: perf +--- +{ + "AlignedAlloc": ( + 1, + 2048, + ), + "Calloc": ( + 1, + 4096, + ), + "Free": ( + 5, + 0, + ), + "Malloc": ( + 3, + 5453, + ), + "Realloc": ( + 1, + 8192, + ), +} diff --git a/src/executor/memory/executor.rs b/src/executor/memory/executor.rs index 24032868..152dbafb 100644 --- a/src/executor/memory/executor.rs +++ b/src/executor/memory/executor.rs @@ -77,11 +77,11 @@ impl MemoryExecutor { Ok((ipc_server, cmd_builder, env_file)) } - /// Ensure memtrack can load its eBPF programs: either run as root or hold the - /// required file capabilities. Tries a one-time capability grant if neither - /// holds, and bails if that fails. + /// Ensure memtrack can load its eBPF programs: either run as root, hold the + /// required file capabilities, or be handed a delegated BPF token. Tries a + /// one-time capability grant if none holds, and bails if that fails. fn ensure_privileges() -> Result<()> { - if is_root_user() || has_memtrack_capabilities() { + if is_root_user() || has_delegated_bpf_token() || has_memtrack_capabilities() { return Ok(()); } @@ -98,6 +98,14 @@ impl MemoryExecutor { } } +/// Whether a delegated BPF token is available (the sandbox's mechanism for +/// loading eBPF without host privileges): `LIBBPF_BPF_TOKEN_PATH` points at a +/// live bpffs, which memtrack's libbpf uses to authorize its `bpf()` calls. +fn has_delegated_bpf_token() -> bool { + std::env::var_os("LIBBPF_BPF_TOKEN_PATH") + .is_some_and(|p| !p.is_empty() && Path::new(&p).is_dir()) +} + #[async_trait(?Send)] impl Executor for MemoryExecutor { fn name(&self) -> ExecutorName { @@ -114,6 +122,11 @@ impl Executor for MemoryExecutor { detail: "running as root".to_string(), }); } + if has_delegated_bpf_token() { + return Some(PrivilegeStatus::Satisfied { + detail: "delegated BPF token".to_string(), + }); + } if has_memtrack_capabilities() { return Some(PrivilegeStatus::Satisfied { detail: "capabilities granted".to_string(), @@ -140,6 +153,11 @@ impl Executor for MemoryExecutor { } fn grant_privileges(&self) -> Result<()> { + // A delegated BPF token already provides the privilege memtrack needs; + // there is nothing to grant (and no sudo to prompt for) in that case. + if has_delegated_bpf_token() { + return Ok(()); + } ensure_memtrack_capabilities() }