MinimizerCorpusScheduler skeleton

libafl/src/corpus/minset.rs

@@ -1,225 +1,127 @@
-use alloc::{borrow::ToOwned, vec::Vec};
+use core::marker::PhantomData;
-use core::{cell::RefCell, iter::Iterator, marker::PhantomData};
-use serde::{Deserialize, Serialize};
 
 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
+pub trait FavFactor<I>
-{
-    fn compute<I>(testcase: &Testcase<I>) -> Result<u64, Error>
-    where
-        I: Input;
-}
-
-pub struct LenTimeMulFavFactor {}
-
-// TODO time as Duration and put len into Testcase
-impl FavFactor for LenTimeMulFavFactor {
-    fn compute<I>(entry: &Testcase<I>) -> Result<u64, Error>
-    where
-        I: Input + HasLen
-    {
-        entry.exec_time().as_usec() * (entry.cached_len()? as u64)
-    }
-}
-
-pub trait CorpusMinimizer {
-    fn update_score<C, I, R>(corpus: &mut C, idx: usize) -> Result<(), Error>
-    where
-        C: Corpus<I, R>,
-        I: Input,
-        R: Rand;
-
-    fn cull<C, I, R>(corpus: &mut C) -> Result<(), Error>
-    where
-        C: Corpus<I, R>,
-        I: Input,
-        R: Rand;
-}
-
-pub struct FavCorpusMinimizer<F>
 where
-    F: FavFactor
-{
-    phantom: PhantomData<F>
-}
-
-impl<F> CorpusMinimizer for FavCorpusMinimizer<F>
-where
-    F: FavFactor
-{
-    fn update_score<C, I, R>(corpus: &mut C, idx: usize) -> Result<(), Error>
-    where
-        C: Corpus<I, R>,
-        I: Input,
-        R: Rand
-    {
-    
-    }
-
-    fn cull<C, I, R>(corpus: &mut C) -> Result<(), Error>
-    where
-        C: Corpus<I, R>,
-        I: Input,
-        R: Rand
-    {
-    
-    }
-}
-
-
-#[derive(Serialize, Deserialize)]
-pub struct NoveltiesMeta {
-    novelties: Vec<usize>,
-}
-// impl Iterator<Item = usize>
-
-/// A Queue-like corpus, wrapping an existing Corpus instance
-#[derive(Serialize, Deserialize, Clone, Debug)]
-#[serde(bound = "I: serde::de::DeserializeOwned")]
-pub struct MinSetCorpus<C, F, I, IT, R, T>
-where
-    C: Corpus<I, R>,
-    F: FavFactor,
     I: Input,
-    IT: Iterator<Item = T>,
-    R: Rand,
 {
-    corpus: C,
+    fn compute(testcase: &mut Testcase<I>) -> Result<u64, Error>;
-    pos: usize,
-    // TODO rebase minset on remove()
-    minset: HashSet<usize>,
-    top_rated: HashMap<T, usize>,
-    phantom: PhantomData<(F, I, IT, R)>,
 }
 
-impl<C, I, R> HasTestcaseVec<I> for MinSetCorpus<C, I, R>
+pub struct LenTimeMulFavFactor<I>
 where
-    C: Corpus<I, R>,
+    I: Input + HasLen,
-    I: Input,
-    R: Rand,
 {
-    #[inline]
+    phantom: PhantomData<I>,
-    fn entries(&self) -> &[RefCell<Testcase<I>>] {
+}
-        self.corpus.entries()
+
-    }
+impl<I> FavFactor<I> for LenTimeMulFavFactor<I>
-    #[inline]
+where
-    fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
+    I: Input + HasLen,
-        self.corpus.entries_mut()
+{
+    fn compute(entry: &mut Testcase<I>) -> Result<u64, Error> {
+        Ok(entry.exec_time().map_or(1, |d| d.as_millis()) as u64 * entry.cached_len()? as u64)
     }
 }
 
-impl<C, I, R> Corpus<I, R> for MinSetCorpus<C, I, R>
+pub struct MinimizerCorpusScheduler<C, CS, F, I, S>
 where
