epbf: Propagate PID namespace to kernel space

ebpf/bpf/pid.h

@@ -1,47 +1,42 @@
-
-#if !defined(PYROSCOPE_PID)
-#define PYROSCOPE_PID
-
-// this should not be used in production, and always be disabled
-// but is useful for running in a privileged context outside host pid namespace, for example wsl2
-//#define PYROSCOPE_PID_NAMESPACED
-
-#if defined(PYROSCOPE_PID_NAMESPACED)
-
-#include "bpf_core_read.h"
-// https://github.com/grafana/beyla/blob/6366275ce2d2c9bdefd47975b389fbcf39cbbea8/bpf/pid.h#L13
-// Good resource on this: https://mozillazg.com/2022/05/ebpf-libbpfgo-get-process-info-en.html
-// Using bpf_get_ns_current_pid_tgid is too restrictive for us
-//static __always_inline void ns_pid_ppid(struct task_struct *task, u32 *pid , int *ppid, u32 *pid_ns_id) {
-static __always_inline void current_pid(u32 *pid) {
-    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
-    if (task == 0) {
-        return;
-    }
-    struct upid upid;
-
-    unsigned int level = BPF_CORE_READ(task, nsproxy, pid_ns_for_children, level);
-    struct pid *ns_pid = (struct pid *)BPF_CORE_READ(task, group_leader, thread_pid);
-    bpf_probe_read_kernel(&upid, sizeof(upid), &ns_pid->numbers[level]);
-
-    *pid = (u32)upid.nr;
-//    unsigned int p_level = BPF_CORE_READ(task, real_parent, nsproxy, pid_ns_for_children, level);
-//
-//    struct pid *ns_ppid = (struct pid *)BPF_CORE_READ(task, real_parent, group_leader, thread_pid);
-//    bpf_probe_read_kernel(&upid, sizeof(upid), &ns_ppid->numbers[p_level]);
-//    *ppid = upid.nr;
-//
-//    struct ns_common ns = BPF_CORE_READ(task, nsproxy, pid_ns_for_children, ns);
-//    *pid_ns_id = ns.inum;
-}
-
-#else // PYROSCOPE_PID_NAMESPACED
-
-static __always_inline void current_pid(u32 *pid) {
-  u64 pid_tgid = bpf_get_current_pid_tgid();
-  *pid = (u32)(pid_tgid >> 32);
-}
-#endif // PYROSCOPE_PID_NAMESPACED
-
-
-#endif // PYROSCOPE_PID
+#if !defined(PYROSCOPE_PID)
+#define PYROSCOPE_PID
+
+#include "bpf_core_read.h"
+#include "bpf_helpers.h"
+#include "vmlinux.h"
+
+#define PID_NESTED_NAMESPACES_MAX 4
+
+static __always_inline void current_pid(uint64_t ns_pid_ino, uint32_t *pid) {
+    struct upid upid;
+    unsigned int inum;
+
+    // fallback to host pid, if no inode provided
+    if (ns_pid_ino == 0) {
+        uint64_t pid_tgid = bpf_get_current_pid_tgid();
+        *pid = (u32)(pid_tgid >> 32);
+        return;
+    }
+
+    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+
+    // retrieve level nested namespaces
+    unsigned int level = BPF_CORE_READ(task, nsproxy, pid_ns_for_children, level);
+
+    // match the level with pid ns inode
+#pragma unroll
+    for (int i = 0; i < PID_NESTED_NAMESPACES_MAX; i++) {
+        if ((level - i) < 0) {
+            break;
+        }
+        upid = BPF_CORE_READ(task, thread_pid, numbers[level - i]);
+        inum = BPF_CORE_READ(upid.ns, ns.inum);
+
+        if (inum == ns_pid_ino) {
+            *pid = upid.nr;
+            break;
+        }
+    }
+}
+
+#endif // PYROSCOPE_PID

ebpf/bpf/profile.bpf.c

@@ -9,10 +9,16 @@
 
 #define PF_KTHREAD 0x00200000
 
