diff --git a/afl/src/executors/inmemory.rs b/afl/src/executors/inmemory.rs index 6acca6b8b3..5963a52d14 100644 --- a/afl/src/executors/inmemory.rs +++ b/afl/src/executors/inmemory.rs @@ -1,11 +1,10 @@ -use alloc::rc::Rc; -use core::cell::RefCell; use core::ffi::c_void; use core::ptr; use crate::executors::{Executor, ExitKind}; use crate::inputs::Input; use crate::metamap::NamedAnyMap; +use crate::observers::Observer; use crate::AflError; type HarnessFunction = fn(&dyn Executor, &[u8]) -> ExitKind; @@ -15,7 +14,7 @@ where I: Input, { harness: HarnessFunction, - observers: NamedAnyMap, + observers: NamedAnyMap, } static mut CURRENT_INMEMORY_EXECUTOR_PTR: *const c_void = ptr::null(); @@ -35,6 +34,14 @@ where } Ok(ret) } + + fn observers(&self) -> &NamedAnyMap { + &self.observers + } + + fn observers_mut(&mut self) -> &mut NamedAnyMap { + &mut self.observers + } } impl InMemoryExecutor @@ -48,6 +55,7 @@ where } Self { harness: harness_fn, + observers: NamedAnyMap::new(), } } } diff --git a/afl/src/executors/mod.rs b/afl/src/executors/mod.rs index aaf98385b5..e17d13b146 100644 --- a/afl/src/executors/mod.rs +++ b/afl/src/executors/mod.rs @@ -1,8 +1,8 @@ pub mod inmemory; use crate::inputs::Input; -use crate::observers::Observer; use crate::metamap::NamedAnyMap; +use crate::observers::Observer; use crate::AflError; pub enum ExitKind { @@ -22,29 +22,33 @@ where fn run_target(&mut self, input: &I) -> Result; /// Get the linked observers - fn observers(&self) -> &NamedAnyMap; + fn observers(&self) -> &NamedAnyMap; /// Get the linked observers - fn observers_mut(&mut self) -> &mut NamedAnyMap; + fn observers_mut(&mut self) -> &mut NamedAnyMap; /// Add a linked observer - fn add_observer(&mut self, observer: Box, name: &'static str) { - self.observers_mut().push(observer); + fn add_observer(&mut self, observer: Box) { + self.observers_mut().insert(observer, observer.name()); } /// Reset the state of all the observes linked to this executor fn reset_observers(&mut self) -> Result<(), AflError> { - for observer in self.observers_mut() { - observer.reset()?; + for typeid in self.observers().all_typeids() { + for observer in self.observers_mut().all_by_typeid_mut(typeid).unwrap() { + observer.reset()?; + } } Ok(()) } /// Run the post exec hook for all the observes linked to this executor fn post_exec_observers(&mut self) -> Result<(), AflError> { - self.observers_mut() - .iter() - .map(|x| x.post_exec()) - .fold(Ok(()), |acc, x| if x.is_err() { x } else { acc }) + for typeid in self.observers().all_typeids() { + for observer in self.observers_mut().all_by_typeid_mut(typeid).unwrap() { + observer.post_exec()?; + } + } + Ok(()) } } diff --git a/afl/src/feedbacks/mod.rs b/afl/src/feedbacks/mod.rs index c3cb878b80..7fe07f7efd 100644 --- a/afl/src/feedbacks/mod.rs +++ b/afl/src/feedbacks/mod.rs @@ -7,6 +7,7 @@ use num::Integer; use serde::{Deserialize, Serialize}; use crate::corpus::{Testcase, TestcaseMetadata}; +use crate::executors::Executor; use crate::inputs::Input; use crate::observers::MapObserver; use crate::serde_anymap::SerdeAny; @@ -17,7 +18,7 @@ where I: Input, { /// is_interesting should return the "Interestingness" from 0 to 255 (percent times 2.55) - fn is_interesting(&mut self, input: &I) -> Result; + fn is_interesting(&mut self, input: &I, executor: &dyn Executor) -> Result; /// Append to the testcase the generated metadata in case of a new corpus item fn append_metadata(&mut self, _testcase: &mut Testcase) -> Result<(), AflError> { @@ -28,6 +29,8 @@ where fn discard_metadata(&mut self, _input: &I) -> Result<(), AflError> { Ok(()) } + + fn name(&self) -> &'static str; } /// A Reducer function is used to aggregate values for the novelty search @@ -78,16 +81,6 @@ where } } -/// Returns a usable history map of the given size -pub fn create_history_map(map_size: usize) -> Rc>> -where - T: Default + Clone, -{ - { - Rc::new(RefCell::new(vec![T::default(); map_size])) - } -} - /// The most common AFL-like feedback type pub struct MapFeedback where @@ -96,11 +89,11 @@ where O: MapObserver, { /// Contains information about untouched entries - history_map: Rc>>, - /// The observer this feedback struct observes - map_observer: Rc>, + history_map: Vec, + /// Name identifier of this instance + name: &'static str, /// Phantom Data of Reducer - phantom: PhantomData, + phantom: PhantomData<(R, O)>, } impl Feedback for MapFeedback @@ -110,9 +103,8 @@ where O: MapObserver, I: Input, { - fn is_interesting(&mut self, _input: &I) -> Result { + fn is_interesting(&mut self, _input: &I, executor: &dyn Executor) -> Result { let mut interesting = 0; - // TODO optimize let size = self.map_observer.borrow().map().len(); let mut history_map = self.history_map.borrow_mut(); @@ -129,6 +121,10 @@ where Ok(interesting) } + + fn name(&self) -> &'static str { + self.name + } } impl MapFeedback @@ -137,11 +133,11 @@ where R: Reducer, O: MapObserver, { - /// Create new MapFeedback using a map observer - pub fn new(map_observer: Rc>, map_size: usize) -> Self { + /// Create new MapFeedback + pub fn new(name: &'static str, map_size: usize) -> Self { Self { - map_observer: map_observer, - history_map: create_history_map::(map_size), + history_map: vec![T::default(); map_size], + name: name, phantom: PhantomData, } } @@ -155,13 +151,10 @@ where { /// Create new MapFeedback using a map observer, and a map. /// The map can be shared. - pub fn with_history_map( - map_observer: Rc>, - history_map: Rc>>, - ) -> Self { + pub fn with_history_map(name: &'static str, history_map: Vec) -> Self { Self { - map_observer: map_observer, history_map: history_map, + name: name, phantom: PhantomData, } } @@ -206,11 +199,11 @@ where O: MapObserver, { /// Contains information about untouched entries - history_map: Rc>>, - /// The observer this feedback struct observes - map_observer: Rc>, + history_map: Vec, + /// Name identifier of this instance + name: &'static str, /// Phantom Data of Reducer - phantom: PhantomData, + phantom: PhantomData<(R, O)>, /// Track novel entries indexes novelties: Vec, } @@ -254,6 +247,10 @@ where self.novelties.clear(); Ok(()) } + + fn name(&self) -> &'static str { + self.name + } } impl MapTrackerFeedback diff --git a/afl/src/metamap.rs b/afl/src/metamap.rs index 979e641fcf..42d5dc23c1 100644 --- a/afl/src/metamap.rs +++ b/afl/src/metamap.rs @@ -2,7 +2,7 @@ use alloc::boxed::Box; use alloc::vec::Vec; use core::any::{Any, TypeId}; use core::slice::{Iter, IterMut}; -use hashbrown::hash_map::{Values, ValuesMut}; +use hashbrown::hash_map::{Keys, Values, ValuesMut}; use hashbrown::HashMap; pub struct MetaMap { @@ -241,11 +241,22 @@ impl MetaInstanceMap { } } -pub struct NamedAnyMap { - map: HashMap>>, +pub trait AsAny { + fn as_any(&self) -> &dyn Any; + fn as_any_mut(&mut self) -> &mut dyn Any; } -impl NamedAnyMap { +pub struct NamedAnyMap +where + B: ?Sized + Any + AsAny, +{ + map: HashMap>>, +} + +impl NamedAnyMap +where + B: ?Sized + Any + AsAny, +{ pub fn get(&self, name: &'static str) -> Option<&T> where T: Any, @@ -254,7 +265,14 @@ impl NamedAnyMap { None => None, Some(h) => h .get(&name) - .map(|x| x.as_ref().downcast_ref::().unwrap()), + .map(|x| x.as_any().downcast_ref::().unwrap()), + } + } + + pub fn by_typeid(&self, name: &'static str, typeid: &TypeId) -> Option<&B> { + match self.map.get(typeid) { + None => None, + Some(h) => h.get(&name).map(|x| x.as_ref()), } } @@ -266,27 +284,42 @@ impl NamedAnyMap { None => None, Some(h) => h .get_mut(&name) - .map(|x| x.as_mut().downcast_mut::().unwrap()), + .map(|x| x.as_any_mut().downcast_mut::().unwrap()), + } + } + + pub fn by_typeid_mut(&mut self, name: &'static str, typeid: &TypeId) -> Option<&mut B> { + match self.map.get_mut(typeid) { + None => None, + Some(h) => h.get(&name).map(|x| x.as_mut()), } } pub fn get_all( &self, - ) -> Option>, fn(&Box) -> &T>> + ) -> Option>, fn(&Box) -> &T>> where T: Any, { match self.map.get(&TypeId::of::()) { None => None, - Some(h) => Some(h.values().map(|x| x.as_ref().downcast_ref::().unwrap())), + Some(h) => Some(h.values().map(|x| x.as_any().downcast_ref::().unwrap())), + } + } + + pub fn all_by_typeid( + &self, + typeid: &TypeId, + ) -> Option>, fn(&Box) -> &B>> { + match self.map.get(typeid) { + None => None, + Some(h) => Some(h.values().map(|x| x.as_ref())), } } pub fn get_all_mut( &mut self, - ) -> Option< - core::iter::Map>, fn(&mut Box) -> &mut T>, - > + ) -> Option>, fn(&mut Box) -> &mut T>> where T: Any, { @@ -294,23 +327,32 @@ impl NamedAnyMap { None => None, Some(h) => Some( h.values_mut() - .map(|x| x.as_mut().downcast_mut::().unwrap()), + .map(|x| x.as_any_mut().downcast_mut::().unwrap()), ), } } - pub fn insert(&mut self, t: T, name: &'static str) - where - T: Any, + pub fn all_by_typeid_mut( + &mut self, + typeid: &TypeId, + ) -> Option>, fn(&mut Box) -> &mut B>> { - let typeid = TypeId::of::(); + match self.map.get_mut(typeid) { + None => None, + Some(h) => Some(h.values_mut().map(|x| x.as_mut())), + } + } + + pub fn all_typeids(&self) -> Keys<'_, TypeId, HashMap<&'static str, Box>> { + self.map.keys() + } + + pub fn insert(&mut self, val: Box, name: &'static str) { + let typeid = val.type_id(); if !self.map.contains_key(&typeid) { self.map.insert(typeid, HashMap::default()); } - self.map - .get_mut(&typeid) - .unwrap() - .insert(name, Box::new(t)); + self.map.get_mut(&typeid).unwrap().insert(name, val); } pub fn len(&self) -> usize { diff --git a/afl/src/observers/mod.rs b/afl/src/observers/mod.rs index f790e6a15e..54ae4097e4 100644 --- a/afl/src/observers/mod.rs +++ b/afl/src/observers/mod.rs @@ -1,16 +1,15 @@ extern crate num; -use alloc::rc::Rc; -use core::cell::RefCell; -use core::slice::from_raw_parts_mut; use core::any::Any; +use core::slice::from_raw_parts_mut; use num::Integer; +use crate::metamap::AsAny; use crate::AflError; /// Observers observe different information about the target. /// They can then be used by various sorts of feedback. -pub trait Observer: Any { +pub trait Observer: Any + AsAny { fn flush(&mut self) -> Result<(), AflError> { Ok(()) } @@ -20,6 +19,8 @@ pub trait Observer: Any { fn post_exec(&mut self) -> Result<(), AflError> { Ok(()) } + + fn name(&self) -> &'static str; } /// A MapObserver observes the static map, as oftentimes used for afl-like coverage information @@ -52,24 +53,41 @@ where } } -pub struct StdMapObserver<'a, T> +pub struct StdMapObserver where - T: Integer + Copy, + T: Integer + Copy + 'static, { - map: &'a mut [T], + map: &'static mut [T], initial: T, + name: &'static str, } -impl<'a, T> Observer for StdMapObserver<'a, T> +impl Observer for StdMapObserver where T: Integer + Copy, { fn reset(&mut self) -> Result<(), AflError> { self.reset_map() } + + fn name(&self) -> &'static str { + self.name + } } -impl<'a, T> MapObserver for StdMapObserver<'a, T> +impl AsAny for StdMapObserver +where + T: Integer + Copy, +{ + fn as_any(&self) -> &dyn Any { + self + } + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +impl MapObserver for StdMapObserver where T: Integer + Copy, { @@ -94,36 +112,29 @@ where } } -impl<'a, T> StdMapObserver<'a, T> +impl StdMapObserver where T: Integer + Copy, { /// Creates a new MapObserver - pub fn new(map: &'a mut [T]) -> Self { + pub fn new(name: &'static str, map: &'static mut [T]) -> Self { let initial = if map.len() > 0 { map[0] } else { T::zero() }; Self { map: map, initial: initial, + name: name, } } /// Creates a new MapObserver from a raw pointer - pub fn new_from_ptr(map_ptr: *mut T, len: usize) -> Self { + pub fn new_from_ptr(name: &'static str, map_ptr: *mut T, len: usize) -> Self { unsafe { let initial = if len > 0 { *map_ptr } else { T::zero() }; StdMapObserver { map: from_raw_parts_mut(map_ptr, len), initial: initial, + name: name, } } } } - -impl<'a, T> Into>> for StdMapObserver<'a, T> -where - T: Integer + Copy, -{ - fn into(self) -> Rc> { - Rc::new(RefCell::new(self)) - } -}