Change GuestTracking to use a BTreeSet to remove performance bottleneck (#3112)

Co-authored-by: Your Name <[email protected]>

Author: WorksButNotTested

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,41 +66,29 @@ 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(())
     }
 
     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),
 }

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
             ))

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(()));