+struct global_config_t {
+    uint64_t ns_pid_ino;
+};
+
+const volatile struct global_config_t global_config;
+
 SEC("perf_event")
 int do_perf_event(struct bpf_perf_event_data *ctx) {
     u32 tgid = 0;
-    current_pid(&tgid);
+    current_pid(global_config.ns_pid_ino, &tgid);
 
     struct sample_key key = {};
     u32 *val, one = 1;
@@ -88,7 +94,7 @@ int BPF_KPROBE(disassociate_ctty, int on_exit) {
         return 0;
     }
     u32 pid = 0;
-    current_pid(&pid);
+    current_pid(global_config.ns_pid_ino, &pid);
     if (pid == 0) {
         return 0;
     }
@@ -104,7 +110,7 @@ int BPF_KPROBE(disassociate_ctty, int on_exit) {
 SEC("kprobe/exec")
 int BPF_KPROBE(exec, void *_) {
     u32 pid = 0;
-    current_pid(&pid);
+    current_pid(global_config.ns_pid_ino, &pid);
     if (pid == 0) {
         return 0;
     }

ebpf/bpf/pyperf.bpf.c

@@ -40,6 +40,7 @@ enum {
 struct global_config_t {
     uint8_t bpf_log_err;
     uint8_t bpf_log_debug;
+    uint64_t ns_pid_ino;
 };
 
 const volatile struct global_config_t global_config;
@@ -278,7 +279,7 @@ static __always_inline int pyperf_collect_impl(struct bpf_perf_event_data* ctx,
 SEC("perf_event")
 int pyperf_collect(struct bpf_perf_event_data *ctx) {
     u32 pid;
-    current_pid(&pid);
+    current_pid(global_config.ns_pid_ino, &pid);
     if (pid == 0) {
         return 0;
     }

ebpf/pyrobpf/gen.go

@@ -1,4 +1,4 @@
 package pyrobpf
 
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type pid_event -target amd64 -cc clang -cflags "-O2 -Wall -Werror -fpie -Wno-unused-variable -Wno-unused-function" Profile ../bpf/profile.bpf.c -- -I../bpf/libbpf -I../bpf/vmlinux/
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type pid_event -target arm64 -cc clang -cflags "-O2 -Wall -Werror -fpie -Wno-unused-variable -Wno-unused-function" Profile ../bpf/profile.bpf.c -- -I../bpf/libbpf -I../bpf/vmlinux/
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type global_config_t -type pid_event -target amd64 -cc clang -cflags "-O2 -Wall -Werror -fpie -Wno-unused-variable -Wno-unused-function" Profile ../bpf/profile.bpf.c -- -I../bpf/libbpf -I../bpf/vmlinux/
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -type global_config_t -type pid_event -target arm64 -cc clang -cflags "-O2 -Wall -Werror -fpie -Wno-unused-variable -Wno-unused-function" Profile ../bpf/profile.bpf.c -- -I../bpf/libbpf -I../bpf/vmlinux/

ebpf/session.go

@@ -15,6 +15,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"syscall"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/btf"
@@ -148,7 +149,25 @@ func (s *session) Start() error {
 	opts := &ebpf.CollectionOptions{
 		Programs: s.progOptions(),
 	}
-	if err := pyrobpf.LoadProfileObjects(&s.bpf, opts); err != nil {
+	spec, err := pyrobpf.LoadProfile()
+	if err != nil {
+		return fmt.Errorf("pyrobpf load %w", err)
+	}
+
+	_, nsIno, err := getPIDNamespace()
+	if err != nil {
+		return fmt.Errorf("unable to get pid namespace %w", err)
+	}
+	err = spec.RewriteConstants(map[string]interface{}{
+		"global_config": pyrobpf.ProfileGlobalConfigT{
+			NsPidIno: nsIno,
+		},
+	})
+	if err != nil {
+		return fmt.Errorf("pyrobpf rewrite constants %w", err)
+	}
+	err = spec.LoadAndAssign(&s.bpf, opts)
+	if err != nil {
 		s.logVerifierError(err)
 		s.stopLocked()
 		return fmt.Errorf("load bpf objects: %w", err)
@@ -905,3 +924,14 @@ func (s *stackBuilder) reset() {
 func (s *stackBuilder) append(sym string) {
 	s.stack = append(s.stack, sym)
 }
+
+func getPIDNamespace() (dev uint64, ino uint64, err error) {
+	stat, err := os.Stat("/proc/self/ns/pid")
+	if err != nil {
+		return 0, 0, err
+	}
+	if st, ok := stat.Sys().(*syscall.Stat_t); ok {
+		return st.Dev, st.Ino, nil
+	}
+	return 0, 0, fmt.Errorf("could not determine pid namespace")
+}

ebpf/session_python.go

@@ -120,10 +120,15 @@ func (s *session) loadPyPerf(cause *sd.Target) (*python.Perf, error) {
 	if err != nil {
 		return nil, fmt.Errorf("pyperf load %w", err)
 	}
+	_, nsIno, err := getPIDNamespace()
+	if err != nil {
+		return nil, fmt.Errorf("unable to get pid namespace %w", err)
+	}
 	err = spec.RewriteConstants(map[string]interface{}{
 		"global_config": python.PerfGlobalConfigT{
 			BpfLogErr:   boolToU8(s.pythonBPFErrorLogEnabled(cause)),
 			BpfLogDebug: boolToU8(s.pythonBPFDebugLogEnabled(cause)),
+			NsPidIno:    nsIno,
 		},
 	})
 	if err != nil {