rand and queue corpus schedulers

libafl/src/corpus/inmemory.rs

@@ -1,74 +1,74 @@
 //! In-memory corpus, keeps all test cases in memory at all times
 
-use alloc::{borrow::ToOwned, vec::Vec};
-use core::{cell::RefCell, marker::PhantomData};
+use alloc::vec::Vec;
+use core::cell::RefCell;
 use serde::{Deserialize, Serialize};
 
-use crate::{
-    corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, Error,
-};
+use crate::{corpus::Corpus, corpus::Testcase, inputs::Input, Error};
 
-/// A corpus handling all important fuzzing in memory.
-#[derive(Serialize, Deserialize, Clone, Debug)]
+/// A corpus handling all in memory.
+#[derive(Default, Serialize, Deserialize, Clone, Debug)]
 #[serde(bound = "I: serde::de::DeserializeOwned")]
-pub struct InMemoryCorpus<I, R>
+pub struct InMemoryCorpus<I>
 where
     I: Input,
-    R: Rand,
 {
     entries: Vec<RefCell<Testcase<I>>>,
-    pos: usize,
-    phantom: PhantomData<R>,
+    current: Option<usize>,
 }
 
-impl<I, R> HasTestcaseVec<I> for InMemoryCorpus<I, R>
+impl<I> Corpus<I> for InMemoryCorpus<I>
 where
     I: Input,
