diff --git a/libafl_targets/Cargo.toml b/libafl_targets/Cargo.toml index 21ddd34336..962a0143b4 100644 --- a/libafl_targets/Cargo.toml +++ b/libafl_targets/Cargo.toml @@ -26,6 +26,7 @@ sancov_8bit = [] sancov_cmplog = [] sancov_pcguard = ["sancov_pcguard_hitcounts"] clippy = [] # Ignore compiler warnings during clippy +observers = ["intervaltree", "ahash"] [build-dependencies] cc = { version = "1.0", features = ["parallel"] } @@ -36,4 +37,6 @@ log = "0.4.17" rangemap = "1.0" serde = { version = "1.0", default-features = false, features = ["alloc"] } # serialization lib +intervaltree = { version = "0.2.7", default-features = false, features = ["serde"], optional = true } +ahash = { version = "0.8.3", default-features = false, optional = true } # serde-big-array = "0.3.2" diff --git a/libafl_targets/src/sancov_8bit.rs b/libafl_targets/src/sancov_8bit.rs index 7cb7e4e5cf..1a745bcb96 100644 --- a/libafl_targets/src/sancov_8bit.rs +++ b/libafl_targets/src/sancov_8bit.rs @@ -19,3 +19,302 @@ pub extern "C" fn __sanitizer_cov_8bit_counters_init(start: *mut u8, stop: *mut )); } } + +#[cfg(feature = "observers")] +pub use self::observers::{counters_maps_observer, CountersMultiMapObserver}; + +#[cfg(feature = "observers")] +mod observers { + use alloc::{ + string::{String, ToString}, + vec::Vec, + }; + use core::{ + fmt::Debug, + hash::{BuildHasher, Hasher}, + iter::Flatten, + slice::{from_raw_parts, Iter, IterMut}, + }; + + use ahash::RandomState; + use intervaltree::IntervalTree; + use libafl::{ + bolts::{ + ownedref::OwnedMutSlice, tuples::Named, AsIter, AsIterMut, AsMutSlice, AsSlice, HasLen, + }, + inputs::UsesInput, + observers::{DifferentialObserver, MapObserver, Observer, ObserversTuple}, + Error, + }; + use serde::{Deserialize, Serialize}; + + use super::COUNTERS_MAPS; + + #[must_use] + #[export_name = "counters_maps_observer"] + /// Create a new [`CountersMultiMapObserver`] of the [`COUNTERS_MAP`]. + /// + /// This is a special [`MultiMapObserver`] for the [`COUNTERS_MAP`] and may be used when + /// 8-bit counters are used for `SanitizerCoverage`. You can utilize this observer in a + /// [`HitcountsIterableMapObserver`] like so: + /// + /// ```rust,ignore + /// use libafl::{ + /// observers::HitcountsIterableMapObserver, + /// feedbacks::MaxMapFeedback, + /// }; + /// use libafl_targets::sancov_8bit::counters_maps_observer; + /// + /// let counters_maps_observer = unsafe { counters_maps_observer("counters-maps") }; + /// let counters_maps_hitcounts_observer = HitcountsIterableMapObserver::new(counters_maps_observer); + /// let counters_maps_feedback = MaxMapFeedback::new(&counters_maps_hitcounts_observer); + /// ``` + /// + /// # Safety + /// + /// This function instantiates an observer of a `static mut` map whose contents are mutated by + /// `SanitizerCoverage` instrumentation. This is unsafe, and data in the map may be mutated from + /// under us at any time. It should never be assumed constant. + pub unsafe fn counters_maps_observer(name: &'static str) -> CountersMultiMapObserver { + CountersMultiMapObserver::new(name) + } + + /// The [`CountersMultiMapObserver`] observes all the counters that may be set by + /// `SanitizerCoverage` in [`COUNTERS_MAPS`] + #[derive(Serialize, Deserialize, Debug)] + #[allow(clippy::unsafe_derive_deserialize)] + pub struct CountersMultiMapObserver { + intervals: IntervalTree, + len: usize, + initial: u8, + name: String, + iter_idx: usize, + } + + impl Observer for CountersMultiMapObserver + where + S: UsesInput, + Self: MapObserver, + { + #[inline] + fn pre_exec(&mut self, _state: &mut S, _input: &S::Input) -> Result<(), Error> { + self.reset_map() + } + } + + impl Observer for CountersMultiMapObserver + where + S: UsesInput, + Self: MapObserver, + { + // in differential mode, we are *not* responsible for resetting the map! + } + + impl Named for CountersMultiMapObserver { + #[inline] + fn name(&self) -> &str { + self.name.as_str() + } + } + + impl HasLen for CountersMultiMapObserver { + #[inline] + fn len(&self) -> usize { + self.len + } + } + + impl MapObserver for CountersMultiMapObserver { + type Entry = u8; + + #[inline] + fn get(&self, idx: usize) -> &u8 { + let elem = self.intervals.query_point(idx).next().unwrap(); + let i = elem.value; + let j = idx - elem.range.start; + unsafe { &COUNTERS_MAPS[i].as_slice()[j] } + } + + #[inline] + fn get_mut(&mut self, idx: usize) -> &mut u8 { + let elem = self.intervals.query_point(idx).next().unwrap(); + let i = elem.value; + let j = idx - elem.range.start; + unsafe { &mut COUNTERS_MAPS[i].as_mut_slice()[j] } + } + + #[inline] + fn initial(&self) -> u8 { + self.initial + } + + fn count_bytes(&self) -> u64 { + let initial = self.initial(); + let mut res = 0; + for map in unsafe { &COUNTERS_MAPS } { + for x in map.as_slice() { + if *x != initial { + res += 1; + } + } + } + res + } + + fn hash(&self) -> u64 { + let mut hasher = RandomState::with_seeds(0, 0, 0, 0).build_hasher(); + for map in unsafe { &COUNTERS_MAPS } { + let slice = map.as_slice(); + let ptr = slice.as_ptr() as *const u8; + let map_size = slice.len() / core::mem::size_of::(); + unsafe { + hasher.write(from_raw_parts(ptr, map_size)); + } + } + hasher.finish() + } + + fn reset_map(&mut self) -> Result<(), Error> { + let initial = self.initial(); + for map in unsafe { &mut COUNTERS_MAPS } { + for x in map.as_mut_slice() { + *x = initial; + } + } + Ok(()) + } + + fn usable_count(&self) -> usize { + self.len() + } + + fn to_vec(&self) -> Vec { + let cnt = self.usable_count(); + let mut res = Vec::with_capacity(cnt); + for i in 0..cnt { + res.push(*self.get(i)); + } + res + } + + /// Get the number of set entries with the specified indexes + fn how_many_set(&self, indexes: &[usize]) -> usize { + let initial = self.initial(); + let cnt = self.usable_count(); + let mut res = 0; + for i in indexes { + if *i < cnt && *self.get(*i) != initial { + res += 1; + } + } + res + } + } + + impl CountersMultiMapObserver { + /// Creates a new [`CountersMultiMapObserver`], maybe in differential mode + #[must_use] + fn maybe_differential(name: &'static str) -> Self { + let mut idx = 0; + let mut builder = vec![]; + for (v, x) in unsafe { &COUNTERS_MAPS }.iter().enumerate() { + let l = x.as_slice().len(); + let r = (idx..(idx + l), v); + idx += l; + builder.push(r); + } + Self { + intervals: builder.into_iter().collect::>(), + len: idx, + name: name.to_string(), + initial: u8::default(), + iter_idx: 0, + } + } + } + + impl CountersMultiMapObserver { + /// Creates a new [`CountersMultiMapObserver`] in differential mode + #[must_use] + pub fn differential(name: &'static str) -> Self { + Self::maybe_differential(name) + } + } + + impl CountersMultiMapObserver { + /// Creates a new [`CountersMultiMapObserver`] + #[must_use] + pub fn new(name: &'static str) -> Self { + Self::maybe_differential(name) + } + + /// Creates a new [`CountersMultiMapObserver`] with an owned map + #[must_use] + pub fn owned(name: &'static str) -> Self { + let mut idx = 0; + let mut v = 0; + let mut builder = vec![]; + unsafe { &mut COUNTERS_MAPS }.iter_mut().for_each(|m| { + let l = m.as_mut_slice().len(); + let r = (idx..(idx + l), v); + idx += l; + builder.push(r); + v += 1; + }); + Self { + intervals: builder.into_iter().collect::>(), + len: idx, + name: name.to_string(), + initial: u8::default(), + iter_idx: 0, + } + } + } + + impl<'it, const DIFFERENTIAL: bool> AsIter<'it> for CountersMultiMapObserver { + type Item = u8; + type IntoIter = Flatten>>; + + fn as_iter(&'it self) -> Self::IntoIter { + unsafe { COUNTERS_MAPS.iter().flatten() } + } + } + + impl<'it, const DIFFERENTIAL: bool> AsIterMut<'it> for CountersMultiMapObserver { + type Item = u8; + type IntoIter = Flatten>>; + + fn as_iter_mut(&'it mut self) -> Self::IntoIter { + unsafe { COUNTERS_MAPS.iter_mut().flatten() } + } + } + + impl<'it, const DIFFERENTIAL: bool> IntoIterator for &'it CountersMultiMapObserver { + type Item = as Iterator>::Item; + type IntoIter = Flatten>>; + + fn into_iter(self) -> Self::IntoIter { + unsafe { &COUNTERS_MAPS }.iter().flatten() + } + } + + impl<'it, const DIFFERENTIAL: bool> IntoIterator + for &'it mut CountersMultiMapObserver + { + type Item = as Iterator>::Item; + type IntoIter = Flatten>>; + + fn into_iter(self) -> Self::IntoIter { + unsafe { &mut COUNTERS_MAPS }.iter_mut().flatten() + } + } + + impl DifferentialObserver for CountersMultiMapObserver + where + Self: MapObserver, + OTA: ObserversTuple, + OTB: ObserversTuple, + S: UsesInput, + { + } +}