-    C: Corpus<I, R>,
+CS: CorpusScheduler<I, S>,
-    I: Input,
+F: FavFactor<I>,
-    R: Rand,
+I: Input,
+S: HasCorpus<C, I>,
+C: Corpus<I>
 {
-    /// Returns the number of elements
+    base: CS,
-    #[inline]
+    phantom: PhantomData<(C, F, I, S)>,
-    fn count(&self) -> usize {
+}
-        self.corpus.count()
+
+impl<C, CS, F, I, S> CorpusScheduler<I, S> for MinimizerCorpusScheduler<C, CS, F, I, S>
+where
+CS: CorpusScheduler<I, S>,
+F: FavFactor<I>,
+I: Input,
+S: HasCorpus<C, I>,
+C: Corpus<I>
+{
+    /// 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
+    /// Replaces the testcase at the given idx
-    #[inline]
+    fn on_replace(&self, state: &mut S, idx: usize, testcase: &Testcase<I>) -> Result<(), Error> {
-    fn add(&mut self, entry: Testcase<I>) -> usize {
+        self.base.on_replace(state, idx, testcase)
-        let idx = self.corpus.add(entry);
-        self.update_score(idx).unwrap();
-        idx
     }
 
     /// Removes an entry from the corpus, returning it if it was present.
-    #[inline]
+    fn on_remove(
-    fn remove(&mut self, entry: &Testcase<I>) -> Option<Testcase<I>> {
+        &self,
-        self.corpus.remove(entry)
+        state: &mut S,
-    }
+        idx: usize,
-
+        testcase: &Option<Testcase<I>>,
-    /// Gets a random entry
+    ) -> Result<(), Error> {
-    #[inline]
+        self.base.on_remove(state, idx, testcase)
-    fn random_entry(&self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, usize), Error> {
-        self.corpus.random_entry(rand)
-    }
-
-    /// Returns the testacase we currently use
-    #[inline]
-    fn current_testcase(&self) -> (&RefCell<Testcase<I>>, usize) {
-        (self.get(self.pos - 1), self.pos - 1)
     }
 
+    // TODO: IntoIter
     /// Gets the next entry
-    #[inline]
+    fn next(&self, state: &mut S) -> Result<usize, Error> {
-    fn next(&mut self, _rand: &mut R) -> Result<(&RefCell<Testcase<I>>, usize), Error> {
+        self.base.next(state)
-        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))
     }
 }
 
-impl<C, I, R> MinSetCorpus<C, I, R>
+impl<C, CS, F, I, S> MinimizerCorpusScheduler<C, CS, F, I, S>
 where
-    C: Corpus<I, R>,
+    CS: CorpusScheduler<I, S>,
+    F: FavFactor<I>,
     I: Input,
-    R: Rand,
+    S: HasCorpus<C, I>,
+    C: Corpus<I>
 {
-    pub fn new(corpus: C) -> Self {
+    /*pub fn update_score(&self, state: &mut S, idx: usize) -> Result<(), Error> {
-        Self {
+        let entry = state.corpus().get(idx)?.borrow_mut();
-            corpus: corpus,
+        let factor = F::compute(&mut *entry)?;
-            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<LenTimeMulFavFactor>>
-    
-    pub fn update_score(&mut self, idx: usize) -> Result<(), Error> {
-        let factor = F::compute(self.entries()[idx].borrow())?;
         for elem in entry.get::<IT>() {
             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() {
+    pub fn new(base: CS) -> Self {
-            if !acc.contains(key) {
+        Self {
-                let idx = self.top_rated.get(key).unwrap();
+            base: base,
-                let entry = self.entries()[idx].borrow();
+            phantom: PhantomData,
-                for elem in entry.get::<IT>() {
-                    acc.insert(elem);
-                }
-                
-                self.minset.insert(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::<IT>() {
+                acc.insert(elem);
+            }
+
+            self.minset.insert(idx);
+        }
+    }
+}
+*/

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<I>) -> Result<(), Error> {
+    fn on_add(&self, _state: &mut S, _idx: usize) -> Result<(), Error> {
         Ok(())
     }
 

libafl/src/corpus/testcase.rs

@@ -29,7 +29,7 @@ where
     /// Time needed to execute the input
     exec_time: Option<Duration>,
     /// Cached len of the input, if any
-    cached_len: Option<usize>
+    cached_len: Option<usize>,
 }
 
 /// 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<Duration> {
+        &self.exec_time
+    }
+
+    /// Get the execution time of the testcase (mut)
+    pub fn exec_time_mut(&mut self) -> &mut Option<Duration> {
+        &mut self.exec_time
+    }
+
     /// Create a new Testcase instace given an input
     #[inline]
     pub fn new<T>(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
+                }
+            },
         })
     }
 }