Merge pull request #67 from gamemann/20250322-rl-changes

Add Flow And IP-Based Rate Limiting

README.md

@@ -173,8 +173,10 @@ Here are more details on the layout of the runtime configuration.
 | log | bool | `false` | Whether to log packets that are matched. |
 | action | int | `1` | The value of `0` drops or blocks the packet while `1` allows/passes the packet through. |
 | block_time | int | `1` | The amount of seconds to block the source IP for if matched. |
-| pps | int64 | `NULL` | Matches if this threshold of packets per second is exceeded for a source IP. |
+| ip_pps | int64 | `NULL` | Matches if this threshold of packets per second is exceeded for a source IP. |
-| bps | int64 | `NULL` | Matches if this threshold of bytes per second is exceeded for a source IP. |
+| ip_bps | int64 | `NULL` | Matches if this threshold of bytes per second is exceeded for a source IP. |
+| flow_pps | int64 | `NULL` | Matches if this threshold of packets per second is exceeded for a source flow (IP and port). |
+| flow_bps | int64 | `NULL` | Matches if this threshold of bytes per second is exceeded for a source flow (IP and port). |
 
 #### IP Options
 | Name | Type | Default | Description |
@@ -211,8 +213,8 @@ You may additionally specified UDP header options for a filter rule which start
 | Name | Type | Default | Description |
 | ---- | ---- | ------- | ----------- |
 | udp_enabled | bool | `false` | Whether to enable UDP on this filter rule. |
-| udp_sport | int | `NULL` | The UDP source port to match with single range support (e.g., `"27000-27015"`). |
+| udp_sport | int \| string | `NULL` | The UDP source port to match with single range support (e.g., `"27000-27015"`). |
-| udp_dport | int | `NULL` | The UDP destination port to match with single range support (e.g., `"27000-27015"`). |
+| udp_dport | int \| string | `NULL` | The UDP destination port to match with single range support (e.g., `"27000-27015"`). |
 
 #### ICMP Options
 You may additionally specified UDP header options for a filter rule which start with `icmp_`.

src/common/config.h

@@ -15,10 +15,6 @@
 // Decrease this value if you receive errors related to the BPF program being too large.
 #define MAX_FILTERS 60
 
-// The maximum amount of IPs/flows to track stats for.
-// The higher this value is, the more memory that'll be used.
-#define MAX_TRACK_IPS 100000
-
 // Feel free to comment this out if you don't want the `blocked` entry on the stats map to be incremented every single time a packet is dropped from the source IP being on the blocked map.
 // Commenting this line out should increase performance when blocking malicious traffic.
 #define DO_STATS_ON_BLOCK_MAP
@@ -31,14 +27,26 @@
 // 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 ALLOW_SINGLE_IP_V4_V6
 
-// 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
-
 // Enables filter logging through XDP.
 // If performance is a concern, it is best to disable this feature by commenting out the below line with //.
 #define ENABLE_FILTER_LOGGING
 
 // Maximum interfaces the firewall can attach to.
 #define MAX_INTERFACES 6
+
+// NOTE - If you're receiving a high volume of spoofed packets, it is recommended you disable rate limiting below.
+// This is because the PPS/BPS counters are updated for every packet and with a spoofed attack, the LRU map will recycle a lot of entries resulting in additional load on the CPU.
+// Enable source IP rate limiting.
+#define ENABLE_RL_IP
+
+// Enable source flow rate limiting.
+#define ENABLE_RL_FLOW
+
+// Maximum entries in source IP rate limit map.
+#define MAX_RL_IP 100000
+
+// Maximum entries in source flow rate limit map.
+#define MAX_RL_FLOW 100000
+
+// Maximum entries in block map.
+#define MAX_BLOCK 100000

src/common/types.h

@@ -108,11 +108,21 @@ struct filter
     u8 action;
     u16 block_time;
 
-    unsigned int do_pps : 1;
+#ifdef ENABLE_RL_IP
-    u64 pps;
+    unsigned int do_ip_pps : 1;
+    u64 ip_pps;
 
-    unsigned int do_bps : 1;
+    unsigned int do_ip_bps : 1;
-    u64 bps;
+    u64 ip_bps;
+#endif
+
+#ifdef ENABLE_RL_FLOW
+    unsigned int do_flow_pps : 1;
+    u64 flow_pps;
+
+    unsigned int do_flow_bps : 1;
+    u64 flow_bps;
+#endif
     
     filter_ip_t ip;
 
