chore: wip [skip ci]

Author: not-matthias

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

@@ -10,14 +10,42 @@
 #define EVENT_TYPE_MUNMAP 7
 #define EVENT_TYPE_BRK 8
 
-/* Event structure - shared between BPF and userspace */
-struct event {
+/* Common header shared by all event types */
+struct event_header {
     uint8_t event_type; /* See EVENT_TYPE_* constants above */
     uint64_t timestamp; /* monotonic time in nanoseconds (CLOCK_MONOTONIC) */
     uint32_t pid;
     uint32_t tid;
-    uint64_t addr; /* address returned/freed */
-    uint64_t size; /* size requested */
+};
+
+/* Tagged union event structure */
+struct event {
+    struct event_header header;
+    union {
+        /* Allocation events (malloc, calloc, aligned_alloc) */
+        struct {
+            uint64_t addr; /* address returned */
+            uint64_t size; /* size requested */
+        } alloc;
+
+        /* Deallocation event (free) */
+        struct {
+            uint64_t addr; /* address to free */
+        } free;
+
+        /* Reallocation event - includes both old and new addresses */
+        struct {
+            uint64_t old_addr; /* previous address (can be NULL) */
+            uint64_t new_addr; /* new address returned */
+            uint64_t size;     /* new size requested */
+        } realloc;
+
+        /* Memory mapping events (mmap, munmap, brk) */
+        struct {
+            uint64_t addr; /* address of mapping */
+            uint64_t size; /* size of mapping */
+        } mmap;
+    } data;
 };
 
 #endif /* __EVENT_H__ */

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

@@ -123,29 +123,64 @@ static __always_inline __u64* take_param(void* map) {
     return value;
 }
 
-/* Helper to submit an event to the ring buffer */
-static __always_inline int submit_event(__u64 addr, __u64 size, __u8 event_type) {
-    __u64 tid = bpf_get_current_pid_tgid();
-    __u32 pid = tid >> 32;
-
-    if (!is_tracked(pid) || !is_enabled()) {
-        return 0;
+/* Macro to handle common event submission boilerplate
+ * Usage: SUBMIT_EVENT(event_type, { e->data.foo = bar; })
+ */
+#define SUBMIT_EVENT(evt_type, fill_data)                              \
+    {                                                                   \
+        __u64 tid = bpf_get_current_pid_tgid();                        \
+        __u32 pid = tid >> 32;                                         \
+                                                                        \
+        if (!is_tracked(pid) || !is_enabled()) {                       \
+            return 0;                                                   \
+        }                                                               \
+                                                                        \
+        struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0); \
+        if (!e) {                                                       \
+            return 0;                                                   \
+        }                                                               \
+                                                                        \
+        e->header.timestamp = bpf_ktime_get_ns();                      \
+        e->header.pid = pid;                                            \
+        e->header.tid = tid & 0xFFFFFFFF;                              \
+        e->header.event_type = evt_type;                               \
+                                                                        \
+        fill_data;                                                      \
+                                                                        \
+        bpf_ringbuf_submit(e, 0);                                       \
+        return 0;                                                       \
     }
 
-    struct event* e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
-    if (!e) {
-        return 0;
-    }
+/* Helper to submit an allocation event (malloc, calloc, aligned_alloc) */
+static __always_inline int submit_alloc_event(__u64 addr, __u64 size, __u8 event_type) {
+    SUBMIT_EVENT(event_type, {
+        e->data.alloc.addr = addr;
+        e->data.alloc.size = size;
+    });
+}
 
-    e->timestamp = bpf_ktime_get_ns();
-    e->pid = pid;
-    e->tid = tid & 0xFFFFFFFF;
-    e->event_type = event_type;
-    e->addr = addr;
-    e->size = size;
-    bpf_ringbuf_submit(e, 0);
+/* Helper to submit a free event */
+static __always_inline int submit_free_event(__u64 addr, __u8 event_type) {
+    SUBMIT_EVENT(event_type, {
+        e->data.free.addr = addr;
+    });
+}
 
-    return 0;
+/* Helper to submit a realloc event with both old and new addresses */
+static __always_inline int submit_realloc_event(__u64 old_addr, __u64 new_addr, __u64 size) {
+    SUBMIT_EVENT(EVENT_TYPE_REALLOC, {
+        e->data.realloc.old_addr = old_addr;
+        e->data.realloc.new_addr = new_addr;
+        e->data.realloc.size = size;
+    });
+}
+
+/* Helper to submit a memory mapping event (mmap, munmap, brk) */
+static __always_inline int submit_mmap_event(__u64 addr, __u64 size, __u8 event_type) {
+    SUBMIT_EVENT(event_type, {
+        e->data.mmap.addr = addr;
+        e->data.mmap.size = size;
+    });
 }
 
 /* Macro to generate uprobe/uretprobe pairs for allocation functions */
@@ -163,7 +198,7 @@ static __always_inline int submit_event(__u64 addr, __u64 size, __u8 event_type)
         if (addr == 0) {                                                                    \
             return 0;                                                                       \
         }                                                                                   \
-        return submit_event(addr, *size_ptr, event_type);                                   \
+        return submit_alloc_event(addr, *size_ptr, event_type);                             \
     }
 
 /* Macro for simple address-only functions like free */
@@ -174,7 +209,7 @@ static __always_inline int submit_event(__u64 addr, __u64 size, __u8 event_type)
         if (addr == 0) {                              \
             return 0;                                 \
         }                                             \
-        return submit_event(addr, 0, event_type);     \
+        return submit_free_event(addr, event_type);   \
     }
 
 /* malloc: allocates with size parameter */
