Merge branch 'richiejp-cgroup'

Author: dkorunic

.editorconfig

@@ -0,0 +1,24 @@
+# EditorConfig file
+# https://editorconfig.org/
+
+# C settings
+[*.c]
+indent_style = space
+indent_size = 2
+tab_width = 2
+end_of_line = lf
+insert_final_newline = true
+trim_trailing_whitespace = true
+max_line_length = 80
+charset = utf-8
+
+# Go settings
+[*.go]
+indent_style = tab
+indent_size = 2
+tab_width = 2
+end_of_line = lf
+insert_final_newline = true
+trim_trailing_whitespace = true
+max_line_length = 80
+charset = utf-8

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"
 
@@ -45,6 +45,7 @@
 
 #define OK 1
 #define NOK 0
+#define ALLOW_PKT 1
 
 // Map key struct for IP traffic
 typedef struct statkey_t {
@@ -237,6 +238,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;
   }
 
@@ -265,6 +267,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;
   }
 
@@ -281,6 +284,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.
  *
@@ -352,6 +412,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
+}