-    R: Rand,
 {
-    fn entries(&self) -> &[RefCell<Testcase<I>>] {
-        &self.entries
-    }
-    fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
-        &mut self.entries
-    }
-}
-
-impl<I, R> Corpus<I, R> for InMemoryCorpus<I, R>
-where
-    I: Input,
-    R: Rand,
-{
-    /// Gets the next entry
+    /// Returns the number of elements
     #[inline]
-    fn next(&mut self, rand: &mut R) -> Result<(&RefCell<Testcase<I>>, usize), Error> {
-        if self.count() == 0 {
-            Err(Error::Empty("No entries in corpus".to_owned()))
+    fn count(&self) -> usize {
+        self.entries.len()
+    }
+
+    /// Add an entry to the corpus and return its index
+    #[inline]
+    fn add(&mut self, testcase: Testcase<I>) -> Result<usize, Error> {
+        self.entries.push(RefCell::new(testcase));
+        Ok(self.entries.len() - 1)
+    }
+
+    /// Replaces the testcase at the given idx
+    #[inline]
+    fn replace(&mut self, idx: usize, testcase: Testcase<I>) -> Result<(), Error> {
+        if idx >= self.entries.len() {
+            return Err(Error::KeyNotFound(format!("Index {} out of bounds", idx)));
+        }
+        self.entries[idx] = RefCell::new(testcase);
+        Ok(())
+    }
+
+    /// Removes an entry from the corpus, returning it if it was present.
+    #[inline]
+    fn remove(&mut self, idx: usize) -> Result<Option<Testcase<I>>, Error> {
+        if idx >= self.entries.len() {
+            Ok(None)
         } else {
-            let len = { self.entries().len() };
-            let id = rand.below(len as u64) as usize;
-            self.pos = id;
-            Ok((self.get(id), id))
+            Ok(Some(self.entries.remove(idx).into_inner()))
         }
     }
 
-    /// Returns the testacase we currently use
+    /// Get by id
     #[inline]
-    fn current_testcase(&self) -> (&RefCell<Testcase<I>>, usize) {
-        (self.get(self.pos), self.pos)
-    }
+    fn get(&self, idx: usize) -> Result<&RefCell<Testcase<I>>, Error> {
+        Ok(&self.entries[idx])
     }
 
-impl<I, R> InMemoryCorpus<I, R>
-where
-    I: Input,
-    R: Rand,
-{
-    pub fn new() -> Self {
-        Self {
-            entries: vec![],
-            pos: 0,
-            phantom: PhantomData,
-        }
+    /// Current testcase scheduled
+    #[inline]
+    fn current(&self) -> &Option<usize> {
+        &self.current
+    }
+
+    /// Current testcase scheduled (mut)
+    #[inline]
+    fn current_mut(&mut self) -> &mut Option<usize> {
+        &mut self.current
     }
 }

libafl/src/corpus/mod.rs

@@ -3,11 +3,21 @@
 pub mod testcase;
 pub use testcase::Testcase;
 
-use alloc::vec::Vec;
-use core::cell::RefCell;
-use serde::{Deserialize, Serialize};
+pub mod inmemory;
+pub use inmemory::InMemoryCorpus;
 
-use crate::{inputs::Input, Error};
+pub mod queue;
+pub use queue::QueueCorpusScheduler;
+
+use core::cell::RefCell;
+use core::marker::PhantomData;
+
+use crate::{
+    inputs::Input,
+    state::{HasCorpus, HasRand},
+    utils::Rand,
+    Error,
+};
 
 /// Corpus with all current testcases
 pub trait Corpus<I>: serde::Serialize + serde::de::DeserializeOwned
@@ -46,7 +56,12 @@ where
     }
 
     /// Replaces the testcase at the given idx
-    fn on_replace(&self, _state: &mut S, _idx: usize, _testcase: &Testcase<I>) -> Result<(), Error> {
+    fn on_replace(
+        &self,
+        _state: &mut S,
+        _idx: usize,
+        _testcase: &Testcase<I>,
+    ) -> Result<(), Error> {
         Ok(())
     }
 
@@ -65,89 +80,34 @@ where
     fn next(&self, state: &mut S) -> Result<usize, Error>;
 }
 
-/*
-pub struct RandCorpusScheduler {}
-
-impl CorpusScheduler for RandCorpusScheduler {
-    /// Gets the next entry at random
-    fn next<C, I, R, S>(state: &mut S) -> Result<usize, Error>
+pub struct RandCorpusScheduler<C, I, R, S>
 where
     S: HasCorpus<C, I> + HasRand<R>,
     C: Corpus<I>,
     I: Input,
     R: Rand,
 {
+    phantom: PhantomData<(C, I, R, S)>,
+}
+
+impl<C, I, R, S> CorpusScheduler<I, S> for RandCorpusScheduler<C, I, R, S>
+where
+    S: HasCorpus<C, I> + HasRand<R>,
+    C: Corpus<I>,
+    I: Input,
+    R: Rand,
+{
+    /// Gets the next entry at random
+    fn next(&self, state: &mut S) -> Result<usize, Error> {
         if state.corpus().count() == 0 {
             Err(Error::Empty("No entries in corpus".to_owned()))
         } else {
             let len = state.corpus().count();
             let id = state.rand_mut().below(len as u64) as usize;
+            *state.corpus_mut().current_mut() = Some(id);
             Ok(id)
         }
     }
 }
-*/
 
-#[derive(Default, Serialize, Deserialize, Clone, Debug)]
-#[serde(bound = "I: serde::de::DeserializeOwned")]
-pub struct InMemoryCorpus<I>
-where
-    I: Input,
-{
-    entries: Vec<RefCell<Testcase<I>>>,
-    current: Option<usize>,
-}
-
-impl<I> Corpus<I> for InMemoryCorpus<I>
-where
-    I: Input,
-{
-    /// Returns the number of elements
-    #[inline]
-    fn count(&self) -> usize {
-        self.entries.len()
-    }
-
-    /// Add an entry to the corpus and return its index
-    #[inline]
-    fn add(&mut self, testcase: Testcase<I>) -> Result<usize, Error> {
-        self.entries.push(RefCell::new(testcase));
-        Ok(self.entries.len() - 1)
-    }
-
-    /// Replaces the testcase at the given idx
-    #[inline]
-    fn replace(&mut self, idx: usize, testcase: Testcase<I>) -> Result<(), Error> {
-        if idx >= self.entries.len() {
-            return Err(Error::KeyNotFound(format!("Index {} out of bounds", idx)));
-        }
-        self.entries[idx] = RefCell::new(testcase);
-        Ok(())
-    }
-
-    /// Removes an entry from the corpus, returning it if it was present.
-    #[inline]
-    fn remove(&mut self, idx: usize) -> Result<Option<Testcase<I>>, Error> {
-        if idx >= self.entries.len() {
-            Ok(None)
-        } else {
-            Ok(Some(self.entries.remove(idx).into_inner()))
-        }
-    }
-
-    /// Get by id
-    #[inline]
-    fn get(&self, idx: usize) -> Result<&RefCell<Testcase<I>>, Error> {
-        Ok(&self.entries[idx])
-    }
-
-    /// Current testcase scheduled
-    fn current(&self) -> &Option<usize> {
-        &self.current
-    }
-
-    /// Current testcase scheduled (mut)
-    fn current_mut(&mut self) -> &mut Option<usize> {
-        &mut self.current
-    }
-}
+pub type StdCorpusScheduler<C, I, R, S> = RandCorpusScheduler<C, I, R, S>;

libafl/src/corpus/queue.rs

@@ -1,132 +1,52 @@
 //! The queue corpus implements an afl-like queue mechanism
 
-use alloc::{borrow::ToOwned, vec::Vec};
-use core::{cell::RefCell, marker::PhantomData};
-use serde::{Deserialize, Serialize};
+use alloc::borrow::ToOwned;
+use core::marker::PhantomData;
 
 use crate::{
-    corpus::Corpus, corpus::HasTestcaseVec, corpus::Testcase, inputs::Input, utils::Rand, Error,
+    corpus::{Corpus, CorpusScheduler},
+    inputs::Input,
+    state::HasCorpus,
+    Error,
 };
 
-/// A Queue-like corpus, wrapping an existing Corpus instance
-#[derive(Serialize, Deserialize, Clone, Debug)]
-#[serde(bound = "I: serde::de::DeserializeOwned")]
-pub struct QueueCorpus<C, I, R>
+pub struct QueueCorpusScheduler<C, I, S>
 where
-    C: Corpus<I, R>,
+    S: HasCorpus<C, I>,
+    C: Corpus<I>,
     I: Input,
-    R: Rand,
 {
-    corpus: C,
-    pos: usize,
-    cycles: u64,
-    phantom: PhantomData<(I, R)>,
+    phantom: PhantomData<(C, I, S)>,
 }
 
-impl<C, I, R> HasTestcaseVec<I> for QueueCorpus<C, I, R>
+impl<C, I, S> CorpusScheduler<I, S> for QueueCorpusScheduler<C, I, S>
 where
-    C: Corpus<I, R>,
+    S: HasCorpus<C, I>,
+    C: Corpus<I>,
     I: Input,
-    R: Rand,
 {
-    #[inline]
-    fn entries(&self) -> &[RefCell<Testcase<I>>] {
-        self.corpus.entries()
+    /// Gets the next entry at random
+    fn next(&self, state: &mut S) -> Result<usize, Error> {
+        if state.corpus().count() == 0 {
+            Err(Error::Empty("No entries in corpus".to_owned()))
+        } else {
+            let id = match state.corpus().current() {
+                Some(cur) => {
+                    if *cur + 1 > state.corpus().count() {
+                        0
+                    } else {
+                        *cur + 1
+                    }
+                }
+                None => 0,
+            };
+            *state.corpus_mut().current_mut() = Some(id);
+            Ok(id)
         }
-    #[inline]
-    fn entries_mut(&mut self) -> &mut Vec<RefCell<Testcase<I>>> {
-        self.corpus.entries_mut()
-    }
-}
-
-impl<C, I, R> Corpus<I, R> for QueueCorpus<C, I, R>
-where
-    C: Corpus<I, R>,
-    I: Input,
-    R: Rand,
-{
-    /// Returns the number of elements
-    #[inline]
-    fn count(&self) -> usize {
-        self.corpus.count()
-    }
-
-    #[inline]
-    fn add(&mut self, entry: Testcase<I>) -> usize {
-        self.corpus.add(entry)
-    }
-
-    /// Removes an entry from the corpus, returning it if it was present.
-    #[inline]
-    fn remove(&mut self, entry: &Testcase<I>) -> Option<Testcase<I>> {
-        self.corpus.remove(entry)
-    }
-
-    /// Gets a random entry
-    #[inline]
-    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)
-    }
-
-    /// Gets the next entry
-    #[inline]
-    fn next(&mut self, _rand: &mut R) -> Result<(&RefCell<Testcase<I>>, usize), Error> {
-        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> QueueCorpus<C, I, R>
-where
-    C: Corpus<I, R>,
-    I: Input,
-    R: Rand,
-{
-    pub fn new(corpus: C) -> Self {
-        Self {
-            corpus: corpus,
-            phantom: PhantomData,
-            cycles: 0,
-            pos: 0,
-        }
-    }
-
-    #[inline]
-    pub fn cycles(&self) -> u64 {
-        self.cycles
-    }
-
-    #[inline]
-    pub fn pos(&self) -> usize {
-        self.pos
-    }
-    
-    // TODO maybe impl HasCorpus
-    #[inline]
-    pub fn corpus(&self) -> &C {
-        &self.corpus
-    }
-    
-    #[inline]
-    pub fn corpus_mut(&mut self) -> &mut C {
-        &mut self.corpus
     }
 }
 
+/*
 #[cfg(test)]
 #[cfg(feature = "std")]
 mod tests {
@@ -170,3 +90,4 @@ mod tests {
         assert_eq!(filename, "fancyfile");
     }
 }
+*/

libafl/src/executors/inprocess.rs

@@ -249,7 +249,10 @@ pub mod unix_signals {
             .is_interesting_all(&input, observers, ExitKind::Crash)
             .expect("In crash handler objectives failure.".into());
         if obj_fitness > 0 {
-            state.solutions_mut().add(Testcase::new(input.clone())).expect("In crash handler solutions failure.".into());
+            state
+                .solutions_mut()
+                .add(Testcase::new(input.clone()))
+                .expect("In crash handler solutions failure.".into());
             mgr.fire(
                 state,
                 Event::Objective {
@@ -302,7 +305,10 @@ pub mod unix_signals {
             .is_interesting_all(&input, observers, ExitKind::Crash)
             .expect("In timeout handler objectives failure.".into());
         if obj_fitness > 0 {
-            state.solutions_mut().add(Testcase::new(input.clone())).expect("In timeout handler solutions failure.".into());
+            state
+                .solutions_mut()
+                .add(Testcase::new(input.clone()))
+                .expect("In timeout handler solutions failure.".into());
             mgr.fire(
                 state,
                 Event::Objective {

libafl/src/mutators/mutations.rs

@@ -711,7 +711,13 @@ where
         }
     }
 
-    let other_size = state.corpus().get(idx)?.borrow_mut().load_input()?.bytes().len();
+    let other_size = state
+        .corpus()
+        .get(idx)?
+        .borrow_mut()
+        .load_input()?
+        .bytes()
+        .len();
     if other_size < 2 {
         return Ok(MutationResult::Skipped);
     }
@@ -762,7 +768,13 @@ where
         }
     }
 
-    let other_size = state.corpus().get(idx)?.borrow_mut().load_input()?.bytes().len();
+    let other_size = state
+        .corpus()
+        .get(idx)?
+        .borrow_mut()
+        .load_input()?
+        .bytes()
+        .len();
     if other_size < 2 {
         return Ok(MutationResult::Skipped);
     }

libafl/src/stages/mod.rs

@@ -29,14 +29,7 @@ pub trait StagesTuple<E, EM, F, S> {
 }
 
 impl<E, EM, F, S> StagesTuple<E, EM, F, S> for () {
-    fn perform_all(
-        &self,
-        _: &F,
-        _: &mut S,
-        _: &mut E,
-        _: &mut EM,
-        _: usize,
-    ) -> Result<(), Error> {
+    fn perform_all(&self, _: &F, _: &mut S, _: &mut E, _: &mut EM, _: usize) -> Result<(), Error> {
         Ok(())
     }
 }

libafl/src/stages/mutational.rs

@@ -59,8 +59,7 @@ where
 
             let fitness = state.evaluate_input(input_mut, executor, manager)?;
 
-            self.mutator()
-                .post_exec(fuzzer, state, fitness, i as i32)?;
+            self.mutator().post_exec(fuzzer, state, fitness, i as i32)?;
         }
         Ok(())
     }

libafl/src/state/mod.rs

@@ -583,12 +583,12 @@ where
         executor.post_exec_observers()?;
 
         let observers = executor.observers();
-        let fitness = self
-            .feedbacks_mut()
+        let fitness =
+            self.feedbacks_mut()
                 .is_interesting_all(&input, observers, exit_kind.clone())?;
 
-        let is_solution =
-            self.objectives_mut()
+        let is_solution = self
+            .objectives_mut()
             .is_interesting_all(&input, observers, exit_kind)?
             > 0;
         Ok((fitness, is_solution))