@@ -186,8 +221,49 @@ UPROBE_ADDR_ONLY(free, PT_REGS_PARM1(ctx), EVENT_TYPE_FREE)
 /* calloc: allocates with nmemb * size */
 UPROBE_WITH_ARGS(calloc, PT_REGS_PARM1(ctx) * PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_CALLOC)
 
-/* realloc: reallocates with new size */
-UPROBE_WITH_ARGS(realloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_REALLOC)
+/* realloc: reallocates with new size (custom implementation to capture old_addr) */
+struct realloc_params {
+    __u64 old_addr;
+    __u64 size;
+};
+BPF_HASH_MAP(realloc_params, __u64, struct realloc_params, 10000);
+
+SEC("uprobe")
+int uprobe_realloc(struct pt_regs* ctx) {
+    __u64 tid = bpf_get_current_pid_tgid();
+    __u32 pid = tid >> 32;
+
+    if (!is_tracked(pid)) {
+        return 0;
+    }
+
+    struct realloc_params params = {
+        .old_addr = PT_REGS_PARM1(ctx),
+        .size = PT_REGS_PARM2(ctx),
+    };
+    bpf_map_update_elem(&realloc_params, &tid, &params, BPF_ANY);
+    return 0;
+}
+
+SEC("uretprobe")
+int uretprobe_realloc(struct pt_regs* ctx) {
+    __u64 tid = bpf_get_current_pid_tgid();
+    struct realloc_params* params = bpf_map_lookup_elem(&realloc_params, &tid);
+    if (!params) {
+        return 0;
+    }
+
+    /* Take params and delete from map */
+    struct realloc_params p = *params;
+    bpf_map_delete_elem(&realloc_params, &tid);
+
+    __u64 new_addr = PT_REGS_RC(ctx);
+    if (new_addr == 0) {
+        return 0;
+    }
+
+    return submit_realloc_event(p.old_addr, new_addr, p.size);
+}
 
 /* aligned_alloc: allocates with alignment and size */
 UPROBE_WITH_ARGS(aligned_alloc, PT_REGS_PARM2(ctx), PT_REGS_RC(ctx), EVENT_TYPE_ALIGNED_ALLOC)
@@ -234,7 +310,7 @@ int tracepoint_sys_exit_mmap(struct trace_event_raw_sys_exit* ctx) {
         return 0;
     }
 
-    return submit_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
+    return submit_mmap_event((__u64)ret, args->len, EVENT_TYPE_MMAP);
 }
 
 /* munmap tracking */
@@ -247,7 +323,7 @@ int tracepoint_sys_enter_munmap(struct trace_event_raw_sys_enter* ctx) {
         return 0;
     }
 
-    return submit_event(addr, len, EVENT_TYPE_MUNMAP);
+    return submit_mmap_event(addr, len, EVENT_TYPE_MUNMAP);
 }
 
 /* brk tracking - adjusts the program break (heap boundary) */
@@ -281,5 +357,5 @@ int tracepoint_sys_exit_brk(struct trace_event_raw_sys_exit* ctx) {
         return 0;
     }
 
-    return submit_event(new_brk, 0, EVENT_TYPE_BRK);
+    return submit_mmap_event(new_brk, 0, EVENT_TYPE_BRK);
 }