Detect more VMA types

Golang heaps can be determined by a pattern in the address, dictated by address hints supplied to mmap while allocating memory for them.
Thread stacks can be identified by tracking the stack VMA for all newly created threads.

pkg/ebpf/c/common/memory.h

@@ -7,10 +7,14 @@
 
 enum vma_type
 {
+    VMA_FILE_BACKED,
     VMA_STACK,
     VMA_HEAP,
+    VMA_GOLANG_HEAP,
+    VMA_THREAD_STACK,
+    VMA_VDSO,
     VMA_ANON,
-    VMA_OTHER
+    VMA_UNKNOWN,
 };
 
 // PROTOTYPES
@@ -22,11 +26,14 @@ statfunc unsigned long get_env_start_from_mm(struct mm_struct *);
 statfunc unsigned long get_env_end_from_mm(struct mm_struct *);
 statfunc unsigned long get_vma_flags(struct vm_area_struct *);
 statfunc struct vm_area_struct *find_vma(void *ctx, struct task_struct *task, u64 addr);
-statfunc bool vma_is_stack(struct vm_area_struct *vma);
-statfunc bool vma_is_heap(struct vm_area_struct *vma);
+statfunc bool vma_is_file_backed(struct vm_area_struct *vma);
+statfunc bool vma_is_initial_stack(struct vm_area_struct *vma);
+statfunc bool vma_is_initial_heap(struct vm_area_struct *vma);
 statfunc bool vma_is_anon(struct vm_area_struct *vma);
+statfunc bool vma_is_golang_heap(struct vm_area_struct *vma);
+statfunc bool vma_is_thread_stack(struct task_struct *task, struct vm_area_struct *vma);
 statfunc bool vma_is_vdso(struct vm_area_struct *vma);
-statfunc enum vma_type get_vma_type(struct vm_area_struct *vma);
+statfunc enum vma_type get_vma_type(struct task_struct *task, struct vm_area_struct *vma);
 
 // FUNCTIONS
 
@@ -121,7 +128,12 @@ statfunc struct vm_area_struct *find_vma(void *ctx, struct task_struct *task, u6
     return vma;
 }
 
-statfunc bool vma_is_stack(struct vm_area_struct *vma)
+statfunc bool vma_is_file_backed(struct vm_area_struct *vma)
+{
+    return BPF_CORE_READ(vma, vm_file) != NULL;
+}
+
+statfunc bool vma_is_initial_stack(struct vm_area_struct *vma)
 {
     struct mm_struct *vm_mm = BPF_CORE_READ(vma, vm_mm);
     if (vm_mm == NULL)
@@ -138,7 +150,7 @@ statfunc bool vma_is_stack(struct vm_area_struct *vma)
     return false;
 }
 