@@ -128,12 +138,12 @@ struct stats
     u64 passed;
 } typedef stats_t;
 
-struct ip_stats
+struct cl_stats
 {
     u64 pps;
     u64 bps;
     u64 next_update;
-} typedef ip_stats_t ;
+} typedef cl_stats_t;
 
 struct flow
 {
@@ -168,8 +178,11 @@ struct filter_log_event
 
     u8 protocol;
 
-    u64 pps;
+    u64 ip_pps;
-    u64 bps;
+    u64 ip_bps;
+
+    u64 flow_pps;
+    u64 flow_bps;
 } typedef filter_log_event_t;
 
 struct lpm_trie_key

src/loader/utils/config.c

@@ -404,20 +404,36 @@ int parse_cfg(config__t *cfg, const char* data, config_overrides_t* overrides)
                 filter->block_time = block_time;
             }
 
-            // PPS (not required).
+            // IP PPS (not required).
-            s64 pps;
+            s64 ip_pps;
 
-            if (config_setting_lookup_int64(filter_cfg, "pps", &pps) == CONFIG_TRUE)
+            if (config_setting_lookup_int64(filter_cfg, "ip_pps", &ip_pps) == CONFIG_TRUE)
             {
-                filter->pps = pps;
+                filter->ip_pps = ip_pps;
             }
 
-            // BPS (not required).
+            // IP BPS (not required).
-            s64 bps;
+            s64 ip_bps;
 
-            if (config_setting_lookup_int64(filter_cfg, "bps", &bps) == CONFIG_TRUE)
+            if (config_setting_lookup_int64(filter_cfg, "ip_bps", &ip_bps) == CONFIG_TRUE)
             {
-                filter->bps = bps;
+                filter->ip_bps = ip_bps;
+            }
+
+            // Flow PPS (not required).
+            s64 flow_pps;
+
+            if (config_setting_lookup_int64(filter_cfg, "flow_pps", &flow_pps) == CONFIG_TRUE)
+            {
+                filter->flow_pps = flow_pps;
+            }
+
+            // Flow BPS (not required).
+            s64 flow_bps;
+
+            if (config_setting_lookup_int64(filter_cfg, "flow_bps", &flow_bps) == CONFIG_TRUE)
+            {
+                filter->flow_bps = flow_bps;
             }
 
             /* IP Options */
@@ -874,18 +890,32 @@ int save_cfg(config__t* cfg, const char* file_path)
                     config_setting_set_int(block_time, filter->block_time);
                 }
 
-                // Add PPS.
+                // Add IP PPS.
-                if (filter->pps > -1)
+                if (filter->ip_pps > -1)
                 {
-                    config_setting_t* pps = config_setting_add(filter_cfg, "pps", CONFIG_TYPE_INT64);
+                    config_setting_t* pps = config_setting_add(filter_cfg, "ip_pps", CONFIG_TYPE_INT64);
-                    config_setting_set_int64(pps, filter->pps);
+                    config_setting_set_int64(pps, filter->ip_pps);
                 }
 
-                // Add BPS.
+                // Add IP BPS.
-                if (filter->bps > -1)
+                if (filter->ip_bps > -1)
                 {
-                    config_setting_t* bps = config_setting_add(filter_cfg, "bps", CONFIG_TYPE_INT64);
+                    config_setting_t* bps = config_setting_add(filter_cfg, "ip_bps", CONFIG_TYPE_INT64);
-                    config_setting_set_int64(bps, filter->bps);
+                    config_setting_set_int64(bps, filter->ip_bps);
+                }
+
+                // Add flow PPS.
+                if (filter->flow_pps > -1)
+                {
+                    config_setting_t* pps = config_setting_add(filter_cfg, "flow_pps", CONFIG_TYPE_INT64);
+                    config_setting_set_int64(pps, filter->flow_pps);
+                }
+
+                // Add flow BPS.
+                if (filter->flow_bps > -1)
+                {
+                    config_setting_t* bps = config_setting_add(filter_cfg, "flow_bps", CONFIG_TYPE_INT64);
+                    config_setting_set_int64(bps, filter->flow_bps);
                 }
 
                 // Add source IPv4.
@@ -1130,8 +1160,10 @@ void set_filter_defaults(filter_rule_cfg_t* filter)
     filter->action = 1;
     filter->block_time = 1;
 
