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

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

@@ -0,0 +1,24 @@
+#pragma once
+
+#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.
+ * 
+ * @param src_ip The source/main IP to check against.
+ * @param net_ip The network IP.
+ * @param cidr The CIDR range.
+ * 
+ * @return 1 on yes, 0 on no.
+*/
+static __always_inline __u8 IsIpInRange(__u32 src_ip, __u32 net_ip, __u8 cidr)
+{
+    return !((src_ip ^ net_ip) & htonl(0xFFFFFFFFu << (32 - cidr)));
+}