-statfunc bool vma_is_heap(struct vm_area_struct *vma)
+statfunc bool vma_is_initial_heap(struct vm_area_struct *vma)
 {
     struct mm_struct *vm_mm = BPF_CORE_READ(vma, vm_mm);
     if (vm_mm == NULL)
@@ -158,7 +170,46 @@ statfunc bool vma_is_heap(struct vm_area_struct *vma)
 
 statfunc bool vma_is_anon(struct vm_area_struct *vma)
 {
-    return BPF_CORE_READ(vma, vm_file) == NULL;
+    return !vma_is_file_backed(vma);
+}
+
+// The golang heap consists of arenas which are memory regions mapped using mmap.
+// When allocating areans, golang supplies mmap with an address hint, which is an
+// address that the kernel should place the mapping at.
+// Hints are constant and vary between architectures, see `mallocinit()` in
+// https://github.com/golang/go/blob/master/src/runtime/malloc.go
+// From observation, when allocating arenas the MAP_FIXED flag is used which forces
+// the kernel to use the specified address or fail the mapping, so it is safe to
+// rely on the address pattern to determine if it belongs to a heap arena.
+#define GOLANG_ARENA_HINT_MASK 0x80ff00000000UL
+#if defined(bpf_target_x86)
+    #define GOLANG_ARENA_HINT (0xc0UL << 32)
+#elif defined(bpf_target_arm64)
+    #define GOLANG_ARENA_HINT (0x40UL << 32)
+#else
+    #error Unsupported architecture
+#endif
+
+statfunc bool vma_is_golang_heap(struct vm_area_struct *vma)
+{
+    u64 vm_start = BPF_CORE_READ(vma, vm_start);
+
+    return (vm_start & GOLANG_ARENA_HINT_MASK) == GOLANG_ARENA_HINT;
+}
+
+statfunc bool vma_is_thread_stack(struct task_struct *task, struct vm_area_struct *vma)
+{
+    // Look up the stack VMA for this task
+    pid_t pid = BPF_CORE_READ(task, pid);
+    address_range_t *stack = bpf_map_lookup_elem(&thread_stacks, &pid);
+    if (stack == NULL)
+        // This thread's stack isn't tracked
+        return false;
+
+    // Check if the VMA is **contained** in the thread stack range.
+    // We don't check exact address range match because a change to the permissions
+    // of part of the stack VMA will split it into multiple VMAs.
+    return BPF_CORE_READ(vma, vm_start) >= stack->start && BPF_CORE_READ(vma, vm_end) <= stack->end;
 }
 
 statfunc bool vma_is_vdso(struct vm_area_struct *vma)
@@ -174,19 +225,33 @@ statfunc bool vma_is_vdso(struct vm_area_struct *vma)
     return strncmp("[vdso]", mapping_name, 7) == 0;
 }
 
