From 36c748100ea14314057968b262d0c4bc739d3687 Mon Sep 17 00:00:00 2001
From: WorksButNotTested <62701594+WorksButNotTested@users.noreply.github.com>
Date: Fri, 28 Mar 2025 19:13:34 +0000
Subject: [PATCH] Change GuestTracking to use a BTreeSet to remove performance
 bottleneck (#3112)

Co-authored-by: Your Name <you@example.com>
---
 .../librasan/asan/src/tracking/guest.rs       | 40 +++++++------------
 .../librasan/asan/tests/guest_tracking.rs     | 14 ++-----
 .../fuzz/fuzz_targets/guest_tracking.rs       |  4 +-
 3 files changed, 18 insertions(+), 40 deletions(-)

diff --git a/libafl_qemu/librasan/asan/src/tracking/guest.rs b/libafl_qemu/librasan/asan/src/tracking/guest.rs
index 9d3db62176..4dd8e6f38b 100644
--- a/libafl_qemu/librasan/asan/src/tracking/guest.rs
+++ b/libafl_qemu/librasan/asan/src/tracking/guest.rs
@@ -2,7 +2,7 @@
 //! This implementation performs guest memory tracking by use of a simple sorted
 //! list residing in the guest's user space. Hence no interaction with the host
 //! is required.
-use alloc::vec::Vec;
+use alloc::collections::BTreeSet;
 use core::cmp::Ordering;
 
 use log::debug;
@@ -44,7 +44,7 @@ impl Ord for Range {
     }
 }
 
-type Ranges = Vec<Range>;
+type Ranges = BTreeSet<Range>;
 
 #[derive(Debug)]
 pub struct GuestTracking {
@@ -66,20 +66,8 @@ impl Tracking for GuestTracking {
 
         let item = Range { start, len };
 
-        let pos = self.ranges.binary_search(&item);
-        match pos {
-            Ok(pos) => {
-                let conflict = &self.ranges[pos];
-                Err(GuestTrackingError::TrackingConflict(
-                    conflict.start,
-                    conflict.len,
-                    item.start,
-                    item.len,
-                ))?;
-            }
-            Err(pos) => {
-                self.ranges.insert(pos, item);
-            }
+        if !self.ranges.insert(item) {
+            Err(GuestTrackingError::TrackingConflict(start, len))?;
         }
 
         Ok(())
@@ -87,20 +75,20 @@ impl Tracking for GuestTracking {
 
     fn untrack(&mut self, start: GuestAddr) -> Result<(), Self::Error> {
         debug!("dealloc - start: 0x{:x}", start);
-        let pos = self.ranges.binary_search_by(|item| item.start.cmp(&start));
-        match pos {
-            Ok(pos) => {
-                self.ranges.remove(pos);
-                Ok(())
-            }
-            Err(_pos) => Err(GuestTrackingError::AllocationNotFound(start)),
+        let item = Range { start, len: 1 };
+
+        if !self.ranges.remove(&item) {
+            Err(GuestTrackingError::AllocationNotFound(start))?;
         }
+        Ok(())
     }
 }
 
 impl GuestTracking {
     pub fn new() -> Result<Self, GuestTrackingError> {
-        Ok(GuestTracking { ranges: Vec::new() })
+        Ok(GuestTracking {
+            ranges: BTreeSet::new(),
+        })
     }
 
     pub fn is_out_of_bounds(addr: GuestAddr, len: usize) -> bool {
@@ -118,8 +106,8 @@ pub enum GuestTrackingError {
     AddressRangeOverflow(GuestAddr, usize),
     #[error("Allocation not found: {0:x}")]
     AllocationNotFound(GuestAddr),
-    #[error("Tracking conflict")]
-    TrackingConflict(GuestAddr, usize, GuestAddr, usize),
+    #[error("Tracking conflict: {0:x}, len: {1:x}")]
+    TrackingConflict(GuestAddr, usize),
     #[error("Zero Length")]
     ZeroLength(GuestAddr),
 }
diff --git a/libafl_qemu/librasan/asan/tests/guest_tracking.rs b/libafl_qemu/librasan/asan/tests/guest_tracking.rs
index 8828a6ae16..36254bbd9f 100644
--- a/libafl_qemu/librasan/asan/tests/guest_tracking.rs
+++ b/libafl_qemu/librasan/asan/tests/guest_tracking.rs
@@ -46,9 +46,7 @@ mod tests {
         assert_eq!(tracking.track(0x1000, 0x1000), Ok(()));
         assert_eq!(
             tracking.track(0x1000, 0x1000),
-            Err(GuestTrackingError::TrackingConflict(
-                0x1000, 0x1000, 0x1000, 0x1000
-            ))
+            Err(GuestTrackingError::TrackingConflict(0x1000, 0x1000))
         );
     }
 
@@ -72,9 +70,7 @@ mod tests {
         assert_eq!(tracking.track(0x1000, 0x1000), Ok(()));
         assert_eq!(
             tracking.track(0x0000, 0x1001),
-            Err(GuestTrackingError::TrackingConflict(
-                0x1000, 0x1000, 0x0000, 0x1001
-            ))
+            Err(GuestTrackingError::TrackingConflict(0x0000, 0x1001))
         );
     }
 
@@ -84,9 +80,7 @@ mod tests {
         assert_eq!(tracking.track(0x1000, 0x1000), Ok(()));
         assert_eq!(
             tracking.track(0x1fff, 0x1001),
-            Err(GuestTrackingError::TrackingConflict(
-                0x1000, 0x1000, 0x1fff, 0x1001
-            ))
+            Err(GuestTrackingError::TrackingConflict(0x1fff, 0x1001))
         );
     }
 
@@ -100,8 +94,6 @@ mod tests {
         assert_eq!(
             tracking.track(0xffffffffb5b60107, 0xdb),
             Err(GuestTrackingError::TrackingConflict(
-                0xffffffffb5b5ff21,
-                0x3ff,
                 0xffffffffb5b60107,
                 0xdb
             ))
diff --git a/libafl_qemu/librasan/fuzz/fuzz_targets/guest_tracking.rs b/libafl_qemu/librasan/fuzz/fuzz_targets/guest_tracking.rs
index 807599e2cf..bed4d97e81 100644
--- a/libafl_qemu/librasan/fuzz/fuzz_targets/guest_tracking.rs
+++ b/libafl_qemu/librasan/fuzz/fuzz_targets/guest_tracking.rs
@@ -86,9 +86,7 @@ fuzz_target!(|data: Vec<GuestAddr>| {
         if overlaps {
             assert_eq!(
                 test_result,
-                Err(GuestTrackingError::TrackingConflict(
-                    start, len, test_start, test_len
-                ))
+                Err(GuestTrackingError::TrackingConflict(test_start, test_len))
             );
         } else {
             assert_eq!(test_result, Ok(()));