Add flow-based client stats by default for rate limits and organize/clean code.

.gitignore

@@ -2,3 +2,4 @@
 build/
 xdpfw
 xdpfw.s
+xdpfw.conf

Makefile

@@ -47,7 +47,7 @@ OBJS = $(BUILDDIR)/$(CONFIGOBJ) $(BUILDDIR)/$(CMDLINEOBJ)
 
 # LD flags and includes.
 LDFLAGS += -lconfig -lelf -lz
-INCS = -I $(LIBBPFSRC)
+INCS = -I $(SRCDIR) -I $(LIBBPFSRC)
 INCS += -I /usr/include -I /usr/local/include
 
 # All chain.
@@ -67,8 +67,8 @@ xdpfw_filter:
 # Utils chain.
 utils:
 	mkdir -p $(BUILDDIR)/
-	$(CC) -O2 -c -o $(BUILDDIR)/$(CONFIGOBJ) $(SRCDIR)/$(CONFIGSRC)
+	$(CC) $(INCS) -O2 -c -o $(BUILDDIR)/$(CONFIGOBJ) $(SRCDIR)/$(CONFIGSRC)
-	$(CC) -O2 -c -o $(BUILDDIR)/$(CMDLINEOBJ) $(SRCDIR)/$(CMDLINESRC)
+	$(CC) $(INCS) -O2 -c -o $(BUILDDIR)/$(CMDLINEOBJ) $(SRCDIR)/$(CMDLINESRC)
 
 # LibXDP chain. We need to install objects here since our program relies on installed object files and such.
 libxdp:

README.md

@@ -239,6 +239,13 @@ There is a possibility I may make this firewall stateful in the future *when* I
 
 You may also be interested in this awesome project called [FastNetMon](https://github.com/pavel-odintsov/fastnetmon)!
 
+### Rate Limits
+By default, client stats including packets and bytes per second are calculated per *partial* flow (source IP/port and protocol). This is useful if you want to specify connection-specific rate limits inside of your filtering rules using the `pps` and `bps` settings. However, if you want to calculate client stats using only the source IP, you may comment out [this](https://github.com/gamemann/XDP-Firewall/blob/master/src/xdpfw.h#L25) line.
+
+```C
+//#define USE_FLOW_RL
+```
+
 ## My Other XDP Projects
 I just wanted to share other open source projects I've made which also utilize XDP (or AF_XDP sockets) for those interested. I hope code from these other projects help programmers trying to utilize XDP in their own projects!
 

src/config.c

@@ -6,9 +6,9 @@
 
 #include <arpa/inet.h>
 
-#include "xdpfw.h"
+#include <xdpfw.h>
-#include "config.h"
+#include <config.h>
-#include "utils.h"
+#include <utils.h>
 
 FILE *file;
 

src/xdp/helpers.h

@@ -0,0 +1,8 @@
+#pragma once
+
+#include <linux/bpf.h>
+#include <linux/bpf_common.h>
+
+#include <bpf_helpers.h>
+#include <xdp/xdp_helpers.h>
+#include <xdp/prog_dispatcher.h>

src/xdp/maps.h

@@ -0,0 +1,67 @@
+#pragma once
+
+#include <xdpfw.h>
+
+#include <xdp/helpers.h>
+
+struct 
+{
+    __uint(priority, 10);
+    __uint(XDP_PASS, 1);
+} XDP_RUN_CONFIG(xdp_prog_main);
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
+    __uint(max_entries, MAX_FILTERS);
+    __type(key, __u32);
+    __type(value, struct filter);
+} filters_map SEC(".maps");
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
+    __uint(max_entries, 1);
+    __type(key, __u32);
+    __type(value, struct stats);
+} stats_map SEC(".maps");
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, MAX_TRACK_IPS);
+#ifdef USE_FLOW_RL
+    __type(key, struct flow);
+#else
+    __type(key, __u32);
+#endif
+    __type(value, struct ip_stats);
+} ip_stats_map SEC(".maps");
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, MAX_TRACK_IPS);
+    __type(key, __u32);
+    __type(value, __u64);
+} ip_blacklist_map SEC(".maps");
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, MAX_TRACK_IPS);
+#ifdef USE_FLOW_RL
+    __type(key, struct flow6);
+#else
+    __type(key, __u128);
+#endif
+    __type(value, struct ip_stats);
+} ip6_stats_map SEC(".maps");
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, MAX_TRACK_IPS);
+    __type(key, __u128);
+    __type(value, __u64);
+} ip6_blacklist_map SEC(".maps");

src/xdp/rl.h

