Add cgroup v1/v2 detection and fix multi-CPU map update races

Introduce a `counter_cfg` eBPF map to pass the cgroup filesystem magic
from userland to eBPF programs, enabling `get_current_cgroup_id()` to
select the correct cgroup ID retrieval path (v2 via bpf helper vs. v1
via task_struct walk). Detect cgroup v1, v2 and hybrid mode in
`getCgroupFsMagic()` and push the magic value at startup for KProbes
and CGroup SKB modes.

Consolidate `get_cgroupid()` into `counter_common.h` (renamed to
`get_cgroup_id()`) so it is shared across all BPF programs; update
`cgroup.bpf.c` to include `counter_common.h` instead of `cgroup.h`.

Fix a multi-CPU TOCTOU race in `update_val()` and the cgroup_skb path:
when `BPF_NOEXIST` insert fails (concurrent insert on another CPU),
retry the lookup and increment atomically with `__sync_fetch_and_add`
to avoid silently dropping packets.

Scope interface resolution to only TC and XDP modes, avoiding a
spurious fatal error when running in KProbes or CGroup mode without a
meaningful `--iface` argument.

Fix `icmpv6_rcv` kprobe to guard `update_val()` with `msglen > 0`,
matching the convention used by all other probes.

Fix cgroup path mapping to emit "/" instead of an empty string when the
cgroup root itself is the target.

Add a TUI race fix: check the `done` channel after `processMap` returns
to avoid queuing a draw after the app has been stopped.

Author: dkorunic

bpf/cgroup.bpf.c

@@ -28,7 +28,7 @@
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
 