-    filter->pps = -1;
+    filter->ip_pps = -1;
-    filter->bps = -1;
+    filter->ip_bps = -1;
+    filter->flow_pps = -1;
+    filter->flow_bps = -1;
 
     if (filter->ip.src_ip)
     {
@@ -1299,8 +1331,11 @@ void print_filter(filter_rule_cfg_t* filter, int idx)
     printf("\t\tAction => %d (0 = Block, 1 = Allow).\n", filter->action);
     printf("\t\t\tBlock Time => %d\n\n", filter->block_time);
 
-    printf("\t\t\tPPS => %lld\n", filter->pps);
+    printf("\t\t\tIP PPS => %lld\n", filter->ip_pps);
-    printf("\t\t\tBPS => %lld\n\n", filter->bps);
+    printf("\t\t\tIP BPS => %lld\n", filter->ip_bps);
+
+    printf("\t\t\tFlow PPS => %lld\n", filter->flow_pps);
+    printf("\t\t\tFlow BPS => %lld\n", filter->flow_bps);
 
     // IP Options.
     printf("\t\tIP Options\n");

src/loader/utils/config.h

@@ -72,8 +72,11 @@ struct filter_rule_cfg
     int action;
     int block_time;
 
-    s64 pps;
+    s64 ip_pps;
-    s64 bps;
+    s64 ip_bps;
+
+    s64 flow_pps;
+    s64 flow_bps;
 
     filter_rule_ip_opts_t ip;
     

src/loader/utils/logging.c

@@ -163,7 +163,7 @@ int hdl_filters_rb_event(void* ctx, void* data, size_t sz)
 
     const char* protocol_str = get_protocol_str_by_id(e->protocol);
 
-    log_msg(cfg, 0, 0, "[FILTER %d] %s %s packet '%s:%d' => '%s:%d' (PPS => %llu, BPS => %llu, Filter Block Time => %llu, length => %d)...", e->filter_id + 1, action, protocol_str, src_ip_str, htons(e->src_port), dst_ip_str, htons(e->dst_port), e->pps, e->bps, filter->block_time, e->length);
+    log_msg(cfg, 0, 0, "[FILTER %d] %s %s packet '%s:%d' => '%s:%d' (IP PPS => %llu, IP BPS => %llu, Flow PPS => %llu, Flow BPS => %llu Filter Block Time => %llu, length => %d)...", e->filter_id + 1, action, protocol_str, src_ip_str, htons(e->src_port), dst_ip_str, htons(e->dst_port), e->ip_pps, e->ip_bps, e->flow_pps, e->flow_bps, filter->block_time, e->length);
 
     return 0;
 }

src/loader/utils/xdp.c

@@ -257,19 +257,37 @@ int update_filter(int map_filters, filter_rule_cfg_t* filter_cfg, int idx)
         filter.block_time = filter_cfg->block_time;
     }
 
-    if (filter_cfg->pps > -1)
+#ifdef ENABLE_RL_IP
+    if (filter_cfg->ip_pps > -1)
     {
-        filter.do_pps = 1;
+        filter.do_ip_pps = 1;
 
-        filter.pps = (u64) filter_cfg->pps;
+        filter.ip_pps = (u64) filter_cfg->ip_pps;
     }
 
-    if (filter_cfg->bps > -1)
+    if (filter_cfg->ip_bps > -1)
     {
-        filter.do_bps = 1;
+        filter.do_ip_bps = 1;
 
-        filter.bps = (u64) filter_cfg->bps;
+        filter.ip_bps = (u64) filter_cfg->ip_bps;
     }
