Allow monitoring bandwidth usage of CGroups

This adds the CGroup V2 ingress and egress hook points. The user can attach
to any given CGroup and see bandwidth statistics just for that CGroup.
This also has the advantage of seeing VPN traffic.

Attaching to the root CGroup shows statics for the entire system.

Author: richiejp

counter.c

@@ -19,7 +19,7 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-// go:build ignore
+//go:build ignore
 
 #include "vmlinux.h"
 
@@ -44,6 +44,7 @@
 
 #define OK 1
 #define NOK 0
+#define ALLOW_PKT 1
 
 // Map key struct for IP traffic
 typedef struct statkey_t {
@@ -236,6 +237,7 @@ static inline void process_eth(void *data, void *data_end, __u64 pkt_len) {
 
   // validate Ethernet size
   if ((void *)eth + sizeof(*eth) > data_end) {
+    bpf_printk("size validation failure");
     return;
   }
 
@@ -264,6 +266,7 @@ static inline void process_eth(void *data, void *data_end, __u64 pkt_len) {
     break;
   }
   default:
+      bpf_printk("wrong packet type: %d", eth->h_proto);
     return;
   }
 
@@ -280,6 +283,63 @@ static inline void process_eth(void *data, void *data_end, __u64 pkt_len) {
   }
 }
 
+/**
+ * Process SKB as it is seen by the cgroup, which is without the ethernet
+ * headers
+ *
+ * @param skb The CGroup skb
+ *
+ * @return none
+ *
+ * @throws none
+ */
+static inline void process_cgroup_skb(struct __sk_buff *skb) {
+  void *data = (void *)(long)skb->data;
+  void *data_end = (void *)(long)skb->data_end;
+  __u64 pkt_len = skb->len;
+
+  // initialize key
+  statkey key;
+  __builtin_memset(&key, 0, sizeof(key));
+
+  switch (bpf_ntohs(skb->protocol)) {
+  case ETH_P_IP: {
+    struct iphdr *ip4 = data;
+
+    if (process_ip4(ip4, data_end, &key) == NOK) {
+      return;
+    }
+
+    break;
+  }
+  case ETH_P_IPV6: {
+    struct ipv6hdr *ip6 = data;
+
+    if (process_ip6(ip6, data_end, &key) == NOK) {
+      return;
+    }
+
+    break;
+  }
+  default:
+      bpf_printk("wrong packet type: %d", skb->protocol);
+    return;
+
+  }
+
+  // lookup value in hash
+  statvalue *val = (statvalue *)bpf_map_lookup_elem(&pkt_count, &key);
+  if (val) {
+    // atomic XADD, doesn't need bpf_spin_lock()
+    __sync_fetch_and_add(&val->packets, 1);
+    __sync_fetch_and_add(&val->bytes, pkt_len);
+  } else {
+    statvalue initval = {.packets = 1, .bytes = pkt_len};
+
+    bpf_map_update_elem(&pkt_count, &key, &initval, BPF_NOEXIST);
+  }
+}
+
 /**
  * Process the packet for traffic control and take necessary actions.
  *
@@ -351,6 +411,20 @@ int tc_count_packets(struct __sk_buff *skb) {
   return TC_ACT_UNSPEC;
 }
 
+SEC("cgroup_skb/ingress")
+int cgroup_skb_ingress(struct __sk_buff *skb) {
+  process_cgroup_skb(skb);
+
+  return ALLOW_PKT;
+}
+
+SEC("cgroup_skb/egress")
+int cgroup_skb_egress(struct __sk_buff *skb) {
+  process_cgroup_skb(skb);
+
+  return ALLOW_PKT;
+}
+
 /**
  * Process TCP socket information and populate the key structure with
  * extracted data.

counter_arm64_bpfel.go

@@ -40,7 +40,7 @@ const (
 )
 
 var (
-	ifname, xdpMode                                          *string
+	ifname, xdpMode, useCGroup                               *string
 	jsonOutput, version, help, useXDP, useKProbes, enableTUI *bool
 	timeout, refresh                                         *time.Duration
 	xdpAttachFlags                                           link.XDPAttachFlags
@@ -51,6 +51,7 @@ func parseFags() {
 
 	help = fs.Bool('?', "help", "display help")
 	jsonOutput = fs.Bool('j', "json", "if true, output in JSON format")
+	useCGroup = fs.String('c', "cgroup", "", "the path to a CGroup V2 to measure statistics on")
 	useXDP = fs.Bool('x', "xdp", "if true, use XDP instead of TC (this disables egress statistics)")
 	useKProbes = fs.Bool('k', "kprobes", "if true, use KProbes for per-proces TCP/UDP statistics")
 	enableTUI = fs.Bool('g', "tui", "if true, enable TUI")

counter_arm64_bpfel.o

@@ -68,6 +68,8 @@ func main() {
 	}()
 
 	switch {
+	case *useCGroup != "":
+		links = startCgroup(objs, *useCGroup, links)
 	// KProbes w/ PID tracking
 	case *useKProbes:
 		hooks := []kprobeHook{
@@ -289,3 +291,40 @@ func startTC(objs counterObjects, iface *net.Interface, links []link.Link) []lin
 
 	return links
 }
+
+func startCgroup(objs counterObjects, cgroupPath string, links []link.Link) []link.Link {
+	var l link.Link
+
+	err := features.HaveProgramType(ebpf.CGroupSKB)
+	if errors.Is(err, ebpf.ErrNotSupported) {
+		log.Fatalf("CgroupSKB not supported on this kernel")
+	}
+
+	if err != nil {
+		log.Fatalf("Error checking CGroupSKB support: %v", err)
+	}
+
+	l, err = link.AttachCgroup(link.CgroupOptions{
+		Program: objs.CgroupSkbIngress,
+		Attach:  ebpf.AttachCGroupInetIngress,
+		Path:    cgroupPath,
+	})
+	if err != nil {
+		log.Fatalf("Error attaching CgroupSkbIngress to %s: %v", cgroupPath, err)
+	}
+
+	l, err = link.AttachCgroup(link.CgroupOptions{
+		Program: objs.CgroupSkbEgress,
+		Attach:  ebpf.AttachCGroupInetEgress,
+		Path:    cgroupPath,
+	})
+	if err != nil {
+		log.Fatalf("Error attaching CgroupSkbEgress to %s: %v", cgroupPath, err)
+	}
+
+	links = append(links, l)
+
+	log.Printf("Starting on CGroup %s", cgroupPath)
+
+	return links
+}