From 5aa3270f82a91a61f027de93a1d24cf164279eb5 Mon Sep 17 00:00:00 2001 From: Christian Deacon Date: Sun, 23 Mar 2025 20:35:08 -0400 Subject: [PATCH] Implement both IP and flow-based rate limiting. --- src/common/config.h | 28 ++++-- src/common/types.h | 29 ++++-- src/loader/utils/config.c | 75 ++++++++++---- src/loader/utils/config.h | 7 +- src/loader/utils/logging.c | 2 +- src/loader/utils/xdp.c | 30 ++++-- src/rule_add/prog.c | 25 +++-- src/rule_add/utils/cli.c | 21 +++- src/rule_add/utils/cli.h | 7 +- src/xdp/prog.c | 89 +++++++++++++---- src/xdp/utils/logging.c | 15 ++- src/xdp/utils/logging.h | 2 +- src/xdp/utils/maps.h | 40 +++++--- src/xdp/utils/rl.c | 193 ++++++++++++++++++++++++++++--------- src/xdp/utils/rl.h | 6 +- 15 files changed, 420 insertions(+), 149 deletions(-) diff --git a/src/common/config.h b/src/common/config.h index b14a2cd..045402d 100644 --- a/src/common/config.h +++ b/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 \ No newline at end of file +#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 \ No newline at end of file diff --git a/src/common/types.h b/src/common/types.h index 23ada6a..0267412 100644 --- a/src/common/types.h +++ b/src/common/types.h @@ -108,11 +108,21 @@ struct filter u8 action; u16 block_time; - unsigned int do_pps : 1; - u64 pps; +#ifdef ENABLE_RL_IP + unsigned int do_ip_pps : 1; + u64 ip_pps; - unsigned int do_bps : 1; - u64 bps; + unsigned int do_ip_bps : 1; + 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 bps; + u64 ip_pps; + u64 ip_bps; + + u64 flow_pps; + u64 flow_bps; } typedef filter_log_event_t; struct lpm_trie_key diff --git a/src/loader/utils/config.c b/src/loader/utils/config.c index 52bd412..d300bcd 100644 --- a/src/loader/utils/config.c +++ b/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). - s64 pps; + // IP PPS (not required). + 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). - s64 bps; + // IP BPS (not required). + 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. - if (filter->pps > -1) + // Add IP PPS. + if (filter->ip_pps > -1) { - config_setting_t* pps = config_setting_add(filter_cfg, "pps", CONFIG_TYPE_INT64); - config_setting_set_int64(pps, filter->pps); + config_setting_t* pps = config_setting_add(filter_cfg, "ip_pps", CONFIG_TYPE_INT64); + config_setting_set_int64(pps, filter->ip_pps); } - // Add BPS. - if (filter->bps > -1) + // Add IP BPS. + if (filter->ip_bps > -1) { - config_setting_t* bps = config_setting_add(filter_cfg, "bps", CONFIG_TYPE_INT64); - config_setting_set_int64(bps, filter->bps); + config_setting_t* bps = config_setting_add(filter_cfg, "ip_bps", CONFIG_TYPE_INT64); + 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->bps = -1; + filter->ip_pps = -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\tBPS => %lld\n\n", filter->bps); + printf("\t\t\tIP PPS => %lld\n", filter->ip_pps); + 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"); diff --git a/src/loader/utils/config.h b/src/loader/utils/config.h index 68bc804..2162934 100644 --- a/src/loader/utils/config.h +++ b/src/loader/utils/config.h @@ -72,8 +72,11 @@ struct filter_rule_cfg int action; int block_time; - s64 pps; - s64 bps; + s64 ip_pps; + s64 ip_bps; + + s64 flow_pps; + s64 flow_bps; filter_rule_ip_opts_t ip; diff --git a/src/loader/utils/logging.c b/src/loader/utils/logging.c index 9bdf387..f300793 100644 --- a/src/loader/utils/logging.c +++ b/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; } \ No newline at end of file diff --git a/src/loader/utils/xdp.c b/src/loader/utils/xdp.c index 4e762cb..3c9e51f 100644 --- a/src/loader/utils/xdp.c +++ b/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) { diff --git a/src/rule_add/prog.c b/src/rule_add/prog.c index 2e8d389..6ab68e1 100644 --- a/src/rule_add/prog.c +++ b/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.bps = -1; + cli.ip_pps = -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) diff --git a/src/rule_add/utils/cli.c b/src/rule_add/utils/cli.c index 636995f..cd168a4 100644 --- a/src/rule_add/utils/cli.c +++ b/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 }, - { "bps", required_argument, NULL, 10 }, + { "ip-pps", required_argument, NULL, 9 }, + { "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; diff --git a/src/rule_add/utils/cli.h b/src/rule_add/utils/cli.h index f2eecf5..13ca19e 100644 --- a/src/rule_add/utils/cli.h +++ b/src/rule_add/utils/cli.h @@ -36,8 +36,11 @@ struct cli char* src_ip6; char* dst_ip6; - s64 pps; - s64 bps; + s64 ip_pps; + s64 ip_bps; + + s64 flow_pps; + s64 flow_bps; int min_ttl; int max_ttl; diff --git a/src/xdp/prog.c b/src/xdp/prog.c index 184d773..cea4f77 100644 --- a/src/xdp/prog.c +++ b/src/xdp/prog.c @@ -285,18 +285,41 @@ int xdp_prog_main(struct xdp_md *ctx) } } +#ifdef ENABLE_FILTERS // Update client stats (PPS/BPS). - u64 pps = 0; - u64 bps = 0; - - if (iph6) +#ifdef ENABLE_RL_IP + u64 ip_pps = 0; + u64 ip_bps = 0; +#endif + +#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; @@ -312,6 +335,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 +475,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 +621,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 diff --git a/src/xdp/utils/logging.c b/src/xdp/utils/logging.c index 43cef33..aa7590b 100644 --- a/src/xdp/utils/logging.c +++ b/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 bps The current BPS rate. + * @param ip_pps The current IP PPS 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->bps = bps; + e->ip_pps = ip_pps; + e->ip_bps = ip_bps; + + e->flow_pps = flow_pps; + e->flow_bps = flow_bps; e->length = pkt_len; diff --git a/src/xdp/utils/logging.h b/src/xdp/utils/logging.h index 34301dd..c174556 100644 --- a/src/xdp/utils/logging.h +++ b/src/xdp/utils/logging.h @@ -6,7 +6,7 @@ #include #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). diff --git a/src/xdp/utils/maps.h b/src/xdp/utils/maps.h index 57278c9..ab3773e 100644 --- a/src/xdp/utils/maps.h +++ b/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); -#ifdef USE_FLOW_RL - __type(key, flow_t); -#else + __uint(max_entries, MAX_RL_IP); __type(key, u32); -#endif - __type(value, ip_stats_t); + __type(value, cl_stats_t); } map_ip_stats SEC(".maps"); struct { __uint(type, BPF_MAP_TYPE_LRU_HASH); - __uint(max_entries, MAX_TRACK_IPS); -#ifdef USE_FLOW_RL - __type(key, flow6_t); -#else + __uint(max_entries, MAX_RL_IP); __type(key, u128); -#endif - __type(value, ip_stats_t); + __type(value, cl_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 { diff --git a/src/xdp/utils/rl.c b/src/xdp/utils/rl.c index 103e557..50d169f 100644 --- a/src/xdp/utils/rl.c +++ b/src/xdp/utils/rl.c @@ -1,8 +1,116 @@ #include #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); -#else - ip_stats_t *ip_stats = bpf_map_lookup_elem(&map_ip_stats, &ip); -#endif + cl_stats_t* stats = bpf_map_lookup_elem(&map_flow_stats, &key); - if (ip_stats) + if (stats) { // Check for next update. - if (now > ip_stats->next_update) + if (now > stats->next_update) { - ip_stats->pps = 1; - ip_stats->bps = pkt_len; - ip_stats->next_update = now + NANO_TO_SEC; + 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(&ip_stats->pps, 1); - __sync_fetch_and_add(&ip_stats->bps, pkt_len); + __sync_fetch_and_add(&stats->pps, 1); + __sync_fetch_and_add(&stats->bps, pkt_len); } - *pps = ip_stats->pps; - *bps = ip_stats->bps; + *pps = stats->pps; + *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_ip_stats, &key, &new, BPF_ANY); -#else - bpf_map_update_elem(&map_ip_stats, &ip, &new, BPF_ANY); -#endif + bpf_map_update_elem(&map_flow_stats, &key, &new, BPF_ANY); } + + 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); -#else - ip_stats_t *ip_stats = bpf_map_lookup_elem(&map_ip6_stats, ip); -#endif + cl_stats_t* stats = bpf_map_lookup_elem(&map_flow6_stats, &key); - if (ip_stats) + if (stats) { // Check for next update. - if (now > ip_stats->next_update) + if (now > stats->next_update) { - ip_stats->pps = 1; - ip_stats->bps = pkt_len; - ip_stats->next_update = now + NANO_TO_SEC; + 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(&ip_stats->pps, 1); - __sync_fetch_and_add(&ip_stats->bps, pkt_len); + __sync_fetch_and_add(&stats->pps, 1); + __sync_fetch_and_add(&stats->bps, pkt_len); } - *pps = ip_stats->pps; - *bps = ip_stats->bps; + *pps = stats->pps; + *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_ip6_stats, &key, &new, BPF_ANY); -#else - bpf_map_update_elem(&map_ip6_stats, ip, &new, BPF_ANY); -#endif + bpf_map_update_elem(&map_flow6_stats, &key, &new, BPF_ANY); } + + return 0; } +#endif #endif \ No newline at end of file diff --git a/src/xdp/utils/rl.h b/src/xdp/utils/rl.h index 857a8f0..abc5f80 100644 --- a/src/xdp/utils/rl.h +++ b/src/xdp/utils/rl.h @@ -7,8 +7,10 @@ #include #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 void update_ip6_stats(u64 *pps, u64 *bps, u128 *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 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).