Added IPv6 packets counter

Author: suchit07-git

counter.c

@@ -3,14 +3,22 @@
 #include <linux/bpf.h>
 #include <linux/if_ether.h>
 #include <linux/ip.h>
+#include <linux/ipv6.h>
 #include "headers/bpf_helpers.h"
 
 struct {
     __uint(type, BPF_MAP_TYPE_HASH);
     __type(key, __be32);
     __type(value, __u64);
     __uint(max_entries, 1024);
-} pkt_count SEC(".maps");
+} ipv4_counter_map SEC(".maps");
+
+struct {
+    __uint(type, BPF_MAP_TYPE_HASH);
+    __type(key, struct in6_addr);
+    __type(value, __u64);
+    __uint(max_entries, 1024);
+} ipv6_counter_map SEC(".maps");
 
 SEC("xdp")
 int xdp_ip_packet_counter(struct xdp_md* ctx) {
@@ -20,19 +28,30 @@ int xdp_ip_packet_counter(struct xdp_md* ctx) {
     if ((void*)(eth + 1) > data_end) {
         return XDP_PASS;
     }
-    if (eth->h_proto != __constant_htons(ETH_P_IP)) {
-        return XDP_PASS;
-    }
-    struct iphdr *ip = (struct iphdr*)(eth + 1);
-    if ((void*)(ip + 1) > data_end) {
-        return XDP_PASS;
-    }
-    __u64 *count, init_val = 1;
-    count = bpf_map_lookup_elem(&pkt_count, &ip->saddr);
-    if (count) {
-        __sync_fetch_and_add(count, 1);
-    } else {
-        bpf_map_update_elem(&pkt_count, &ip->saddr, &init_val, BPF_ANY);
+    if (eth->h_proto == __constant_htons(ETH_P_IP)) {
+        struct iphdr *ip = (struct iphdr*)(eth + 1);
+        if ((void*)(ip + 1) > data_end) {
+            return XDP_PASS;
+        }
+        __u64 *count, init_val = 1;
+        count = bpf_map_lookup_elem(&ipv4_counter_map, &ip->saddr);
+        if (count) {
+            __sync_fetch_and_add(count, 1);
+        } else {
+            bpf_map_update_elem(&ipv4_counter_map, &ip->saddr, &init_val, BPF_ANY);
+        }
+    } else if (eth->h_proto == __constant_htons(ETH_P_IPV6)) {
+        struct ipv6hdr *ip6 = (struct ipv6hdr*)(eth + 1);
+        if ((void*)(ip6 + 1) > data_end) {
+            return XDP_PASS;
+        }
+        __u64 *count, init_val = 1;
+        count = bpf_map_lookup_elem(&ipv6_counter_map, &ip6->saddr);
+        if (count) {
+            __sync_fetch_and_add(count, 1);
+        } else {
+            bpf_map_update_elem(&ipv6_counter_map, &ip6->saddr, &init_val, BPF_ANY);
+        }
     }
     return XDP_PASS;
 }

counter.o

@@ -16,12 +16,52 @@ import (
 
 const bpfProgramPath = "counter.o"
 
-func ipToString(ip uint32) net.IP {
+func ip4ToString(ip uint32) net.IP {
 	address := make(net.IP, 4)
 	binary.LittleEndian.PutUint32(address, ip)
 	return address
 }
 
+func ip6ToString(ip [16]byte) string {
+	return net.IP(ip[:]).String()
+}
+
+func printMap(m *ebpf.Map, ipToString func(interface{}) string) {
+	iter := m.Iterate()
+	var ip interface{}
+	var count uint64
+	for iter.Next(&ip, &count) {
+		fmt.Printf("%s: %d packets\n", ipToString(ip), count)
+	}
+	if err := iter.Err(); err != nil {
+		fmt.Printf("Error iterating map: %v\n", err)
+	}
+}
+
+func printIPv4Map(m *ebpf.Map) {
+	iter := m.Iterate()
+	var ip uint32
+	var count uint64
+	for iter.Next(&ip, &count) {
+		fmt.Printf("  %s: %d packets\n", ip4ToString(ip), count)
+	}
+	if err := iter.Err(); err != nil {
+		fmt.Printf("Error iterating map: %v\n", err)
+	}
+}
+
+func printIPv6Map(m *ebpf.Map) {
+	iter := m.Iterate()
+	var ip [16]byte
+	var count uint64
+	for iter.Next(&ip, &count) {
+		fmt.Printf("  %s: %d packets\n", ip6ToString(ip), count)
+	}
+	if err := iter.Err(); err != nil {
+		fmt.Printf("Error iterating map: %v\n", err)
+	}
+}
+
 func main() {
 	if len(os.Args) < 3 || os.Args[1] != "-count" {
 		fmt.Println("Usage: ./ebpf-packet-filter -count <interface_name>")
@@ -39,9 +79,10 @@ func main() {
 		log.Fatal("Failed to create BPF collection:", err)
 	}
 	defer collection.Close()
-	ipMap := collection.Maps["pkt_count"]
-	if ipMap == nil {
-		log.Fatal("Failed to find pkt_count map in the BPF program")
+	ip4Map := collection.Maps["ipv4_counter_map"]
+	ip6Map := collection.Maps["ipv6_counter_map"]
+	if ip4Map == nil || ip6Map == nil {
+		log.Fatal("Failed to find required map in the BPF program")
 	}
 	ifname := os.Args[2]
 	iface, err := net.InterfaceByName(ifname)
@@ -68,15 +109,10 @@ func main() {
 	for {
 		select {
 		case <-tick:
-			iter := ipMap.Iterate()
-			var ip uint32
-			var count uint64
-			for iter.Next(&ip, &count) {
-				log.Printf("%s: %d packets\n", ipToString(ip), count)
-			}
-			if err := iter.Err(); err != nil {
-				log.Fatal("Error iterating map:", err)
-			}
+			fmt.Println("IPv4 Packets count:")
+			printIPv4Map(ip4Map)
+			fmt.Println("IPv6 Packets count:")
+			printIPv6Map(ip6Map)
 		case <-stop:
 			log.Println("Received signal, exiting...")
 			return