-statfunc enum vma_type get_vma_type(struct vm_area_struct *vma)
+statfunc enum vma_type get_vma_type(struct task_struct *task, struct vm_area_struct *vma)
 {
-    if (vma_is_stack(vma))
+    // The check order is a balance between how expensive the check is and how likely it is to pass
+
+    if (vma_is_file_backed(vma))
+        return VMA_FILE_BACKED;
+
+    if (vma_is_initial_stack(vma))
         return VMA_STACK;
 
-    if (vma_is_heap(vma))
+    if (vma_is_initial_heap(vma))
         return VMA_HEAP;
 
-    if (vma_is_anon(vma) && !vma_is_vdso(vma)) {
+    if (vma_is_anon(vma)) {
+        if (vma_is_golang_heap(vma))
+            return VMA_GOLANG_HEAP;
+
+        if (vma_is_thread_stack(task, vma))
+            return VMA_THREAD_STACK;
+
+        if (vma_is_vdso(vma))
+            return VMA_VDSO;
+
         return VMA_ANON;
     }
 
-    return VMA_OTHER;
+    return VMA_UNKNOWN;
 }
 
 #endif

pkg/ebpf/c/maps.h

@@ -395,6 +395,16 @@ struct elf_files_map {
 
 typedef struct elf_files_map elf_files_map_t;
 
+// keep track of thread stacks
+struct thread_stacks {
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, 16384);
+    __type(key, pid_t);
+    __type(value, address_range_t);
+} thread_stacks SEC(".maps");
+
+typedef struct thread_stacks thread_stacks_t;
+
 //
 // versioned maps (map of maps)
 //

pkg/ebpf/c/tracee.bpf.c

@@ -1294,6 +1294,14 @@ int lkm_seeker_new_mod_only_tail(struct pt_regs *ctx)
 SEC("raw_tracepoint/sched_process_exec")
 int tracepoint__sched__sched_process_exec(struct bpf_raw_tracepoint_args *ctx)
 {
+    // Thread stacks map upkeeping
+    pid_t pid = bpf_get_current_pid_tgid();
+    bpf_map_delete_elem(&thread_stacks, &pid);
+    pid_t old_pid = ctx->args[1];
+    if (old_pid != pid)
+        // execve was called from a thread and it inherited the main thread's PID, remove the old PID as well
+        bpf_map_delete_elem(&thread_stacks, &old_pid);
+
     program_data_t p = {};
     if (!init_program_data(&p, ctx, SCHED_PROCESS_EXEC))
         return 0;
@@ -1432,6 +1440,10 @@ int sched_process_exec_event_submit_tail(struct bpf_raw_tracepoint_args *ctx)
 SEC("raw_tracepoint/sched_process_exit")
 int tracepoint__sched__sched_process_exit(struct bpf_raw_tracepoint_args *ctx)
 {
+    // Thread stacks map upkeeping
+    pid_t pid = bpf_get_current_pid_tgid();
+    bpf_map_delete_elem(&thread_stacks, &pid);
+
     program_data_t p = {};
     if (!init_program_data(&p, ctx, SCHED_PROCESS_EXIT))
         return 0;
@@ -5184,6 +5196,40 @@ int BPF_KPROBE(trace_chmod_common)
     return events_perf_submit(&p, 0);
 }
 
+// Keep track of new threads' stacks
+SEC("kprobe/wake_up_new_task")
+int BPF_KPROBE(trace_wake_up_new_task)
+{
+    struct task_struct *task = (struct task_struct *) PT_REGS_PARM1(ctx);
+
+    if (get_task_flags(task) & PF_KTHREAD)
+        return 0;
+
+        // Get user SP of new thread
+#if defined(bpf_target_x86)
+    struct fork_frame *fork_frame = (struct fork_frame *) BPF_CORE_READ(task, thread.sp);
+    u64 thread_sp = BPF_CORE_READ(fork_frame, regs.sp);
+#elif defined(bpf_target_arm64)
+    struct pt_regs *thread_regs = (struct pt_regs *) BPF_CORE_READ(task, thread.cpu_context.sp);
+    u64 thread_sp = BPF_CORE_READ(thread_regs, sp);
+#else
+    #error Unsupported architecture
+#endif
+
+    // Find VMA which contains the SP
+    struct vm_area_struct *vma = find_vma(ctx, task, thread_sp);
+    if (unlikely(vma == NULL))
+        return 0;
+
+    // Add the VMA address range to the thread stacks map
+    pid_t pid = BPF_CORE_READ(task, pid);
+    address_range_t range = {.start = BPF_CORE_READ(vma, vm_start),
+                             .end = BPF_CORE_READ(vma, vm_end)};
+    bpf_map_update_elem(&thread_stacks, &pid, &range, BPF_ANY);
+
+    return 0;
+}
+
 //
 // Syscall checkers
 //
@@ -5215,11 +5261,13 @@ statfunc void check_suspicious_syscall_source(void *ctx, struct pt_regs *regs, u
     if (unlikely(vma == NULL))
         return;
 
-    // Get VMA type and make sure it's abnormal (stack/heap/anonymous VMA)
-    enum vma_type vma_type = get_vma_type(vma);
-    if (vma_type == VMA_OTHER)
+    // If the VMA is file-backed, the syscall is determined to be legitimate
+    if (vma_is_file_backed(vma))
         return;
 
+    // Get VMA type
+    enum vma_type vma_type = get_vma_type(task, vma);
+
     // Build a key that identifies the combination of syscall,
     // source VMA and process so we don't submit it multiple times
     syscall_source_key_t key = {.syscall = syscall,
@@ -5237,17 +5285,27 @@ statfunc void check_suspicious_syscall_source(void *ctx, struct pt_regs *regs, u
 
     switch (vma_type) {
         case VMA_STACK:
-            vma_type_str = "stack";
+            vma_type_str = "main stack";
+            break;
+        case VMA_THREAD_STACK:
+            vma_type_str = "thread stack";
             break;
         case VMA_HEAP:
             vma_type_str = "heap";
             break;
+        case VMA_GOLANG_HEAP:
+            // Goroutine stacks are allocated on the golang heap
+            vma_type_str = "golang heap/stack";
+            break;
         case VMA_ANON:
             vma_type_str = "anonymous";
             break;
-        // shouldn't happen
+        case VMA_VDSO:
+            vma_type_str = "vdso";
+            break;
         default:
-            return;
+            vma_type_str = "unknown";
+            break;
     }
 
     unsigned long vma_start = BPF_CORE_READ(vma, vm_start);
@@ -5271,7 +5329,7 @@ int BPF_KPROBE(syscall_checker)
     struct pt_regs *regs = ctx;
     if (get_kconfig(ARCH_HAS_SYSCALL_WRAPPER))
         regs = (struct pt_regs *) PT_REGS_PARM1(ctx);
-    
+
     // Get syscall ID
     u32 syscall = get_syscall_id_from_regs(regs);
 

pkg/ebpf/c/types.h

@@ -577,4 +577,9 @@ typedef struct {
     u64 vma_addr;
 } syscall_source_key_t;
 
+typedef struct {
+    u64 start;
+    u64 end;
+} address_range_t;
+
 #endif

pkg/ebpf/c/vmlinux.h

@@ -256,6 +256,28 @@ typedef struct {
     uid_t val;
 } kuid_t;
 
+#if defined(__TARGET_ARCH_x86)
+
+struct thread_struct {
+    unsigned long sp;
+};
+
+struct fork_frame {
+    struct pt_regs regs;
+};
+
+#elif defined(__TARGET_ARCH_arm64)
+
+struct cpu_context {
+    unsigned long sp;
+};
+
+struct thread_struct {
+    struct cpu_context cpu_context;
+};
+
+#endif
+
 struct task_struct {
     struct thread_info thread_info;
     unsigned int flags;
@@ -278,6 +300,7 @@ struct task_struct {
     struct signal_struct *signal;
     void *stack;
     struct sighand_struct *sighand;
+    struct thread_struct thread;
 };
 
 typedef struct {

pkg/ebpf/probes/probe_group.go

@@ -257,6 +257,7 @@ func NewDefaultProbeGroup(module *bpf.Module, netEnabled bool) (*ProbeGroup, err
 		Dup3:                       NewTraceProbe(SyscallEnter, "dup3", "trace_dup3"),
 		Dup3Ret:                    NewTraceProbe(SyscallExit, "dup3", "trace_ret_dup3"),
 		ChmodCommon:                NewTraceProbe(KProbe, "chmod_common", "trace_chmod_common"),
+		WakeUpNewTask:              NewTraceProbe(KProbe, "wake_up_new_task", "trace_wake_up_new_task"),
 
 		TestUnavailableHook: NewTraceProbe(KProbe, "non_existing_func", "empty_kprobe"),
 		ExecTest:            NewTraceProbe(RawTracepoint, "raw_syscalls:sched_process_exec", "tracepoint__exec_test"),

pkg/ebpf/probes/probes.go

@@ -162,6 +162,7 @@ const (
 	Dup3
 	Dup3Ret
 	ChmodCommon
+	WakeUpNewTask
 )
 
 // Test probe handles

pkg/events/core.go

@@ -13065,7 +13065,14 @@ var CoreEvents = map[ID]Definition{
 		id:      SuspiciousSyscallSource,
 		id32Bit: Sys32Undefined,
 		name:    "suspicious_syscall_source",
-		sets:    []string{},
+		dependencies: Dependencies{
+			probes: []Probe{
+				{handle: probes.WakeUpNewTask, required: false},    // for thread stack tracking
+				{handle: probes.SchedProcessExec, required: false}, // for thread stack tracking
+				{handle: probes.SchedProcessExit, required: false}, // for thread stack tracking
+			},
+		},
+		sets: []string{},
 		fields: []trace.ArgMeta{
 			{Type: "int", Name: "syscall"},
 			{Type: "void*", Name: "ip"},