From e9e9027fe79bbb7f466a7985d103a4366aee7b1b Mon Sep 17 00:00:00 2001
From: Christian Deacon <christianmdeacon@gmail.com>
Date: Thu, 13 Mar 2025 08:08:22 -0400
Subject: [PATCH] Fix issues with IP range drop map.

---
 src/common/types.h      |  2 +-
 src/loader/utils/xdp.c  |  8 ++++----
 src/xdp/prog.c          |  4 ++--
 src/xdp/utils/helpers.c | 10 +++-------
 src/xdp/utils/helpers.h |  2 ++
 src/xdp/utils/maps.h    |  5 +++--
 6 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/src/common/types.h b/src/common/types.h
index 67cfd9a..c02477b 100644
--- a/src/common/types.h
+++ b/src/common/types.h
@@ -164,4 +164,4 @@ struct lpm_trie_key
 {
     u32 prefix_len;
     u32 data;
-} typedef LpmTrieKey;
\ No newline at end of file
+} typedef lpm_trie_key_t;
\ No newline at end of file
diff --git a/src/loader/utils/xdp.c b/src/loader/utils/xdp.c
index f892a7b..a736d12 100644
--- a/src/loader/utils/xdp.c
+++ b/src/loader/utils/xdp.c
@@ -628,10 +628,10 @@ int add_block6(int map_block6, u128 ip, u64 expires)
  */
 int delete_range_drop(int map_range_drop, u32 net, u8 cidr)
 {
-    u32 bit_mask = ( ~( (1 << (32 - cidr) ) - 1) );
+    u32 bit_mask = htonl(( ~( (1 << (32 - cidr) ) - 1) ));
     u32 start = net & bit_mask;
 
-    LpmTrieKey key = {0};
+    lpm_trie_key_t key = {0};
     key.prefix_len = cidr;
     key.data = start;
 
@@ -649,10 +649,10 @@ int delete_range_drop(int map_range_drop, u32 net, u8 cidr)
  */
 int add_range_drop(int map_range_drop, u32 net, u8 cidr)
 {
-    u32 bit_mask = ( ~( (1 << (32 - cidr) ) - 1) );
+    u32 bit_mask = htonl(( ~( (1 << (32 - cidr) ) - 1) ));
     u32 start = net & bit_mask;
 
-    LpmTrieKey key = {0};
+    lpm_trie_key_t key = {0};
     key.prefix_len = cidr;
     key.data = start;
 
diff --git a/src/xdp/prog.c b/src/xdp/prog.c
index 2016629..d88461d 100644
--- a/src/xdp/prog.c
+++ b/src/xdp/prog.c
@@ -10,13 +10,13 @@
 
 #include <common/all.h>
 
-#include <xdp/utils/maps.h>
-
 #include <xdp/utils/rl.h>
 #include <xdp/utils/logging.h>
 #include <xdp/utils/stats.h>
 #include <xdp/utils/helpers.h>
 
+#include <xdp/utils/maps.h>
+
 struct 
 {
     __uint(priority, 10);
diff --git a/src/xdp/utils/helpers.c b/src/xdp/utils/helpers.c
index 98b03bc..e7e1b3d 100644
--- a/src/xdp/utils/helpers.c
+++ b/src/xdp/utils/helpers.c
@@ -1,7 +1,5 @@
 #include <xdp/utils/helpers.h>
 
-#include <xdp/utils/maps.h>
-
 /**
  * Checks if an IP is within a specific CIDR range.
  * 
@@ -26,11 +24,9 @@ static __always_inline int is_ip_in_range(u32 src_ip, u32 net_ip, u8 cidr)
  */
 static __always_inline int check_ip_range_drop(u32 ip)
 {
-    LpmTrieKey key =
-    {
-        .prefix_len = 32,
-        .data = ip
-    };
+    lpm_trie_key_t key = {0};
+    key.prefix_len = 32;
+    key.data = ip;
 
     u64 *lookup = bpf_map_lookup_elem(&map_range_drop, &key);
 
diff --git a/src/xdp/utils/helpers.h b/src/xdp/utils/helpers.h
index 87269c1..7c27982 100644
--- a/src/xdp/utils/helpers.h
+++ b/src/xdp/utils/helpers.h
@@ -13,6 +13,8 @@
 #include <bpf/bpf_helpers.h>
 #endif
 
+#include <xdp/utils/maps.h>
+
 #define likely(x) __builtin_expect(!!(x), 1)
 #define unlikely(x) __builtin_expect(!!(x), 0)
 
diff --git a/src/xdp/utils/maps.h b/src/xdp/utils/maps.h
index a1eebfe..57278c9 100644
--- a/src/xdp/utils/maps.h
+++ b/src/xdp/utils/maps.h
@@ -30,11 +30,12 @@ struct
 } map_block6 SEC(".maps");
 
 #ifdef ENABLE_IP_RANGE_DROP
-struct {
+struct
+{
     __uint(type, BPF_MAP_TYPE_LPM_TRIE);
     __uint(max_entries, MAX_IP_RANGES);
     __uint(map_flags, BPF_F_NO_PREALLOC);
-    __type(key, LpmTrieKey);
+    __type(key, lpm_trie_key_t);
     __type(value, u64);
 } map_range_drop SEC(".maps");
 #endif