Optimize data structures used by librasan (#3227)

libafl_qemu/librasan/asan/Cargo.toml

@@ -78,6 +78,8 @@ spin = { version = "0.9.8", default-features = false, features = [
 ] }
 syscalls = { version = "0.6.18", default-features = false, optional = true }
 thiserror = { version = "2.0.11", default-features = false }
+ahash = { version = "0.8.12", default-features = false }
+hashbrown = { version = "0.15.3", default-features = false }
 
 [build-dependencies]
 cc = { version = "1.2.13" }

libafl_qemu/librasan/asan/src/allocator/frontend/default.rs

@@ -10,15 +10,20 @@
 //! re-used for a period of time.
 use alloc::{
     alloc::{GlobalAlloc, Layout, LayoutError},
-    collections::{BTreeMap, VecDeque},
+    collections::VecDeque,
     fmt::Debug,
 };
+use core::hash::BuildHasherDefault;
 #[cfg(feature = "initialize")]
 use core::ptr::write_bytes;
 
+use ahash::AHasher;
+use hashbrown::HashMap;
 use log::debug;
 use thiserror::Error;
 
+type Hasher = BuildHasherDefault<AHasher>;
+
 use crate::{
     GuestAddr,
     allocator::frontend::AllocatorFrontend,
@@ -38,7 +43,7 @@ pub struct DefaultFrontend<B: GlobalAlloc + Send, S: Shadow, T: Tracking> {
     shadow: S,
     tracking: T,
     red_zone_size: usize,
-    allocations: BTreeMap<GuestAddr, Allocation>,
+    allocations: HashMap<GuestAddr, Allocation, Hasher>,
     quarantine: VecDeque<Allocation>,
     quarantine_size: usize,
     quaratine_used: usize,
@@ -176,7 +181,7 @@ impl<B: GlobalAlloc + Send, S: Shadow, T: Tracking> DefaultFrontend<B, S, T> {
             shadow,
             tracking,
             red_zone_size,
-            allocations: BTreeMap::new(),
+            allocations: HashMap::with_capacity_and_hasher(4096, Hasher::default()),
             quarantine: VecDeque::new(),
             quarantine_size,
             quaratine_used: 0,

libafl_qemu/librasan/asan/src/patch/mod.rs

@@ -3,7 +3,7 @@
 //! to re-direct execution to an alternative address.
 pub mod raw;
 
-use alloc::{collections::BTreeMap, fmt::Debug};
+use alloc::fmt::Debug;
 
 use log::trace;
 use spin::{Mutex, Once};
@@ -20,8 +20,14 @@ pub trait Patch: Debug {
     fn patch(target: GuestAddr, destination: GuestAddr) -> Result<(), Self::Error>;
 }
 
+use core::hash::BuildHasherDefault;
+
+use ahash::AHasher;
+use hashbrown::HashMap;
+type Hasher = BuildHasherDefault<AHasher>;
+
 static PATCHES: Once<Mutex<Patches>> = Once::new();
-static PATCHED: Mutex<BTreeMap<GuestAddr, GuestAddr>> = Mutex::new(BTreeMap::new());
+static PATCHED: Once<Mutex<HashMap<GuestAddr, GuestAddr, Hasher>>> = Once::new();
 
 pub struct Patches {
     maps: Maps,
@@ -31,6 +37,8 @@ impl Patches {
     pub fn init(maps: Maps) {
         let patches = Mutex::new(Patches { maps });
         PATCHES.call_once(|| patches);
+        let patched = Mutex::new(HashMap::with_capacity_and_hasher(4096, Hasher::default()));
+        PATCHED.call_once(|| patched);
     }
 
     pub fn apply<P: Patch, M: Mmap>(
@@ -45,12 +53,15 @@ impl Patches {
             .map_err(PatchesError::MapsError)?;
         P::patch(target, destination).map_err(|e| PatchesError::PatchError(e))?;
         drop(prot);
-        PATCHED.lock().insert(target, destination);
+        let mut patched = PATCHED.get().ok_or(PatchesError::Uninitialized())?.lock();
+        patched.insert(target, destination);
         Ok(())
     }
 
     pub fn is_patched(addr: GuestAddr) -> bool {
-        PATCHED.lock().contains_key(&addr)
+        PATCHED
+            .get()
+            .map_or(false, |p| p.lock().contains_key(&addr))
     }
 }
 

libafl_qemu/librasan/asan/src/tracking/guest_fast.rs

@@ -0,0 +1,78 @@
+//! # guest
+//! This implementation performs guest memory tracking by use of a hash table
+//! and hence requires no interaction with the guest. Unlike GuestTracking, this
+//! faster variant is unable to detect whether a new allocation overlaps an
+//! existing one (though this should be taken care of by the allocator).
+use core::hash::BuildHasherDefault;
+
+use ahash::AHasher;
+use hashbrown::HashSet;
+use log::debug;
+use thiserror::Error;
+
+use crate::{GuestAddr, tracking::Tracking};
+
+type Hasher = BuildHasherDefault<AHasher>;
+
+#[derive(Debug)]
+pub struct GuestFastTracking {
+    ranges: HashSet<GuestAddr, Hasher>,
+}
+
+impl Tracking for GuestFastTracking {
+    type Error = GuestFastTrackingError;
+
+    fn track(&mut self, start: GuestAddr, len: usize) -> Result<(), Self::Error> {
+        debug!("alloc - start: 0x{:x}, len: 0x{:x}", start, len);
+        if Self::is_out_of_bounds(start, len) {
+            Err(GuestFastTrackingError::AddressRangeOverflow(start, len))?;
+        }
+
+        if len == 0 {
+            Err(GuestFastTrackingError::ZeroLength(start))?;
+        }
+
+        if !self.ranges.insert(start) {
+            Err(GuestFastTrackingError::TrackingConflict(start, len))?;
+        }
+
+        Ok(())
+    }
+
+    fn untrack(&mut self, start: GuestAddr) -> Result<(), Self::Error> {
+        debug!("dealloc - start: 0x{:x}", start);
+
+        if !self.ranges.remove(&start) {
+            Err(GuestFastTrackingError::AllocationNotFound(start))?;
+        }
+        Ok(())
+    }
+}
+
+impl GuestFastTracking {
+    pub fn new() -> Result<Self, GuestFastTrackingError> {
+        Ok(GuestFastTracking {
+            ranges: HashSet::with_capacity_and_hasher(4096, Hasher::default()),
+        })
+    }
+
+    pub fn is_out_of_bounds(addr: GuestAddr, len: usize) -> bool {
+        if len == 0 {
+            false
+        } else {
+            GuestAddr::MAX - len + 1 < addr
+        }
+    }
+}
+
+#[derive(Error, Debug, PartialEq)]
+pub enum GuestFastTrackingError {
+    #[error("Address overflow: {0:x}, len: {1:x}")]
+    AddressRangeOverflow(GuestAddr, usize),
+    #[error("Allocation not found: {0:x}")]
+    AllocationNotFound(GuestAddr),
+    #[error("Tracking conflict: {0:x}, len: {1:x}")]
+    TrackingConflict(GuestAddr, usize),
+    #[error("Zero Length")]
+    ZeroLength(GuestAddr),
+}

libafl_qemu/librasan/asan/src/tracking/mod.rs

@@ -8,6 +8,8 @@ use crate::GuestAddr;
 
 #[cfg(feature = "guest")]
 pub mod guest;
+#[cfg(feature = "guest")]
+pub mod guest_fast;
 #[cfg(feature = "host")]
 pub mod host;
 

libafl_qemu/librasan/gasan/src/lib.rs

@@ -22,7 +22,7 @@ use asan::{
         Symbols,
         dlsym::{DlSymSymbols, LookupTypeNext},
     },
-    tracking::{Tracking, guest::GuestTracking},
+    tracking::{Tracking, guest_fast::GuestFastTracking},
 };
 use log::{Level, debug, trace};
 use spin::{Lazy, mutex::Mutex};
@@ -34,7 +34,7 @@ type GasanMmap = LibcMmap<Syms>;
 type GasanBackend = MimallocBackend<DlmallocBackend<GasanMmap>>;
 
 pub type GasanFrontend =
-    DefaultFrontend<GasanBackend, GuestShadow<GasanMmap, DefaultShadowLayout>, GuestTracking>;
+    DefaultFrontend<GasanBackend, GuestShadow<GasanMmap, DefaultShadowLayout>, GuestFastTracking>;
 
 pub type GasanSyms = DlSymSymbols<LookupTypeNext>;
 
@@ -45,7 +45,7 @@ static FRONTEND: Lazy<Mutex<GasanFrontend>> = Lazy::new(|| {
     debug!("init");
     let backend = GasanBackend::new(DlmallocBackend::new(PAGE_SIZE));
     let shadow = GuestShadow::<GasanMmap, DefaultShadowLayout>::new().unwrap();
-    let tracking = GuestTracking::new().unwrap();
+    let tracking = GuestFastTracking::new().unwrap();
     let frontend = GasanFrontend::new(
         backend,
         shadow,

libafl_qemu/librasan/zasan/src/lib.rs

@@ -16,7 +16,7 @@ use asan::{
         guest::{DefaultShadowLayout, GuestShadow},
     },
     symbols::{Symbols, nop::NopSymbols},
-    tracking::{Tracking, guest::GuestTracking},
+    tracking::{Tracking, guest_fast::GuestFastTracking},
 };
 use log::{Level, trace};
 use spin::{Lazy, Mutex};
@@ -24,7 +24,7 @@ use spin::{Lazy, Mutex};
 pub type ZasanFrontend = DefaultFrontend<
     DlmallocBackend<LinuxMmap>,
     GuestShadow<LinuxMmap, DefaultShadowLayout>,
-    GuestTracking,
+    GuestFastTracking,
 >;
 
 pub type ZasanSyms = NopSymbols;
@@ -35,7 +35,7 @@ static FRONTEND: Lazy<Mutex<ZasanFrontend>> = Lazy::new(|| {
     LinuxLogger::initialize(Level::Info);
     let backend = DlmallocBackend::<LinuxMmap>::new(PAGE_SIZE);
     let shadow = GuestShadow::<LinuxMmap, DefaultShadowLayout>::new().unwrap();
-    let tracking = GuestTracking::new().unwrap();
+    let tracking = GuestFastTracking::new().unwrap();
     let frontend = ZasanFrontend::new(
         backend,
         shadow,