+#endif
+
+#ifdef ENABLE_RL_FLOW
+    if (filter_cfg->flow_pps > -1)
+    {
+        filter.do_flow_pps = 1;
+
+        filter.flow_pps = (u64) filter_cfg->flow_pps;
+    }
+
+    if (filter_cfg->flow_bps > -1)
+    {
+        filter.do_flow_bps = 1;
+
+        filter.flow_bps = (u64) filter_cfg->flow_bps;
+    }
+#endif
 
     if (filter_cfg->ip.src_ip)
     {

src/rule_add/prog.c

@@ -29,8 +29,11 @@ int main(int argc, char *argv[])
     cli.action = -1;
     cli.block_time = -1;
 
-    cli.pps = -1;
+    cli.ip_pps = -1;
-    cli.bps = -1;
+    cli.ip_bps = -1;
+
+    cli.flow_pps = -1;
+    cli.flow_bps = -1;
 
     cli.min_ttl = -1;
     cli.max_ttl = -1;
@@ -223,14 +226,24 @@ int main(int argc, char *argv[])
 
         // To Do: See if I can create a macro for below.
         // As long as the naming convention lines up, it should be easily possible.
-        if (cli.pps > -1)
+        if (cli.ip_pps > -1)
         {
-            new_filter.pps = cli.pps;
+            new_filter.ip_pps = cli.ip_pps;
         }
 
-        if (cli.bps > -1)
+        if (cli.ip_bps > -1)
         {
-            new_filter.bps = cli.bps;
+            new_filter.ip_bps = cli.ip_bps;
+        }
+
+        if (cli.flow_pps > -1)
+        {
+            new_filter.flow_pps = cli.flow_pps;
+        }
+
+        if (cli.flow_bps > -1)
+        {
+            new_filter.flow_bps = cli.flow_bps;
         }
 
         if (cli.min_ttl > -1)

src/rule_add/utils/cli.c

@@ -30,8 +30,11 @@ const struct option opts[] =
     { "max-len", required_argument, NULL, 7 },
     { "tos", required_argument, NULL, 8 },
 
-    { "pps", required_argument, NULL, 9 },
+    { "ip-pps", required_argument, NULL, 9 },
-    { "bps", required_argument, NULL, 10 },
+    { "ip-bps", required_argument, NULL, 10 },
+
+    { "ip-pps", required_argument, NULL, 32 },
+    { "ip-bps", required_argument, NULL, 33 },
 
     { "tcp", required_argument, NULL, 11 },
     { "tsport", required_argument, NULL, 12 },
@@ -170,12 +173,22 @@ void parse_cli(cli_t* cli, int argc, char* argv[])
                 break;
 
             case 9:
-                cli->pps = strtoll(optarg, NULL, 10);
+                cli->ip_pps = strtoll(optarg, NULL, 10);
 
                 break;
 
             case 10:
-                cli->bps = strtoll(optarg, NULL, 10);
+                cli->ip_bps = strtoll(optarg, NULL, 10);
+
+                break;
+
+            case 32:
+                cli->flow_pps = strtoll(optarg, NULL, 10);
+
+                break;
+
+            case 33:
+                cli->flow_bps = strtoll(optarg, NULL, 10);
 
                 break;
 

src/rule_add/utils/cli.h

@@ -36,8 +36,11 @@ struct cli
     char* src_ip6;
     char* dst_ip6;
 
-    s64 pps;
+    s64 ip_pps;
-    s64 bps;
+    s64 ip_bps;
+
+    s64 flow_pps;
+    s64 flow_bps;
 
     int min_ttl;
     int max_ttl;

src/xdp/prog.c

@@ -285,22 +285,46 @@ int xdp_prog_main(struct xdp_md *ctx)
         }
     }
 
+#ifdef ENABLE_FILTERS
     // Update client stats (PPS/BPS).
-    u64 pps = 0;
+#ifdef ENABLE_RL_IP
-    u64 bps = 0;
+    u64 ip_pps = 0;
+    u64 ip_bps = 0;
+#endif
 
-    if (iph6)
+#ifdef ENABLE_RL_FLOW
+    u64 flow_pps = 0;
+    u64 flow_bps = 0;
+#endif
+
+#if defined(ENABLE_RL_IP) || defined(ENABLE_RL_FLOW)
+    if (iph)
     {
-        update_ip6_stats(&pps, &bps, &src_ip6, src_port, protocol, pkt_len, now);
+#ifdef ENABLE_RL_IP
+        update_ip_stats(&ip_pps, &ip_bps, iph->saddr, pkt_len, now);
+#endif
+
+#ifdef ENABLE_RL_FLOW
+        update_flow_stats(&flow_pps, &flow_bps, iph->saddr, src_port, protocol, pkt_len, now);
+#endif
     }
-    else if (iph)
+    else if (iph6)
     {
-        update_ip_stats(&pps, &bps, iph->saddr, src_port, protocol, pkt_len, now);
+#ifdef ENABLE_RL_IP
+        update_ip6_stats(&ip_pps, &ip_bps, &src_ip6, pkt_len, now);
+#endif
+
+#ifdef ENABLE_RL_FLOW
+        update_flow6_stats(&flow_pps, &flow_bps, &src_ip6, src_port, protocol, pkt_len, now);
+#endif
     }