@@ -0,0 +1,138 @@
+#pragma once
+
+#include <xdpfw.h>
+
+#include <xdp/maps.h>
+#include <xdp/helpers.h>
+
+/**
+ * Updates IPv4 client stats.
+ * 
+ * @param pps A pointer to the PPS integer.
+ * @param bps A pointer to the BPS integer.
+ * @param ip_stats A pointer to pointer to the IP stats structure value.
+ * @param ip The client's source IP.
+ * @param port The client's source port.
+ * @param protocol The client's protocol.
+ * @param pkt_len The total packet length.
+ * @param now The current time since boot in nanoseconds.alignas
+ * 
+ * @return void
+*/
+static __always_inline void UpdateIpStats(__u64 *pps, __u64 *bps, struct ip_stats **ip_stats, __u32 ip, __u16 port, __u8 protocol, __u16 pkt_len, __u64 now)
+{
+#ifdef USE_FLOW_RL
+    struct flow key = {0};
+    key.ip = ip;
+    key.port = port;
+    key.protocol = protocol;
+
+    *ip_stats = bpf_map_lookup_elem(&ip_stats_map, &key);
+#else
+    *ip_stats = bpf_map_lookup_elem(&ip_stats_map, &ip);
+#endif
+
+    if (*ip_stats)
+    {
+        // Check for next update.
+        if (now > (*ip_stats)->next_update)
+        {
+            (*ip_stats)->pps = 1;
+            (*ip_stats)->bps = pkt_len;
+            (*ip_stats)->next_update = now + NANO_TO_SEC;
+        }
+        else
+        {
+            // Increment PPS and BPS using built-in functions.
+            __sync_fetch_and_add(&(*ip_stats)->pps, 1);
+            __sync_fetch_and_add(&(*ip_stats)->bps, pkt_len);
+        }
+
+        *pps = (*ip_stats)->pps;
+        *bps = (*ip_stats)->bps;
+    }
+    else
+    {
+        // Create new entry.
+        struct ip_stats new = {0};
+
+        new.pps = 1;
+        new.bps = pkt_len;
+        new.next_update = now + NANO_TO_SEC;
+
+        *pps = new.pps;
+        *bps = new.bps;
+
+#ifdef USE_FLOW_RL
+        bpf_map_update_elem(&ip_stats_map, &key, &new, BPF_ANY);
+#else
+        bpf_map_update_elem(&ip_stats_map, &ip, &new, BPF_ANY);
+#endif
+    }
+}
+
+/**
+ * Updates IPv6 client stats.
+ * 
+ * @param pps A pointer to the PPS integer.
+ * @param bps A pointer to the BPS integer.
+ * @param ip_stats A pointer to pointer to the IP stats structure value.
+ * @param ip The client's source IP.
+ * @param port The client's source port.
+ * @param protocol The client's protocol.
+ * @param pkt_len The total packet length.
+ * @param now The current time since boot in nanoseconds.alignas
+ * 
+ * @return void
+*/
+static __always_inline void UpdateIp6Stats(__u64 *pps, __u64 *bps, struct ip_stats **ip_stats, __u128 *ip, __u16 port, __u8 protocol, __u16 pkt_len, __u64 now)
+{
+#ifdef USE_FLOW_RL
+    struct flow6 key = {0};
+    key.ip = *ip;
+    key.port = port;
+    key.protocol = protocol;
+
+    *ip_stats = bpf_map_lookup_elem(&ip_stats_map, &key);
+#else
+    *ip_stats = bpf_map_lookup_elem(&ip_stats_map, ip);
+#endif
+
+    if (*ip_stats)
+    {
+        // Check for next update.
+        if (now > (*ip_stats)->next_update)
+        {
+            (*ip_stats)->pps = 1;
+            (*ip_stats)->bps = pkt_len;
+            (*ip_stats)->next_update = now + NANO_TO_SEC;
+        }
+        else
+        {
+            // Increment PPS and BPS using built-in functions.
+            __sync_fetch_and_add(&(*ip_stats)->pps, 1);
+            __sync_fetch_and_add(&(*ip_stats)->bps, pkt_len);
+        }
+
+        *pps = (*ip_stats)->pps;
+        *bps = (*ip_stats)->bps;
+    }
+    else
+    {
+        // Create new entry.
+        struct ip_stats new = {0};
+
+        new.pps = 1;
+        new.bps = pkt_len;
+        new.next_update = now + NANO_TO_SEC;
+
+        *pps = new.pps;
+        *bps = new.bps;
+
+#ifdef USE_FLOW_RL
+        bpf_map_update_elem(&ip_stats_map, &key, &new, BPF_ANY);
+#else
+        bpf_map_update_elem(&ip_stats_map, ip, &new, BPF_ANY);
+#endif
+    }
+}

src/xdp/utils.h

@@ -1,6 +1,13 @@
 #pragma once
 
