diff --git a/libafl_qemu/librasan/asan/Cargo.toml b/libafl_qemu/librasan/asan/Cargo.toml index 3ad6707e79..36f06165a7 100644 --- a/libafl_qemu/librasan/asan/Cargo.toml +++ b/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" } diff --git a/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs b/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs index 58e8ab8b41..c09a38de7d 100644 --- a/libafl_qemu/librasan/asan/src/allocator/frontend/default.rs +++ b/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; + use crate::{ GuestAddr, allocator::frontend::AllocatorFrontend, @@ -38,7 +43,7 @@ pub struct DefaultFrontend { shadow: S, tracking: T, red_zone_size: usize, - allocations: BTreeMap, + allocations: HashMap, quarantine: VecDeque, quarantine_size: usize, quaratine_used: usize, @@ -176,7 +181,7 @@ impl DefaultFrontend { 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, diff --git a/libafl_qemu/librasan/asan/src/patch/mod.rs b/libafl_qemu/librasan/asan/src/patch/mod.rs index 112ce590a0..ad18292206 100644 --- a/libafl_qemu/librasan/asan/src/patch/mod.rs +++ b/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; + static PATCHES: Once> = Once::new(); -static PATCHED: Mutex> = Mutex::new(BTreeMap::new()); +static PATCHED: Once>> = 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( @@ -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)) } } diff --git a/libafl_qemu/librasan/asan/src/tracking/guest_fast.rs b/libafl_qemu/librasan/asan/src/tracking/guest_fast.rs new file mode 100644 index 0000000000..9f435e7135 --- /dev/null +++ b/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; + +#[derive(Debug)] +pub struct GuestFastTracking { + ranges: HashSet, +} + +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 { + 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), +} diff --git a/libafl_qemu/librasan/asan/src/tracking/mod.rs b/libafl_qemu/librasan/asan/src/tracking/mod.rs index 115c7cbb60..c4ce9dcce4 100644 --- a/libafl_qemu/librasan/asan/src/tracking/mod.rs +++ b/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; diff --git a/libafl_qemu/librasan/gasan/src/lib.rs b/libafl_qemu/librasan/gasan/src/lib.rs index f4d4641cad..0dd38ca208 100644 --- a/libafl_qemu/librasan/gasan/src/lib.rs +++ b/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; type GasanBackend = MimallocBackend>; pub type GasanFrontend = - DefaultFrontend, GuestTracking>; + DefaultFrontend, GuestFastTracking>; pub type GasanSyms = DlSymSymbols; @@ -45,7 +45,7 @@ static FRONTEND: Lazy> = Lazy::new(|| { debug!("init"); let backend = GasanBackend::new(DlmallocBackend::new(PAGE_SIZE)); let shadow = GuestShadow::::new().unwrap(); - let tracking = GuestTracking::new().unwrap(); + let tracking = GuestFastTracking::new().unwrap(); let frontend = GasanFrontend::new( backend, shadow, diff --git a/libafl_qemu/librasan/zasan/src/lib.rs b/libafl_qemu/librasan/zasan/src/lib.rs index 355e18cc46..bf5c4bb5b3 100644 --- a/libafl_qemu/librasan/zasan/src/lib.rs +++ b/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, GuestShadow, - GuestTracking, + GuestFastTracking, >; pub type ZasanSyms = NopSymbols; @@ -35,7 +35,7 @@ static FRONTEND: Lazy> = Lazy::new(|| { LinuxLogger::initialize(Level::Info); let backend = DlmallocBackend::::new(PAGE_SIZE); let shadow = GuestShadow::::new().unwrap(); - let tracking = GuestTracking::new().unwrap(); + let tracking = GuestFastTracking::new().unwrap(); let frontend = ZasanFrontend::new( backend, shadow,