+#endif
+#endif
 
     int action = 0;
     u64 block_time = 1;
 
+#pragma unroll 30
     for (int i = 0; i < MAX_FILTERS; i++)
     {
         u32 key = i;
@@ -312,6 +336,32 @@ int xdp_prog_main(struct xdp_md *ctx)
             break;
         }
 
+#ifdef ENABLE_RL_IP
+        // Check source IP rate limits.
+        if (filter->do_ip_pps && ip_pps < filter->ip_pps)
+        {
+            continue;
+        }
+
+        if (filter->do_ip_bps && ip_bps < filter->ip_bps)
+        {
+            continue;
+        }
+#endif
+
+#ifdef ENABLE_RL_FLOW
+        // Check source flow rate limits.
+        if (filter->do_flow_pps && flow_pps < filter->flow_pps)
+        {
+            continue;
+        }
+
+        if (filter->do_flow_bps && flow_bps < filter->flow_bps)
+        {
+            continue;
+        }
+#endif
+
         // Do specific IPv6.
         if (iph6)
         {
@@ -426,18 +476,6 @@ int xdp_prog_main(struct xdp_md *ctx)
             }
         }
 
-        // PPS.
-        if (filter->do_pps && pps < filter->pps)
-        {
-            continue;
-        }
-
-        // BPS.
-        if (filter->do_bps && bps < filter->bps)
-        {
-            continue;
-        }
-        
         // Do TCP options.
         if (filter->tcp.enabled)
         {
@@ -584,7 +622,19 @@ int xdp_prog_main(struct xdp_md *ctx)
 #ifdef ENABLE_FILTER_LOGGING
         if (filter->log > 0)
         {
-            log_filter_msg(iph, iph6, src_port, dst_port, protocol, now, pps, bps, pkt_len, i);
+#if defined(ENABLE_RL_IP) || defined(ENABLE_RL_FLOW)
+#if defined(ENABLE_RL_IP) && defined(ENABLE_RL_FLOW)
+            log_filter_msg(iph, iph6, src_port, dst_port, protocol, now, ip_pps, ip_bps, flow_pps, flow_bps, pkt_len, i);
+#else
+#ifdef ENABLE_RL_IP
+            log_filter_msg(iph, iph6, src_port, dst_port, protocol, now, ip_pps, ip_bps, 0, 0, pkt_len, i);
+#else
+            log_filter_msg(iph, iph6, src_port, dst_port, protocol, now, 0, 0, flow_pps, flow_bps, pkt_len, i);
+#endif
+#endif
+#else
+            log_filter_msg(iph, iph6, src_port, dst_port, protocol, now, 0, 0, 0, 0, pkt_len, i);
+#endif
         }
 #endif
         

src/xdp/utils/logging.c

@@ -14,14 +14,16 @@
  * @param dst_port The destination port.
  * @param protocol The protocol.
  * @param now The timestamp.
- * @param pps The current PPS rate.
+ * @param ip_pps The current IP PPS rate.
- * @param bps The current BPS rate.
+ * @param ip_bps The current IP BPS rate.
+ * @param flow_pps The current flow PPS rate.
+ * @param flow_bps The current flow BPS rate.
  * @param pkt_len The full packet length.
  * @param filter_id The filter ID that matched.
  * 
  * @return always 0
  */
-static __always_inline int log_filter_msg(struct iphdr* iph, struct ipv6hdr* iph6, u16 src_port, u16 dst_port, u8 protocol, u64 now, u64 pps, u64 bps, int pkt_len, int filter_id)
+static __always_inline int log_filter_msg(struct iphdr* iph, struct ipv6hdr* iph6, u16 src_port, u16 dst_port, u8 protocol, u64 now, u64 ip_pps, u64 ip_bps, u64 flow_pps, u64 flow_bps, int pkt_len, int filter_id)
 {
     filter_log_event_t* e = bpf_ringbuf_reserve(&map_filter_log, sizeof(*e), 0);
 
@@ -45,8 +47,11 @@ static __always_inline int log_filter_msg(struct iphdr* iph, struct ipv6hdr* iph
 
         e->protocol = protocol;
 
-        e->pps = pps;
+        e->ip_pps = ip_pps;
-        e->bps = bps;
+        e->ip_bps = ip_bps;
+
+        e->flow_pps = flow_pps;
+        e->flow_bps = flow_bps;
 
         e->length = pkt_len;
 

src/xdp/utils/logging.h

@@ -6,7 +6,7 @@
 #include <xdp/prog_dispatcher.h>
 
 #if defined(ENABLE_FILTERS) && defined(ENABLE_FILTER_LOGGING)
-static __always_inline int log_filter_msg(struct iphdr* iph, struct ipv6hdr* iph6, u16 src_port, u16 dst_port, u8 protocol, u64 now, u64 pps, u64 bps, int pkt_len, int filter_id);
+static __always_inline int log_filter_msg(struct iphdr* iph, struct ipv6hdr* iph6, u16 src_port, u16 dst_port, u8 protocol, u64 now, u64 ip_pps, u64 ip_bps, u64 flow_pps, u64 flow_bps, int pkt_len, int filter_id);
 #endif
 
 // The source file is included directly below instead of compiled and linked as an object because when linking, there is no guarantee the compiler will inline the function (which is crucial for performance).

src/xdp/utils/maps.h

@@ -16,7 +16,7 @@ struct
 struct 
 {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
+    __uint(max_entries, MAX_BLOCK);
     __type(key, u32);
     __type(value, u64);
 } map_block SEC(".maps");
@@ -24,7 +24,7 @@ struct
 struct 
 {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
+    __uint(max_entries, MAX_BLOCK);
     __type(key, u128);
     __type(value, u64);
 } map_block6 SEC(".maps");
@@ -41,29 +41,41 @@ struct
 #endif
 
 #ifdef ENABLE_FILTERS
+#ifdef ENABLE_RL_IP
 struct 
 {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
+    __uint(max_entries, MAX_RL_IP);
-#ifdef USE_FLOW_RL
-    __type(key, flow_t);
-#else
     __type(key, u32);
-#endif
+    __type(value, cl_stats_t);
-    __type(value, ip_stats_t);
 } map_ip_stats SEC(".maps");
 
 struct 
 {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
-    __uint(max_entries, MAX_TRACK_IPS);
+    __uint(max_entries, MAX_RL_IP);
-#ifdef USE_FLOW_RL
-    __type(key, flow6_t);
-#else
     __type(key, u128);
-#endif
+    __type(value, cl_stats_t);
-    __type(value, ip_stats_t);
 } map_ip6_stats SEC(".maps");
+#endif
+
+#ifdef ENABLE_RL_FLOW
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, MAX_RL_FLOW);
+    __type(key, flow_t);
+    __type(value, cl_stats_t);
+} map_flow_stats SEC(".maps");
+
+struct 
+{
+    __uint(type, BPF_MAP_TYPE_LRU_HASH);
+    __uint(max_entries, MAX_RL_FLOW);
+    __type(key, flow6_t);
+    __type(value, cl_stats_t);
+} map_flow6_stats SEC(".maps");
+#endif
 
 struct 
 {

src/xdp/utils/rl.c

@@ -1,8 +1,116 @@
 #include <xdp/utils/rl.h>
 
 #ifdef ENABLE_FILTERS
+
+#ifdef ENABLE_RL_IP
 /**
- * Updates IPv4 client stats.
+ * Updates source IPv4 address stats.
+ * 
+ * @param pps A pointer to the PPS integer.
+ * @param bps A pointer to the BPS integer.
+ * @param ip The client's source IP.
+ * @param pkt_len The total packet length.
+ * @param now The current time since boot in nanoseconds.alignas
+ * 
+ * @return always 0
+ */
+static __always_inline int update_ip_stats(u64 *pps, u64 *bps, u32 ip, u16 pkt_len, u64 now)
+{
+    cl_stats_t* stats = bpf_map_lookup_elem(&map_ip_stats, &ip);
+
+    if (stats)
+    {
+        // Check for next update.
+        if (now > stats->next_update)
+        {
+            stats->pps = 1;
+            stats->bps = pkt_len;
+            stats->next_update = now + NANO_TO_SEC;
+        }
+        else
+        {
+            // Increment PPS and BPS using built-in functions.
+            __sync_fetch_and_add(&stats->pps, 1);
+            __sync_fetch_and_add(&stats->bps, pkt_len);
+        }
+
+        *pps = stats->pps;
+        *bps = stats->bps;
+    }
+    else
+    {
+        // Create new entry.
+        cl_stats_t new = {0};
+
+        new.pps = 1;
+        new.bps = pkt_len;
+        new.next_update = now + NANO_TO_SEC;
+
+        *pps = new.pps;
+        *bps = new.bps;
+
+        bpf_map_update_elem(&map_ip_stats, &ip, &new, BPF_ANY);
+    }
+
+    return 0;
+}
+
+/**
+ * Updates source IPv6 address stats.
+ * 
+ * @param pps A pointer to the PPS integer.
+ * @param bps A pointer to the BPS integer.
+ * @param ip The client's source IP.
+ * @param pkt_len The total packet length.
+ * @param now The current time since boot in nanoseconds.alignas
+ * 
+ * @return always 0
+ */
+static __always_inline int update_ip6_stats(u64 *pps, u64 *bps, u128 *ip, u16 pkt_len, u64 now)
+{
+    cl_stats_t* stats = bpf_map_lookup_elem(&map_ip6_stats, ip);
+
+    if (stats)
+    {
+        // Check for next update.
+        if (now > stats->next_update)
+        {
+            stats->pps = 1;
+            stats->bps = pkt_len;
+            stats->next_update = now + NANO_TO_SEC;
+        }
+        else
+        {
+            // Increment PPS and BPS using built-in functions.
+            __sync_fetch_and_add(&stats->pps, 1);
+            __sync_fetch_and_add(&stats->bps, pkt_len);
+        }
+
+        *pps = stats->pps;
+        *bps = stats->bps;
+    }
+    else
+    {
+        // Create new entry.
+        cl_stats_t new = {0};
+
+        new.pps = 1;
+        new.bps = pkt_len;
+        new.next_update = now + NANO_TO_SEC;
+
+        *pps = new.pps;
+        *bps = new.bps;
+
+        bpf_map_update_elem(&map_ip6_stats, ip, &new, BPF_ANY);
+    }
+
+    return 0;
+}
+#endif
+
+#ifdef ENABLE_RL_FLOW
+/**
+ * Updates IPv4 flow stats.
  * 
  * @param pps A pointer to the PPS integer.
  * @param bps A pointer to the BPS integer.
@@ -10,46 +118,42 @@
  * @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
+ * @param now The current time since boot in nanoseconds.
  * 
- * @return void
+ * @return always 0
  */
-static __always_inline void update_ip_stats(u64 *pps, u64 *bps, u32 ip, u16 port, u8 protocol, u16 pkt_len, u64 now)
+static __always_inline int update_flow_stats(u64 *pps, u64 *bps, u32 ip, u16 port, u8 protocol, u16 pkt_len, u64 now)
 {
-#ifdef USE_FLOW_RL
     flow_t key = {0};
     key.ip = ip;
     key.port = port;
     key.protocol = protocol;
 
-    ip_stats_t *ip_stats = bpf_map_lookup_elem(&map_ip_stats, &key);
+    cl_stats_t* stats = bpf_map_lookup_elem(&map_flow_stats, &key);
-#else
-    ip_stats_t *ip_stats = bpf_map_lookup_elem(&map_ip_stats, &ip);
-#endif
 
-    if (ip_stats)
+    if (stats)
     {
         // Check for next update.
-        if (now > ip_stats->next_update)
+        if (now > stats->next_update)
         {
-            ip_stats->pps = 1;
+            stats->pps = 1;
-            ip_stats->bps = pkt_len;
+            stats->bps = pkt_len;
-            ip_stats->next_update = now + NANO_TO_SEC;
+            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(&stats->pps, 1);
-            __sync_fetch_and_add(&ip_stats->bps, pkt_len);
+            __sync_fetch_and_add(&stats->bps, pkt_len);
         }
 
-        *pps = ip_stats->pps;
+        *pps = stats->pps;
-        *bps = ip_stats->bps;
+        *bps = stats->bps;
     }
     else
     {
         // Create new entry.
-        ip_stats_t new = {0};
+        cl_stats_t new = {0};
 
         new.pps = 1;
         new.bps = pkt_len;
@@ -58,16 +162,14 @@ static __always_inline void update_ip_stats(u64 *pps, u64 *bps, u32 ip, u16 port
         *pps = new.pps;
         *bps = new.bps;
 
-#ifdef USE_FLOW_RL
+        bpf_map_update_elem(&map_flow_stats, &key, &new, BPF_ANY);
-        bpf_map_update_elem(&map_ip_stats, &key, &new, BPF_ANY);
-#else
-        bpf_map_update_elem(&map_ip_stats, &ip, &new, BPF_ANY);
-#endif
     }
+
+    return 0;
 }
 
 /**
- * Updates IPv6 client stats.
+ * Updates IPv6 flow stats.
  * 
  * @param pps A pointer to the PPS integer.
  * @param bps A pointer to the BPS integer.
@@ -75,46 +177,42 @@ static __always_inline void update_ip_stats(u64 *pps, u64 *bps, u32 ip, u16 port
  * @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
+ * @param now The current time since boot in nanoseconds.
  * 
- * @return void
+ * @return always 0
  */
-static __always_inline void update_ip6_stats(u64 *pps, u64 *bps, u128 *ip, u16 port, u8 protocol, u16 pkt_len, u64 now)
+static __always_inline int update_flow6_stats(u64 *pps, u64 *bps, u128 *ip, u16 port, u8 protocol, u16 pkt_len, u64 now)
 {
-#ifdef USE_FLOW_RL
     flow6_t key = {0};
     key.ip = *ip;
     key.port = port;
     key.protocol = protocol;
 
-    ip_stats_t *ip_stats = bpf_map_lookup_elem(&map_ip6_stats, &key);
+    cl_stats_t* stats = bpf_map_lookup_elem(&map_flow6_stats, &key);
-#else
-    ip_stats_t *ip_stats = bpf_map_lookup_elem(&map_ip6_stats, ip);
-#endif
 
-    if (ip_stats)
+    if (stats)
     {
         // Check for next update.
-        if (now > ip_stats->next_update)
+        if (now > stats->next_update)
         {
-            ip_stats->pps = 1;
+            stats->pps = 1;
-            ip_stats->bps = pkt_len;
+            stats->bps = pkt_len;
-            ip_stats->next_update = now + NANO_TO_SEC;
+            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(&stats->pps, 1);
-            __sync_fetch_and_add(&ip_stats->bps, pkt_len);
+            __sync_fetch_and_add(&stats->bps, pkt_len);
         }
 
-        *pps = ip_stats->pps;
+        *pps = stats->pps;
-        *bps = ip_stats->bps;
+        *bps = stats->bps;
     }
     else
     {
         // Create new entry.
-        ip_stats_t new = {0};
+        cl_stats_t new = {0};
 
         new.pps = 1;
         new.bps = pkt_len;
@@ -123,11 +221,10 @@ static __always_inline void update_ip6_stats(u64 *pps, u64 *bps, u128 *ip, u16 p
         *pps = new.pps;
         *bps = new.bps;
 
-#ifdef USE_FLOW_RL
+        bpf_map_update_elem(&map_flow6_stats, &key, &new, BPF_ANY);
-        bpf_map_update_elem(&map_ip6_stats, &key, &new, BPF_ANY);
-#else
-        bpf_map_update_elem(&map_ip6_stats, ip, &new, BPF_ANY);
-#endif
     }
+
+    return 0;
 }
 #endif
+#endif

src/xdp/utils/rl.h

@@ -7,8 +7,10 @@
 #include <xdp/utils/maps.h>
 
 #ifdef ENABLE_FILTERS
-static __always_inline void update_ip_stats(u64 *pps, u64 *bps, u32 ip, u16 port, u8 protocol, u16 pkt_len, u64 now);
+static __always_inline int update_ip_stats(u64 *pps, u64 *bps, u32 ip, u16 pkt_len, u64 now);
-static __always_inline void update_ip6_stats(u64 *pps, u64 *bps, u128 *ip, u16 port, u8 protocol, u16 pkt_len, u64 now);
+static __always_inline int update_ip6_stats(u64 *pps, u64 *bps, u128 *ip, u16 pkt_len, u64 now);
+static __always_inline int update_flow_stats(u64 *pps, u64 *bps, u32 ip, u16 port, u8 protocol, u16 pkt_len, u64 now);
+static __always_inline int update_flow6_stats(u64 *pps, u64 *bps, u128 *ip, u16 port, u8 protocol, u16 pkt_len, u64 now);
 #endif
 
 // The source file is included directly below instead of compiled and linked as an object because when linking, there is no guarantee the compiler will inline the function (which is crucial for performance).