-#include "cgroup.h"
+#include "counter_common.h"
 
 typedef struct cgroup_event_t {
   char path[PATH_MAX]; // cgroup path
@@ -50,52 +50,6 @@ struct {
   __type(value, __u32);
 } perf_cgroup_event SEC(".maps");
 
-/**
- * get_cgroupid - reads the cgroup ID from the given struct cgroup
- *
- * This function reads the cgroup ID from the given struct cgroup and returns
- * it. The function works on kernels v4.10 and above.
- *
- * @cgrp: the struct cgroup to read the cgroup ID from
- *
- * Returns: the cgroup ID as an unsigned 64-bit integer
- *
- * get_cgroupid() comes from aquasecurity/tracee, license: Apache-2.0
- */
-static inline __attribute__((always_inline)) __u64
-get_cgroupid(struct cgroup *cgrp) {
-  struct kernfs_node *kn = BPF_CORE_READ(cgrp, kn);
-
-  if (kn == NULL)
-    return 0;
-
-  __u64 id; // was union kernfs_node_id before 5.5, can read it as u64 in both
-            // situations
-
-  if (bpf_core_type_exists(union kernfs_node_id)) {
-    struct kernfs_node___older_v55 *kn_old = (void *)kn;
-    struct kernfs_node___rh8 *kn_rh8 = (void *)kn;
-
-    if (bpf_core_field_exists(kn_rh8->id)) {
-      // RHEL8 has both types declared: union and u64:
-      //     kn->id
-      //     rh->rh_kabi_hidden_172->id
-      // pointing to the same data
-      bpf_core_read(&id, sizeof(__u64), &kn_rh8->id);
-      id = id & 0xffffffff; // XXX: u32 is required
-    } else {
-      // all other regular kernels below v5.5
-      bpf_core_read(&id, sizeof(__u64), &kn_old->id);
-      id = id & 0xffffffff; // XXX: u32 is required
-    }
-  } else {
-    // kernel v5.5 and above
-    bpf_core_read(&id, sizeof(__u64), &kn->id);
-  }
-
-  return id;
-}
-
 /**
  * trace_cgroup_mkdir traces the creation of a new cgroup directory.
  *
@@ -116,7 +70,7 @@ int trace_cgroup_mkdir(struct bpf_raw_tracepoint_args *ctx) {
   struct cgroup *dst_cgrp = (struct cgroup *)ctx->args[0];
   char *path = (char *)ctx->args[1];
 
-  __u64 cgroupid = get_cgroupid(dst_cgrp);
+  __u64 cgroupid = get_cgroup_id(dst_cgrp);
   __u32 zero_key = 0;
 
   cgroupevent *val =

bpf/cgroup.h

@@ -21,8 +21,11 @@
 
 //go:build ignore
 
-#define MAX_ENTRIES 1024
+#pragma once
+
 #define PATH_MAX 4096
+#define CGROUP_FSMAGIC 0x27e0eb    // cgroup v1
+#define CGROUP2_FSMAGIC 0x63677270 // cgroup v2
 
 union kernfs_node_id {
   struct {

bpf/cgroup_skb.bpf.c

@@ -94,7 +94,7 @@ process_cgroup_skb(struct __sk_buff *skb) {
     __builtin_memcpy(key.comm, ski->comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   statvalue *val = (statvalue *)bpf_map_lookup_elem(&pkt_count, &key);
   if (val) {
@@ -104,7 +104,16 @@ process_cgroup_skb(struct __sk_buff *skb) {
   } else {
     statvalue initval = {.packets = 1, .bytes = pkt_len};
 
-    bpf_map_update_elem(&pkt_count, &key, &initval, BPF_NOEXIST);
+    // BPF_NOEXIST can race on multi-CPU: another CPU may insert the same key
+    // between our lookup and this insert. On failure, retry the lookup and
+    // increment atomically so the packet is not silently dropped.
+    if (bpf_map_update_elem(&pkt_count, &key, &initval, BPF_NOEXIST) != 0) {
+      val = (statvalue *)bpf_map_lookup_elem(&pkt_count, &key);
+      if (val) {
+        __sync_fetch_and_add(&val->packets, 1);
+        __sync_fetch_and_add(&val->bytes, pkt_len);
+      }
+    }
   }
 }
 

bpf/counter.h

@@ -21,6 +21,8 @@
 
 //go:build ignore
 
+#pragma once
+
 #define s6_addr in6_u.u6_addr8
 #define s6_addr16 in6_u.u6_addr16
 #define s6_addr32 in6_u.u6_addr32

bpf/counter_common.h

@@ -19,11 +19,14 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
+//go:build ignore
+
 #pragma once
 
+#include "cgroup.h"
 #include "counter.h"
 
-// Map key struct for IP traffic
+// Counter map key struct for IP traffic
 typedef struct statkey_t {
   struct in6_addr srcip;    // source IPv6 address
   struct in6_addr dstip;    // destination IPv6 address
@@ -35,25 +38,39 @@ typedef struct statkey_t {
   __u8 proto;               // transport protocol
 } statkey;
 
-// Map value struct with counters
+// Counter map value struct with counters
 typedef struct statvalue_t {
   __u64 packets; // packets ingress + egress
   __u64 bytes;   // bytes ingress + egress
 } statvalue;
 
-// Map definition
+// Counter map definition
 struct {
   __uint(type, BPF_MAP_TYPE_LRU_HASH); // LRU hash requires 4.10 kernel
   __uint(max_entries, MAX_ENTRIES);
   __type(key, statkey);
   __type(value, statvalue);
 } pkt_count SEC(".maps");
 
+// Sockinfo struct
 typedef struct sockinfo_t {
   __u8 comm[TASK_COMM_LEN];
   pid_t pid;
 } sockinfo;
 
+// Configuration map value struct
+typedef struct counter_cfg_t {
+  __u64 cgrpfs_magic; // cgroupv1 or cgroupv2 fs magic
+} counter_cfg_value;
+
+// Configuration map definition
+struct {
+  __uint(type, BPF_MAP_TYPE_ARRAY);
+  __uint(max_entries, 1);
+  __type(key, __u32);
+  __type(value, counter_cfg_value);
+} counter_cfg SEC(".maps");
+
 // IPv4-mapped IPv6 address prefix (for V4MAPPED conversion)
 static const __u8 ip4in6[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff};
 
@@ -271,7 +288,16 @@ static inline __attribute__((always_inline)) void update_val(statkey *key,
   } else {
     statvalue initval = {.packets = 1, .bytes = size};
 
-    bpf_map_update_elem(&pkt_count, key, &initval, BPF_NOEXIST);
+    // BPF_NOEXIST can race on multi-CPU: another CPU may insert the same key
+    // between our lookup and this insert. On failure, retry the lookup and
+    // increment atomically so the packet is not silently dropped.
+    if (bpf_map_update_elem(&pkt_count, key, &initval, BPF_NOEXIST) != 0) {
+      val = (statvalue *)bpf_map_lookup_elem(&pkt_count, key);
+      if (val) {
+        __sync_fetch_and_add(&val->packets, 1);
+        __sync_fetch_and_add(&val->bytes, size);
+      }
+    }
   }
 }
 
@@ -641,3 +667,84 @@ process_udp_send(struct sk_buff *skb, statkey *key, pid_t pid) {
 
   return msglen;
 }
+
+/**
+ * get_cgroupid - reads the cgroup ID from the given struct cgroup
+ *
+ * This function reads the cgroup ID from the given struct cgroup and returns
+ * it. The function works on kernels v4.10 and above.
+ *
+ * @cgrp: the struct cgroup to read the cgroup ID from
+ *
+ * Returns: the cgroup ID as an unsigned 64-bit integer
+ *
+ * get_cgroupid() comes from aquasecurity/tracee, license: Apache-2.0
+ */
+static inline __attribute__((always_inline)) __u64
+get_cgroup_id(struct cgroup *cgrp) {
+  struct kernfs_node *kn = BPF_CORE_READ(cgrp, kn);
+
+  if (kn == NULL)
+    return 0;
+
+  __u64 id; // was union kernfs_node_id before 5.5, can read it as u64 in both
+            // situations
+
+  if (bpf_core_type_exists(union kernfs_node_id)) {
+    struct kernfs_node___older_v55 *kn_old = (void *)kn;
+    struct kernfs_node___rh8 *kn_rh8 = (void *)kn;
+
+    if (bpf_core_field_exists(kn_rh8->id)) {
+      // RHEL8 has both types declared: union and u64:
+      //     kn->id
+      //     rh->rh_kabi_hidden_172->id
+      // pointing to the same data
+      bpf_core_read(&id, sizeof(__u64), &kn_rh8->id);
+      id = id & 0xffffffff; // XXX: u32 is required
+    } else {
+      // all other regular kernels below v5.5
+      bpf_core_read(&id, sizeof(__u64), &kn_old->id);
+      id = id & 0xffffffff; // XXX: u32 is required
+    }
+  } else {
+    // kernel v5.5 and above
+    bpf_core_read(&id, sizeof(__u64), &kn->id);
+  }
+
+  return id;
+}
+
+/**
+ * get_current_cgroup_id - get the current cgroup ID
+ *
+ * This function determines the current cgroup ID of the running process.
+ * It first checks if the cgroup v2 pseudo-filesystem is present and
+ * if so, calls the bpf_get_current_cgroup_id() helper function. If not,
+ * it reads the cgroup ID from the task_struct structure. It returns the
+ * cgroup ID as an unsigned 64-bit integer.
+ *
+ * @return the current cgroup ID as an unsigned 64-bit integer
+ */
+static inline __attribute__((always_inline)) __u64 get_current_cgroup_id(void) {
+  __u32 zero_key = 0;
+  __u64 cgrpfs_magic = 0;
+
+  counter_cfg_value *cfg =
+      (counter_cfg_value *)bpf_map_lookup_elem(&counter_cfg, &zero_key);
+  if (cfg) {
+    cgrpfs_magic = cfg->cgrpfs_magic;
+  }
+
+  // cgroup v2
+  if (bpf_core_enum_value_exists(enum bpf_func_id,
+                                 BPF_FUNC_get_current_cgroup_id) &&
+      cgrpfs_magic == CGROUP2_FSMAGIC) {
+    return bpf_get_current_cgroup_id();
+  }
+
+  // cgroup v1
+  struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+  struct cgroup *cgroup = BPF_CORE_READ(task, cgroups, subsys[0], cgroup);
+
+  return get_cgroup_id(cgroup);
+}

bpf/kprobe.bpf.c

@@ -55,7 +55,7 @@ int BPF_KPROBE(tcp_sendmsg, struct sock *sk, struct msghdr *msg, size_t size) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   if (!process_tcp(false, sk, &key, pid))
     return 0;
@@ -92,7 +92,7 @@ int BPF_KPROBE(tcp_cleanup_rbuf, struct sock *sk, int copied) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   if (!process_tcp(true, sk, &key, pid))
     return 0;
@@ -131,7 +131,7 @@ int BPF_KPROBE(ip_send_skb, struct net *net, struct sk_buff *skb) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   size_t msglen = process_udp_send(skb, &key, pid);
   if (msglen > 0)
@@ -171,7 +171,7 @@ int BPF_KPROBE(ip6_send_skb, struct sk_buff *skb) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   size_t msglen = process_udp_send(skb, &key, pid);
   if (msglen > 0)
@@ -209,7 +209,7 @@ int BPF_KPROBE(skb_consume_udp, struct sock *sk, struct sk_buff *skb, int len) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   if (!process_udp_recv(true, skb, &key, pid))
     return 0;
@@ -246,7 +246,7 @@ int BPF_KPROBE(__icmp_send, struct sk_buff *skb, __u8 type, __u8 code,
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   /* process_icmp4 populates src/dst IPs from the triggering packet's IP
    * header; its returned msglen (triggering payload) is not used here. */
@@ -308,7 +308,7 @@ int BPF_KPROBE(icmp6_send, struct sk_buff *skb, __u8 type, __u8 code,
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   /* process_icmp6 populates src/dst IPs from the triggering packet's IPv6
    * header; its returned msglen (triggering payload_len) is not used here. */
@@ -359,7 +359,7 @@ int BPF_KPROBE(icmp_rcv, struct sk_buff *skb) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   size_t msglen = process_icmp4(skb, &key, pid);
   if (msglen > 0)
@@ -391,10 +391,11 @@ int BPF_KPROBE(icmpv6_rcv, struct sk_buff *skb) {
     bpf_get_current_comm(&key.comm, sizeof(key.comm));
   }
 
-  key.cgroupid = bpf_get_current_cgroup_id();
+  key.cgroupid = get_current_cgroup_id();
 
   size_t msglen = process_icmp6(skb, &key, pid);
-  update_val(&key, msglen);
+  if (msglen > 0)
+    update_val(&key, msglen);
 
   return 0;
 }