From 705196fe323c343c3096f9b01fffc6fe6658c0f7 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Thu, 25 Feb 2021 15:34:35 +0100 Subject: [PATCH] MinimizerCorpusScheduler skeleton --- libafl/src/corpus/minset.rs | 266 +++++++++++----------------------- libafl/src/corpus/mod.rs | 3 +- libafl/src/corpus/testcase.rs | 38 +++-- 3 files changed, 109 insertions(+), 198 deletions(-) diff --git a/libafl/src/corpus/minset.rs b/libafl/src/corpus/minset.rs index 3fac83917d..2111d68d9e 100644 --- a/libafl/src/corpus/minset.rs +++ b/libafl/src/corpus/minset.rs @@ -1,225 +1,127 @@ -use alloc::{borrow::ToOwned, vec::Vec}; -use core::{cell::RefCell, iter::Iterator, marker::PhantomData}; -use serde::{Deserialize, Serialize}; +use core::marker::PhantomData; use crate::{ - corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::{HasLen, Input}, utils::Rand, Error, + corpus::{Corpus, CorpusScheduler, Testcase}, + inputs::{HasLen, Input}, + state::HasCorpus, + Error, }; -pub trait FavFactor: Serialize + serde::de::DeserializeOwned + 'static -{ - fn compute(testcase: &Testcase) -> Result - where - I: Input; -} - -pub struct LenTimeMulFavFactor {} - -// TODO time as Duration and put len into Testcase -impl FavFactor for LenTimeMulFavFactor { - fn compute(entry: &Testcase) -> Result - where - I: Input + HasLen - { - entry.exec_time().as_usec() * (entry.cached_len()? as u64) - } -} - -pub trait CorpusMinimizer { - fn update_score(corpus: &mut C, idx: usize) -> Result<(), Error> - where - C: Corpus, - I: Input, - R: Rand; - - fn cull(corpus: &mut C) -> Result<(), Error> - where - C: Corpus, - I: Input, - R: Rand; -} - -pub struct FavCorpusMinimizer +pub trait FavFactor where - F: FavFactor -{ - phantom: PhantomData -} - -impl CorpusMinimizer for FavCorpusMinimizer -where - F: FavFactor -{ - fn update_score(corpus: &mut C, idx: usize) -> Result<(), Error> - where - C: Corpus, - I: Input, - R: Rand - { - - } - - fn cull(corpus: &mut C) -> Result<(), Error> - where - C: Corpus, - I: Input, - R: Rand - { - - } -} - - -#[derive(Serialize, Deserialize)] -pub struct NoveltiesMeta { - novelties: Vec, -} -// impl Iterator - -/// A Queue-like corpus, wrapping an existing Corpus instance -#[derive(Serialize, Deserialize, Clone, Debug)] -#[serde(bound = "I: serde::de::DeserializeOwned")] -pub struct MinSetCorpus -where - C: Corpus, - F: FavFactor, I: Input, - IT: Iterator, - R: Rand, { - corpus: C, - pos: usize, - // TODO rebase minset on remove() - minset: HashSet, - top_rated: HashMap, - phantom: PhantomData<(F, I, IT, R)>, + fn compute(testcase: &mut Testcase) -> Result; } -impl HasTestcaseVec for MinSetCorpus +pub struct LenTimeMulFavFactor where - C: Corpus, - I: Input, - R: Rand, + I: Input + HasLen, { - #[inline] - fn entries(&self) -> &[RefCell>] { - self.corpus.entries() - } - #[inline] - fn entries_mut(&mut self) -> &mut Vec>> { - self.corpus.entries_mut() + phantom: PhantomData, +} + +impl FavFactor for LenTimeMulFavFactor +where + I: Input + HasLen, +{ + fn compute(entry: &mut Testcase) -> Result { + Ok(entry.exec_time().map_or(1, |d| d.as_millis()) as u64 * entry.cached_len()? as u64) } } -impl Corpus for MinSetCorpus +pub struct MinimizerCorpusScheduler where - C: Corpus, - I: Input, - R: Rand, +CS: CorpusScheduler, +F: FavFactor, +I: Input, +S: HasCorpus, +C: Corpus { - /// Returns the number of elements - #[inline] - fn count(&self) -> usize { - self.corpus.count() + base: CS, + phantom: PhantomData<(C, F, I, S)>, +} + +impl CorpusScheduler for MinimizerCorpusScheduler +where +CS: CorpusScheduler, +F: FavFactor, +I: Input, +S: HasCorpus, +C: Corpus +{ + /// Add an entry to the corpus and return its index + fn on_add(&self, state: &mut S, idx: usize) -> Result<(), Error> { + self.base.on_add(state, idx) } - // TODO change add to return Result - #[inline] - fn add(&mut self, entry: Testcase) -> usize { - let idx = self.corpus.add(entry); - self.update_score(idx).unwrap(); - idx + /// Replaces the testcase at the given idx + fn on_replace(&self, state: &mut S, idx: usize, testcase: &Testcase) -> Result<(), Error> { + self.base.on_replace(state, idx, testcase) } /// Removes an entry from the corpus, returning it if it was present. - #[inline] - fn remove(&mut self, entry: &Testcase) -> Option> { - self.corpus.remove(entry) - } - - /// Gets a random entry - #[inline] - fn random_entry(&self, rand: &mut R) -> Result<(&RefCell>, usize), Error> { - self.corpus.random_entry(rand) - } - - /// Returns the testacase we currently use - #[inline] - fn current_testcase(&self) -> (&RefCell>, usize) { - (self.get(self.pos - 1), self.pos - 1) + fn on_remove( + &self, + state: &mut S, + idx: usize, + testcase: &Option>, + ) -> Result<(), Error> { + self.base.on_remove(state, idx, testcase) } + // TODO: IntoIter /// Gets the next entry - #[inline] - fn next(&mut self, _rand: &mut R) -> Result<(&RefCell>, usize), Error> { - self.cull(); - self.pos += 1; - if self.corpus.count() == 0 { - return Err(Error::Empty("Corpus".to_owned())); - } - if self.pos > self.corpus.count() { - // TODO: Always loop or return informational error? - self.pos = 1; - self.cycles += 1; - } - Ok((&self.corpus.entries()[self.pos - 1], self.pos - 1)) + fn next(&self, state: &mut S) -> Result { + self.base.next(state) } } -impl MinSetCorpus +impl MinimizerCorpusScheduler where - C: Corpus, + CS: CorpusScheduler, + F: FavFactor, I: Input, - R: Rand, + S: HasCorpus, + C: Corpus { - pub fn new(corpus: C) -> Self { - Self { - corpus: corpus, - phantom: PhantomData, - } - } - - #[inline] - pub fn corpus(&self) -> &C { - &self.corpus - } - - #[inline] - pub fn corpus_mut(&mut self) -> &mut C { - &mut self.corpus - } - - // TODO move this functions and top rated to another struct - // create something like MinSetCorpus<.., FavMinimizer> - - pub fn update_score(&mut self, idx: usize) -> Result<(), Error> { - let factor = F::compute(self.entries()[idx].borrow())?; + /*pub fn update_score(&self, state: &mut S, idx: usize) -> Result<(), Error> { + let entry = state.corpus().get(idx)?.borrow_mut(); + let factor = F::compute(&mut *entry)?; for elem in entry.get::() { if let val = self.top_rated.get_mut(elem) { if factor > F::compute(self.entries()[val].borrow())? { continue } } - + let _ = self.top_rated.insert(elem, idx); } - } - - pub fn cull(&mut self) { - let mut acc = HashSet::new(); - self.minset.clear(); + }*/ - for key in self.top_rated.keys() { - if !acc.contains(key) { - let idx = self.top_rated.get(key).unwrap(); - let entry = self.entries()[idx].borrow(); - for elem in entry.get::() { - acc.insert(elem); - } - - self.minset.insert(idx); - } + pub fn new(base: CS) -> Self { + Self { + base: base, + phantom: PhantomData, } } } +/* +pub fn cull(&mut self) { + let mut acc = HashSet::new(); + self.minset.clear(); + + for key in self.top_rated.keys() { + if !acc.contains(key) { + let idx = self.top_rated.get(key).unwrap(); + let entry = self.entries()[idx].borrow(); + for elem in entry.get::() { + acc.insert(elem); + } + + self.minset.insert(idx); + } + } +} +*/ diff --git a/libafl/src/corpus/mod.rs b/libafl/src/corpus/mod.rs index cd194bcfdb..366bce49f3 100644 --- a/libafl/src/corpus/mod.rs +++ b/libafl/src/corpus/mod.rs @@ -14,6 +14,7 @@ pub use ondisk::OnDiskCorpus; pub mod queue; pub use queue::QueueCorpusScheduler; +pub mod minset; use core::cell::RefCell; use core::marker::PhantomData; @@ -56,7 +57,7 @@ where I: Input, { /// Add an entry to the corpus and return its index - fn on_add(&self, _state: &mut S, _idx: usize, _testcase: &Testcase) -> Result<(), Error> { + fn on_add(&self, _state: &mut S, _idx: usize) -> Result<(), Error> { Ok(()) } diff --git a/libafl/src/corpus/testcase.rs b/libafl/src/corpus/testcase.rs index d4150d8a0e..0fd7793eae 100644 --- a/libafl/src/corpus/testcase.rs +++ b/libafl/src/corpus/testcase.rs @@ -29,7 +29,7 @@ where /// Time needed to execute the input exec_time: Option, /// Cached len of the input, if any - cached_len: Option + cached_len: Option, } /// Impl of a testcase @@ -135,6 +135,16 @@ where self.metadatas.insert(meta); } + /// Get the execution time of the testcase + pub fn exec_time(&self) -> &Option { + &self.exec_time + } + + /// Get the execution time of the testcase (mut) + pub fn exec_time_mut(&mut self) -> &mut Option { + &mut self.exec_time + } + /// Create a new Testcase instace given an input #[inline] pub fn new(input: T) -> Self @@ -147,7 +157,7 @@ where fitness: 0, metadatas: SerdeAnyMap::new(), exec_time: None, - cached_len: None + cached_len: None, } } @@ -160,7 +170,7 @@ where fitness: 0, metadatas: SerdeAnyMap::new(), exec_time: None, - cached_len: None + cached_len: None, } } @@ -173,7 +183,7 @@ where fitness: fitness, metadatas: SerdeAnyMap::new(), exec_time: None, - cached_len: None + cached_len: None, } } @@ -185,7 +195,7 @@ where fitness: 0, metadatas: SerdeAnyMap::new(), exec_time: None, - cached_len: None + cached_len: None, } } } @@ -203,17 +213,15 @@ where let l = i.len(); self.cached_len = Some(l); l - }, - None => { - match self.cached_len { - Some(l) => l, - None => { - let l = self.load_input()?.len(); - self.cached_len = Some(l); - l - } - } } + None => match self.cached_len { + Some(l) => l, + None => { + let l = self.load_input()?.len(); + self.cached_len = Some(l); + l + } + }, }) } }