From 84b82662c497a46c9ab1ce3ec2c27de3f90676fc Mon Sep 17 00:00:00 2001
From: Francisco Javier Honduvilla Coto <javierhonduco@gmail.com>
Date: Wed, 1 May 2024 15:28:29 +0100
Subject: [PATCH] tests: Fix off-by-32 bits in LPM tree

The integration tests started failing starting with kernel 6.8, and
after some checks, it pointed to this commit
https://github.com/torvalds/linux/commit/9b75dbeb36fcd9fc7ed51d370310d0518a387769.

The tests fail now because the prefix length we pass is 32 bits higher
(see initialiser) than the maximum prefix length. After this commit was
introduced, any prefix length that's higher than the key size will
return NULL earlier, rather than returning the value with the longest
match, even if it's beyond the size the key can encode.

Test Plan
=========

Ran the tests and some profiles on a box running 6.8.7-200 (Fedora) without issues.

Also added an assertion to ensure that the prefix_len is not bigger than the LPM trie
key size.
---
 src/bpf/profiler_bindings.rs | 12 ++++++++++--
 src/profiler.rs              |  2 +-
 src/util.rs                  | 21 +++++++++++++--------
 3 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/bpf/profiler_bindings.rs b/src/bpf/profiler_bindings.rs
index 43e22a7..c9c7661 100644
--- a/src/bpf/profiler_bindings.rs
+++ b/src/bpf/profiler_bindings.rs
@@ -16,9 +16,17 @@ unsafe impl Plain for exec_mappings_key {}
 unsafe impl Plain for mapping_t {}
 
 impl exec_mappings_key {
-    pub fn new(pid: u32, address: u64, prefix: u32) -> Self {
+    pub fn new(pid: u32, address: u64, prefix_len: u32) -> Self {
+        let key_size_bits = std::mem::size_of::<Self>() * 8;
+        assert!(
+            prefix_len <= key_size_bits.try_into().unwrap(),
+            "prefix_len {} should be <= than the size of exec_mappings_key {}",
+            prefix_len,
+            key_size_bits
+        );
+
         Self {
-            prefix_len: 32 + prefix,
+            prefix_len,
             pid: pid.to_be(),
             data: address.to_be(),
         }
diff --git a/src/profiler.rs b/src/profiler.rs
index eb1dd51..76b1419 100644
--- a/src/profiler.rs
+++ b/src/profiler.rs
@@ -924,7 +924,7 @@ impl Profiler<'_> {
                     let key = exec_mappings_key::new(
                         pid.try_into().unwrap(),
                         address_range.addr,
-                        address_range.range,
+                        32 + address_range.prefix_len,
                     );
 
                     self.add_bpf_mapping(&key, &mapping).unwrap();
diff --git a/src/util.rs b/src/util.rs
index 4875c44..1d5f01b 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -1,7 +1,7 @@
 #[derive(Debug, PartialEq)]
 pub struct AddressBlockRange {
     pub addr: u64,
-    pub range: u32,
+    pub prefix_len: u32,
 }
 
 /// For a given address range, calculate all the prefix ranges to ensure searching
@@ -22,7 +22,7 @@ pub fn summarize_address_range(low: u64, high: u64) -> Vec<AddressBlockRange> {
         );
         res.push(AddressBlockRange {
             addr: curr,
-            range: 64 - number_of_bits,
+            prefix_len: 64 - number_of_bits,
         });
         curr += 1 << number_of_bits;
         if curr - 1 == u64::MAX {
@@ -51,18 +51,21 @@ mod tests {
         assert_eq!(
             summarize_address_range(0, 100),
             vec![
-                AddressBlockRange { addr: 0, range: 58 },
+                AddressBlockRange {
+                    addr: 0,
+                    prefix_len: 58
+                },
                 AddressBlockRange {
                     addr: 64,
-                    range: 59
+                    prefix_len: 59
                 },
                 AddressBlockRange {
                     addr: 96,
-                    range: 62
+                    prefix_len: 62
                 },
                 AddressBlockRange {
                     addr: 100,
-                    range: 64
+                    prefix_len: 64
                 }
             ]
         );
@@ -100,7 +103,8 @@ mod tests {
         assert!(mapping1.end < mapping2.begin);
 
         for address_range in summarize_address_range(mapping1.begin, mapping1.end - 1) {
-            let key = exec_mappings_key::new(510530, address_range.addr, address_range.range);
+            let key =
+                exec_mappings_key::new(510530, address_range.addr, 32 + address_range.prefix_len);
             map.update(
                 unsafe { plain::as_bytes(&key) },
                 unsafe { plain::as_bytes(&mapping1) },
@@ -110,7 +114,8 @@ mod tests {
         }
 
         for address_range in summarize_address_range(mapping2.begin, mapping2.end - 1) {
-            let key = exec_mappings_key::new(510530, address_range.addr, address_range.range);
+            let key =
+                exec_mappings_key::new(510530, address_range.addr, 32 + address_range.prefix_len);
             map.update(
                 unsafe { plain::as_bytes(&key) },
                 unsafe { plain::as_bytes(&mapping2) },