pass Testcase inside Rc+RefCell

src/corpus/mod.rs

@@ -1,5 +1,5 @@
 pub mod testcase;
-pub use testcase::Testcase;
+pub use testcase::{Testcase, TestcaseMetadata};
 
 use crate::utils::{Rand, HasRand};
 use crate::inputs::Input;
@@ -7,13 +7,15 @@ use crate::AflError;
 
 use std::path::PathBuf;
 use std::marker::PhantomData;
+use std::cell::RefCell;
+use std::rc::Rc;
 
 pub trait HasEntriesVec<I> where I: Input {
     /// Get the entries vector field
-    fn entries(&self) -> &Vec<Box<Testcase<I>>>;
+    fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>>;
 
     /// Get the entries vector field (mutable)
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>;
+    fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>;
 }
 
 /// Corpus with all current testcases
@@ -25,17 +27,17 @@ pub trait Corpus<I> : HasEntriesVec<I> + HasRand where I: Input {
 
     /// Add an entry to the corpus
     #[allow(unused_mut)]
-    fn add(&mut self, mut entry: Box<Testcase<I>>) {
+    fn add(&mut self, mut entry: Rc<RefCell<Testcase<I>>>) {
         self.entries_mut().push(entry);
     }
 
     /// Removes an entry from the corpus, returning it if it was present.
-    fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> {
+    fn remove(&mut self, entry: &Testcase<I>) -> Option<Rc<RefCell<Testcase<I>>>> {
         let mut i: usize = 0;
         let mut found = false;
         for x in self.entries() {
             i = i + 1;
-            if x.as_ref() as *const _ == entry as *const _ {
+            if &*x.borrow() as *const _ == entry as *const _ { // TODO check if correct
                 found = true;
                 break;
             }
@@ -47,28 +49,28 @@ pub trait Corpus<I> : HasEntriesVec<I> + HasRand where I: Input {
     }
 
     /// Gets a random entry
-    fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
+    fn random_entry(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
         let len = { self.entries().len() };
         let id = self.rand_mut().below(len as u64) as usize;
         Ok(self.entries_mut().get_mut(id).unwrap())
     }
 
     /// Gets the next entry (random by default)
-    fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
+    fn get(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
         self.random_entry()
     }
 }
 
 pub struct InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
     rand: &'a mut R,
-    entries: Vec<Box<Testcase<I>>>
+    entries: Vec<Rc<RefCell<Testcase<I>>>>
 }
 
 impl<I, R> HasEntriesVec<I> for InMemoryCorpus<'_, I, R> where I: Input, R: Rand {
-    fn entries(&self) -> &Vec<Box<Testcase<I>>> {
+    fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>> {
         &self.entries
     }
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
+    fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>{
         &mut self.entries
     }
 }
@@ -99,15 +101,15 @@ impl<'a, I, R> InMemoryCorpus<'a, I, R> where I: Input, R: Rand {
 
 pub struct OnDiskCorpus<'a, I, R> where I: Input, R: Rand {
     rand: &'a mut R,
-    entries: Vec<Box<Testcase<I>>>,
+    entries: Vec<Rc<RefCell<Testcase<I>>>>,
     dir_path: PathBuf,
 }
 
 impl<I, R> HasEntriesVec<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
-    fn entries(&self) -> &Vec<Box<Testcase<I>>> {
+    fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>> {
         &self.entries
     }
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
+    fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>{
         &mut self.entries
     }
 }
@@ -125,12 +127,12 @@ impl<I, R> HasRand for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
 
 impl<I, R> Corpus<I> for OnDiskCorpus<'_, I, R> where I: Input, R: Rand {
     /// Add an entry and save it to disk
-    fn add(&mut self, mut entry: Box<Testcase<I>>) {
-        if *entry.filename() == None {
+    fn add(&mut self, entry: Rc<RefCell<Testcase<I>>>) {
+        if *entry.borrow().filename() == None {
             // TODO walk entry metadatas to ask for pices of filename (e.g. :havoc in AFL)
             let filename = &(String::from("id:") + &self.entries.len().to_string());
             let filename = self.dir_path.join(filename);
-            *entry.filename_mut() = Some(filename);
+            *entry.borrow_mut().filename_mut() = Some(filename);
         }
         self.entries.push(entry);
     }
@@ -157,10 +159,10 @@ pub struct QueueCorpus<I, C> where I: Input, C: Corpus<I> {
 }
 
 impl<'a, I, C> HasEntriesVec<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
-    fn entries(&self) -> &Vec<Box<Testcase<I>>> {
+    fn entries(&self) -> &Vec<Rc<RefCell<Testcase<I>>>> {
         self.corpus.entries()
     }
-    fn entries_mut(&mut self) -> &mut Vec<Box<Testcase<I>>>{
+    fn entries_mut(&mut self) -> &mut Vec<Rc<RefCell<Testcase<I>>>>{
         self.corpus.entries_mut()
     }
 }
@@ -182,22 +184,22 @@ impl<'a, I, C> Corpus<I> for QueueCorpus<I, C> where I: Input, C: Corpus<I> {
         self.corpus.count()
     }
 
-    fn add(&mut self, entry: Box<Testcase<I>>) {
+    fn add(&mut self, entry: Rc<RefCell<Testcase<I>>>) {
         self.corpus.add(entry);
     }
 
     /// Removes an entry from the corpus, returning it if it was present.
-    fn remove(&mut self, entry: &Testcase<I>) -> Option<Box<Testcase<I>>> {
+    fn remove(&mut self, entry: &Testcase<I>) -> Option<Rc<RefCell<Testcase<I>>>> {
         self.corpus.remove(entry)
     }
 
     /// Gets a random entry
-    fn random_entry(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
+    fn random_entry(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
         self.corpus.random_entry()
     }
 
     /// Gets the next entry
-    fn get(&mut self) -> Result<&Box<Testcase<I>>, AflError> {
+    fn get(&mut self) -> Result<&Rc<RefCell<Testcase<I>>>, AflError> {
         if self.corpus.count() == 0 {
             return Err(AflError::Empty("Testcases".to_string()));
         }
@@ -238,18 +240,19 @@ mod tests {
     use crate::utils::Xoshiro256StarRand;
 
     use std::path::PathBuf;
+    use std::cell::RefCell;
+    use std::rc::Rc;
 
     #[test]
 
     fn test_queuecorpus() {
         let mut rand = Xoshiro256StarRand::new();
         let mut q = QueueCorpus::new(OnDiskCorpus::new(&mut rand, PathBuf::from("fancy/path")));
-        let i = Box::new(BytesInput::new(vec![0; 4]));
-        let mut t = Box::new(Testcase::new(i));
-        *t.filename_mut() = Some(PathBuf::from("fancyfile"));
+        let i = BytesInput::new(vec![0; 4]);
+        let t = Rc::new(RefCell::new(Testcase::new_with_filename(i, PathBuf::from("fancyfile"))));
         q.add(t);
-        let filename = q.get().unwrap().filename().as_ref().unwrap().to_owned();
-        assert_eq!(filename, q.get().unwrap().filename().as_ref().unwrap().to_owned());
+        let filename = q.get().unwrap().borrow().filename().as_ref().unwrap().to_owned();
+        assert_eq!(filename, q.get().unwrap().borrow().filename().as_ref().unwrap().to_owned());
         assert_eq!(filename, PathBuf::from("fancy/path/fancyfile"));
     }
 }

src/corpus/testcase.rs

@@ -9,13 +9,13 @@ pub trait TestcaseMetadata {}
 /*
 pub trait TestcaseTrait<I: Input> {
     /// Make sure to return a valid input instance loading it from disk if not in memory
-    fn load_input(&mut self) -> Result<&Box<I>, AflError>;
+    fn load_input(&mut self) -> Result<&I, AflError>;
 
     /// Get the input, if any
-    fn input(&self) -> &Option<Box<I>>;
+    fn input(&self) -> &Option<I>;
 
     /// Get the input, if any (mutable)
-    fn input_mut(&mut self) -> &mut Option<Box<I>>;
+    fn input_mut(&mut self) -> &mut Option<I>;
 
     /// Get the filename, if any
     fn filename(&self) -> &Option<PathBuf>;
@@ -30,28 +30,31 @@ pub trait TestcaseTrait<I: Input> {
 
 #[derive(Default)]
 pub struct Testcase<I> where I: Input {
-    input: Option<Box<I>>,
+    input: Option<I>, // TODO remove box
     filename: Option<PathBuf>,
     metadatas: HashMap<String, Box<dyn TestcaseMetadata>>,
 }
 
 impl<I> Testcase<I> where I: Input {
     /// Make sure to return a valid input instance loading it from disk if not in memory
-    pub fn load_input(&mut self) -> Result<&Box<I>, AflError> {
+    pub fn load_input(&mut self) -> Result<&I, AflError> {
         // TODO: Implement cache to disk
-        self.input.as_ref().ok_or(AflError::NotImplemented("load_input".to_string()))
+        match self.input.as_ref() {
+            Some(i) => Ok(i),
+            None => Err(AflError::NotImplemented("load_input".to_string()))
+        }
     }
 
     /// Get the input, if any
-    pub fn input(&self) -> &Option<Box<I>> {
+    pub fn input(&self) -> &Option<I> {
         &self.input
     }
     /// Get the input, if any (mutable)
-    pub fn input_mut(&mut self) -> &mut Option<Box<I>> {
+    pub fn input_mut(&mut self) -> &mut Option<I> {
         &mut self.input
     }
     /// Set the input
-    pub fn set_input(&mut self, input: Option<Box<I>>) {
+    pub fn set_input(&mut self, input: Option<I>) {
         self.input = input;
     }
 
@@ -74,7 +77,7 @@ impl<I> Testcase<I> where I: Input {
     }
 
     /// Create a new DefaultTestcase instace given an input
-    pub fn new(input: Box<I>) -> Self {
+    pub fn new(input: I) -> Self {
         Testcase {
             input: Some(input),
             filename: None,
@@ -83,7 +86,7 @@ impl<I> Testcase<I> where I: Input {
     }
 
     /// Create a new DefaultTestcase instace given an input and a filename
-    pub fn new_with_filename(input: Box<I>, filename: PathBuf) -> Self {
+    pub fn new_with_filename(input: I, filename: PathBuf) -> Self {
         Testcase {
             input: Some(input),
             filename: Some(filename),

src/engines/aflengine.rs

@@ -7,7 +7,7 @@ use crate::feedbacks::Feedback;
 use crate::monitors::Monitor;
 use crate::stages::Stage;
 use crate::utils::Rand;
-
+/*
 pub struct AflEngine<'a, I: Input> {
     pub rand: &'a mut dyn Rand,
     pub feedbacks: Vec<Box<dyn Feedback<I>>>,
@@ -27,3 +27,4 @@ pub struct AflEngine<'a, I: Input> {
 }
 
 impl<I: Input> Engine<'_> for AflEngine<'_, I> {}
+*/

src/engines/mod.rs

@@ -1,3 +1,14 @@
 pub mod aflengine;
 
-pub trait Engine<'a> {}
+use crate::AflError;
+use crate::inputs::Input;
+use crate::corpus::testcase::Testcase;
+
+use std::cell::RefCell;
+use std::rc::Rc;
+
+pub trait Engine<'a, I> where I: Input {
+
+    fn execute(&mut self, input: &mut I, entry: Rc<RefCell<Testcase<I>>>) -> Result<bool, AflError>;
+
+}

src/executors/inmemory.rs

@@ -177,6 +177,7 @@ mod tests {
     use crate::observers::Observer;
     use crate::AflError;
 
+    #[derive(Clone)]
     struct NopInput {}
     impl Input for NopInput {
         fn serialize(&self) -> Result<&[u8], AflError> {

src/executors/mod.rs

@@ -1,5 +1,4 @@
 pub mod inmemory;
-pub use inmemory::{InMemoryExecutor};
 
 use crate::inputs::Input;
 use crate::observers::Observer;

src/inputs/mod.rs

@@ -5,10 +5,11 @@ use std::fs::File;
 use std::io::Read;
 use std::io::Write;
 use std::path::PathBuf;
+use std::clone::Clone;
 
 use crate::AflError;
 
-pub trait Input {
+pub trait Input : Clone {
     fn to_file(&self, path: &PathBuf) -> Result<(), AflError> {
         let mut file = File::create(path)?;
         file.write_all(self.serialize()?)?;

src/mutators/mod.rs

@@ -4,7 +4,6 @@ use crate::corpus::Corpus;
 use crate::AflError;
 
 pub mod scheduled;
-pub use scheduled::{ComposedByMutations, ScheduledMutator, HavocBytesMutator};
 
 pub trait HasOptionCorpus<I> where I: Input {
     type C : Corpus<I>;
@@ -19,7 +18,7 @@ pub trait HasOptionCorpus<I> where I: Input {
     fn set_corpus(&mut self, corpus: Option<Box<Self::C>>);
 }
 
-pub trait Mutator<I> : HasRand + HasOptionCorpus<I> where I: Input {
+pub trait Mutator<I> : HasRand where I: Input {
     /// Mutate a given input
     fn mutate(&mut self, input: &mut I, stage_idx: i32) -> Result<(), AflError>;
 

src/mutators/scheduled.rs

@@ -20,7 +20,7 @@ pub trait ComposedByMutations<I> where I: Input {
     fn add_mutation(&mut self, mutation: MutationFunction<Self, I>);
 }
 
-pub trait ScheduledMutator<I>: Mutator<I> + ComposedByMutations<I> where I: Input {
+pub trait ScheduledMutator<I>: Mutator<I> + HasOptionCorpus<I> + ComposedByMutations<I> where I: Input {
     /// Compute the number of iterations used to apply stacked mutations
     fn iterations(&mut self, _input: &I) -> u64 {
         1 << (1 + self.rand_mut().below(7))

src/stages/mod.rs

@@ -1,7 +1,19 @@
+pub mod mutational;
+
 use crate::corpus::Testcase;
 use crate::inputs::Input;
+use crate::engines::Engine;
 use crate::AflError;
-/// Stages
-pub trait Stage<InputT: Input> {
-    fn perform(&mut self, input: &dyn Input, entry: &mut Testcase<InputT>) -> Result<(), AflError>;
+
+pub trait HasEngine<'a, I> where I: Input {
+    type E : Engine<'a, I>;
+
+    fn engine(&self) -> &Self::E;
+
+    fn engine_mut(&mut self) -> &mut Self::E;
+}
+
+pub trait Stage<'a, I> : HasEngine<'a, I> where I: Input {
+    /// Run the stage
+    fn perform(&mut self, entry: &mut Testcase<I>) -> Result<(), AflError>;
 }

src/stages/mutational.rs

@@ -1,16 +1,74 @@
-use std::Vec;
+use crate::AflError;
 use crate::mutators::Mutator;
 use crate::inputs::Input;
+use crate::utils::{Rand, HasRand};
+use crate::stages::{Stage, HasEngine};
+use crate::corpus::testcase::Testcase;
+use crate::engines::Engine;
 
-pub struct MutationalStage {
-    mutators: Vec<Box<dyn Mutator>>;
-}
+use std::cell::RefCell;
+use std::rc::Rc;
 
-impl Stage for MutationalStage {
+pub trait MutationalStage<'a, I> : Stage<'a, I> + HasRand where I: Input {
+    fn mutators(&self) -> &Vec<Box<dyn Mutator<I, R = Self::R>>>;
 
-    fn Perform(&mut self, input: &Input, entry: &mut Entry) -> Result<(), AflError> {
-        // TODO: Implement me
-        Err(AflError::NotImplemented("Stage does not perform yet"));
+    fn mutators_mut(&mut self) -> &mut Vec<Box<dyn Mutator<I, R = Self::R>>>;
+
+    fn add_mutator(&mut self, mutator: Box<dyn Mutator<I, R = Self::R>>) {
+        self.mutators_mut().push(mutator);
     }
 
+    fn iterations(&mut self) -> usize {
+        1 + self.rand_mut().below(128) as usize
+    }
+
+    fn perform_mutational(&mut self, entry: Rc<RefCell<Testcase<I>>>) -> Result<(), AflError> {
+        let num = self.iterations();
+        let mut input = entry.borrow_mut().load_input()?.clone();
+
+        for i in 0..num {
+            for m in self.mutators_mut() {
+                m.mutate(&mut input, i as i32)?;
+            }
+
+            let interesting = self.engine_mut().execute(&mut input, entry.clone())?;
+
+            for m in self.mutators_mut() {
+                m.post_exec(interesting, i as i32)?;
+            }
+
+            input = entry.borrow_mut().load_input()?.clone();
+        }
+        
+        Ok(())
+    }
+}
+
+pub struct DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
+    rand: &'a mut R,
+    engine: &'a mut E,
+    mutators: Vec<Box<dyn Mutator<I, R = R>>>
+}
+
+impl<'a, I, R, E> HasRand for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
+    type R = R;
+
+    fn rand(&self) -> &Self::R {
+        &self.rand
+    }
+    fn rand_mut(&mut self) -> &mut Self::R {
+        &mut self.rand
+    }
+}
+
+impl<'a, I, R, E> HasEngine<'a, I> for DefaultMutationalStage<'a, I, R, E> where I: Input, R: Rand, E: Engine<'a, I> {
+    type E = E;
+
+    fn engine(&self) -> &Self::E {
+        self.engine
+    }
+
+    fn engine_mut(&mut self) -> &mut Self::E {
+        self.engine
+    }
 }