-#include "xdpfw.h"
+#include <xdpfw.h>
+
+#include <xdp/maps.h>
+#include <xdp/helpers.h>
+
+#ifndef memcpy
+#define memcpy(dest, src, n) __builtin_memcpy((dest), (src), (n))
+#endif
 
 /**
  * Checks if an IP is within a specific CIDR range.

src/xdpfw.c

@@ -20,10 +20,9 @@
 #include <libbpf.h>
 #include <xdp/libxdp.h>
 
-#include "xdpfw.h"
+#include <xdpfw.h>
-#include "config.h"
+#include <config.h>
-#include "cmdline.h"
+#include <cmdline.h>
-#include "xdp_utils.h"
 
 // Other variables.
 static __u8 cont = 1;

src/xdpfw.h

@@ -19,6 +19,11 @@
 // When this is defined, a check will occur inside the IPv4 and IPv6 filters. For IPv6 packets, if no IPv6 source/destination IP addresses are set, but there is an IPv4 address, it will ignore the filter. The same goes for IPv4, if there is no IPv4 source/destination IP addresses set, if an IPv6 address is set, it will ignore the filter.
 #define ALLOWSINGLEIPV4V6
 
+// If uncommented, rate limits for clients are determined using the source IP, port, and protocol instead of just the source IP.
+// This allows for more precise rate limits (connection-specific instead of a single source IP).
+// I decided not to include the destination IP/port because the source IP, port, and protocol should be represent a unique connection.
+#define USE_FLOW_RL
+
 #ifdef __BPF__
 #define likely(x) __builtin_expect(!!(x), 1)
 #define unlikely(x) __builtin_expect(!!(x), 0)
@@ -151,3 +156,17 @@ struct ip_stats
     __u64 bps;
     __u64 next_update;
 };
+
+struct flow
+{
+    __u32 ip;
+    __u16 port;
+    __u8 protocol;
+};
+
+struct flow6
+{
+    __u128 ip;
+    __u16 port;
+    __u8 protocol;
+};

src/xdpfw_kern.c

@@ -8,73 +8,13 @@
 #include <linux/in.h>
 #include <stdatomic.h>
 
-#include <linux/bpf.h>
+#include <xdp/helpers.h>
-#include <linux/bpf_common.h>
 
-#include <bpf_helpers.h>
+#include <xdpfw.h>
-#include <xdp/xdp_helpers.h>
-#include <xdp/prog_dispatcher.h>
 
-#include "xdpfw.h"
+#include <xdp/maps.h>
-#include "xdp_utils.h"
+#include <xdp/rl.h>
-
+#include <xdp/utils.h>
-#ifndef memcpy
-#define memcpy(dest, src, n) __builtin_memcpy((dest), (src), (n))
-#endif
-
-struct 
-{
-    __uint(priority, 10);
-    __uint(XDP_PASS, 1);
-} XDP_RUN_CONFIG(xdp_prog_main);
-
-struct 
-{
-    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
-    __uint(max_entries, MAX_FILTERS);
-    __type(key, __u32);
-    __type(value, struct filter);
-} filters_map SEC(".maps");
-
-struct 
-{
-    __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
-    __uint(max_entries, 1);
-    __type(key, __u32);
-    __type(value, struct stats);
-} stats_map SEC(".maps");
-
-struct 
-{
-    __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
-    __type(key, __u32);
-    __type(value, struct ip_stats);
-} ip_stats_map SEC(".maps");
-
-struct 
-{
-    __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
-    __type(key, __u32);
-    __type(value, __u64);
-} ip_blacklist_map SEC(".maps");
-
-struct 
-{
-    __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
-    __type(key, __u128);
-    __type(value, struct ip_stats);
-} ip6_stats_map SEC(".maps");
-
-struct 
-{
-    __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
-    __type(key, __u128);
-    __type(value, __u64);
-} ip6_blacklist_map SEC(".maps");
 
 SEC("xdp_prog")
 int xdp_prog_main(struct xdp_md *ctx)
@@ -154,10 +94,6 @@ int xdp_prog_main(struct xdp_md *ctx)
     
     if (blocked != NULL && *blocked > 0)
     {
-        #ifdef DEBUG
-        bpf_printk("Checking for blocked packet... Block time %llu.\n", *blocked);
-        #endif
-
         if (now > *blocked)
         {
             // Remove element from map.
@@ -185,72 +121,22 @@ int xdp_prog_main(struct xdp_md *ctx)
         }
     }
 
-    // Update IP stats (PPS/BPS).
+    // Retrieve total packet length.
-    __u64 pps = 0;
-    __u64 bps = 0;
-
-    struct ip_stats *ip_stats = NULL;
-    
-    if (iph6)
-    {
-        ip_stats = bpf_map_lookup_elem(&ip6_stats_map, &src_ip6);
-    }
-    else if (iph)
-    {
-        ip_stats = bpf_map_lookup_elem(&ip_stats_map, &iph->saddr);
-    }
-    
     __u16 pkt_len = data_end - data;
 
-    if (ip_stats)
+    // Parse layer-4 headers and determine source port and protocol.
-    {
-        // Check for next update.
-        if (now > ip_stats->next_update)
-        {
-            ip_stats->pps = 1;
-            ip_stats->bps = pkt_len;
-            ip_stats->next_update = now + NANO_TO_SEC;
-        }
-        else
-        {
-            // Increment PPS and BPS using built-in functions.
-            __sync_fetch_and_add(&ip_stats->pps, 1);
-            __sync_fetch_and_add(&ip_stats->bps, pkt_len);
-        }
-
-        pps = ip_stats->pps;
-        bps = ip_stats->bps;
-    }
-    else
-    {
-        // Create new entry.
-        struct ip_stats new = {0};
-
-        new.pps = 1;
-        new.bps = pkt_len;
-        new.next_update = now + NANO_TO_SEC;
-
-        pps = new.pps;
-        bps = new.bps;
-
-        if (iph6)
-        {
-            bpf_map_update_elem(&ip6_stats_map, &src_ip6, &new, BPF_ANY);
-        }
-        else if (iph)
-        {
-            bpf_map_update_elem(&ip_stats_map, &iph->saddr, &new, BPF_ANY);
-        } 
-    }
-
     struct tcphdr *tcph = NULL;
     struct udphdr *udph = NULL;
     struct icmphdr *icmph = NULL;
     struct icmp6hdr *icmp6h = NULL;
 
-    // Check protocol.
+    __u16 src_port = 0;
+    __u8 protocol = 0;
+    
     if (iph6)
     {
+        protocol = iph6->nexthdr;
+
         switch (iph6->nexthdr)
         {
             case IPPROTO_TCP:
@@ -263,6 +149,8 @@ int xdp_prog_main(struct xdp_md *ctx)
                     return XDP_DROP;
                 }
 
+                src_port = tcph->source;
+
                 break;
 
             case IPPROTO_UDP:
@@ -275,6 +163,8 @@ int xdp_prog_main(struct xdp_md *ctx)
                     return XDP_DROP;
                 }
 
+                src_port = udph->source;
+
                 break;
 
             case IPPROTO_ICMPV6:
@@ -292,6 +182,8 @@ int xdp_prog_main(struct xdp_md *ctx)
     }
     else if (iph)
     {
+        protocol = iph->protocol;
+
         switch (iph->protocol)
         {
             case IPPROTO_TCP:
@@ -304,6 +196,8 @@ int xdp_prog_main(struct xdp_md *ctx)
                     return XDP_DROP;
                 }
 
+                src_port = tcph->source;
+
                 break;
 
             case IPPROTO_UDP:
@@ -316,6 +210,8 @@ int xdp_prog_main(struct xdp_md *ctx)
                     return XDP_DROP;
                 }
 
+                src_port = udph->source;
+
                 break;
 
             case IPPROTO_ICMP:
@@ -332,6 +228,23 @@ int xdp_prog_main(struct xdp_md *ctx)
         }
     }
 
+    // Update client stats (PPS/BPS).
+    __u64 pps = 0;
+    __u64 bps = 0;
+
+    struct ip_stats *ip_stats = NULL;
+    
+    if (iph6)
+    {
+        UpdateIp6Stats(&pps, &bps, &ip_stats, &src_ip6, src_port, protocol, pkt_len, now);
+    }
+    else if (iph)
+    {
+        UpdateIpStats(&pps, &bps, &ip_stats, iph->saddr, src_port, protocol, pkt_len, now);
+    }
+
+    bpf_printk("PPS => %llu. BPS => %llu.\n", pps, bps);
+    
     for (__u8 i = 0; i < MAX_FILTERS; i++)
     {
         __u32 key = i;
@@ -601,10 +514,6 @@ int xdp_prog_main(struct xdp_md *ctx)
         }
         
         // Matched.
-        #ifdef DEBUG
-        bpf_printk("Matched rule ID #%d.\n", filter->id);
-        #endif
-        
         action = filter->action;
         blocktime = filter->blocktime;
 
@@ -621,10 +530,6 @@ int xdp_prog_main(struct xdp_md *ctx)
     matched:
     if (action == 0)
     {
-        #ifdef DEBUG
-        //bpf_printk("Matched with protocol %d and sAddr %lu.\n", iph->protocol, iph->saddr);
-        #endif
-
         // Before dropping, update the blacklist map.
         if (blocktime > 0)
         {