scheduled mutator

src/inputs/bytes.rs

@@ -3,7 +3,7 @@ use crate::AflError;
 
 #[derive(Clone, Debug, Default)]
 pub struct BytesInput {
-    bytes: Vec<u8>,
+    pub bytes: Vec<u8>,
 }
 
 impl Input for BytesInput {

src/inputs/mod.rs

@@ -1,4 +1,5 @@
 pub mod bytes;
+pub use bytes::BytesInput;
 
 use std::fs::File;
 use std::io::Read;

src/mutators/mod.rs

@@ -1,24 +1,21 @@
 use crate::inputs::Input;
+use crate::utils::Rand;
+use crate::corpus::Corpus;
 use crate::AflError;
 
 pub mod scheduled;
 pub use scheduled::ScheduledMutator;
 
-pub trait Mutator {
+pub trait Mutator<InputT : Input> {
-    //fn rand(&self) -> &Box<dyn Rand>;
-    //fn rand_mut(&self) -> &mut Box<dyn Rand>;
 
-    fn mutate(&mut self, input: &mut dyn Input) -> Result<(), AflError> {
+    fn rand(&mut self) -> &mut Box<dyn Rand>;
-        self.mutate_at((-1) as i32, input)
-    }
 
-    fn mutate_at(&mut self, stage_idx: i32, input: &mut dyn Input) -> Result<(), AflError>;
+    fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError>;
 
-    fn post_exec(&mut self, is_interesting: bool) -> Result<(), AflError> {
+    fn post_exec(&mut self, _is_interesting: bool, _stage_idx: i32) -> Result<(), AflError> {
-        self.post_exec_at((-1) as i32, is_interesting)
-    }
-
-    fn post_exec_at(&mut self, _stage_idx: i32, _is_interesting: bool) -> Result<(), AflError> {
         Ok(())
     }
+
+    fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>>;
+
 }

src/mutators/scheduled.rs

@@ -1,34 +1,45 @@
 use crate::mutators::Mutator;
-use crate::inputs::Input;
+use crate::utils::Rand;
+use crate::corpus::Corpus;
+use crate::inputs::{Input, BytesInput};
 use crate::AflError;
 
-type MutationFunction = fn(&mut dyn ScheduledMutator, &mut dyn Input) -> Result<(), AflError>;
+type MutationFunction<MutatorT, InputT> = fn(&mut MutatorT, &mut InputT) -> Result<(), AflError>;
 
-pub trait ScheduledMutator: Mutator {
+pub trait ScheduledMutator<InputT : Input>: Mutator<InputT> {
-    fn iterations(&self, _input: &dyn Input) -> u64 {
+    fn iterations(&mut self, _input: &InputT) -> u64 {
-        //1 << (1 + self.rand_mut().below(7))
+        1 << (1 + self.rand().below(7))
-        return 1;
     }
 
-    fn schedule(&self, _input: &dyn Input) -> Result<MutationFunction, AflError> {
+    fn schedule(&mut self, _input: &InputT) -> Result<MutationFunction<Self, InputT>, AflError> {
-        if self.mutations_count() == 0 {
+        let count = self.mutations_count() as u64;
+        if count == 0 {
             return Err(AflError::Empty("no mutations".to_string()));
         }
-        self.get_mutation_by_idx(1 /* self.rand_mut().below(self.mutations_count()) */)
+        let idx;
+        {
+            idx = self.rand().below(count) as usize;
+        }
+        self.mutation_by_idx(idx)
     }
 
-    fn get_mutation_by_idx(&self, index: usize) -> Result<MutationFunction, AflError>;
+    fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, InputT>, AflError>;
 
     fn mutations_count(&self) -> usize;
+
+    fn add_mutation(&mut self, mutation: MutationFunction<Self, InputT>);
+
 }
 
-pub struct BaseScheduledMutator {
+pub struct DefaultScheduledMutator<InputT : Input> {
-    mutations: Vec<MutationFunction>
+    rand: Box<dyn Rand>,
+    corpus: Option<Box<dyn Corpus>>,
+    mutations: Vec<MutationFunction<Self, InputT>>
 }
 
-impl ScheduledMutator for BaseScheduledMutator {
+impl<InputT : Input> ScheduledMutator<InputT> for DefaultScheduledMutator<InputT> {
 
-    fn get_mutation_by_idx(&self, index: usize) -> Result<MutationFunction, AflError> {
+    fn mutation_by_idx(&self, index: usize) -> Result<MutationFunction<Self, InputT>, AflError> {
         if index >= self.mutations.len() {
             return Err(AflError::Unknown("oob".to_string()));
         }
@@ -39,11 +50,19 @@ impl ScheduledMutator for BaseScheduledMutator {
         self.mutations.len()
     }
 
+    fn add_mutation(&mut self, mutation: MutationFunction<Self, InputT>) {
+        self.mutations.push(mutation)
     }
 
-impl Mutator for BaseScheduledMutator {
+}
 
-    fn mutate_at(&mut self, _stage_idx: i32, input: &mut dyn Input) -> Result<(), AflError> {
+impl<InputT : Input> Mutator<InputT> for DefaultScheduledMutator<InputT> {
+
+    fn rand(&mut self) -> &mut Box<dyn Rand> {
+        &mut self.rand
+    }
+
+    fn mutate(&mut self, input: &mut InputT, _stage_idx: i32) -> Result<(), AflError> {
         let num = self.iterations(input);
         for _ in 0..num {
             self.schedule(input)?(self, input)?;
@@ -51,4 +70,71 @@ impl Mutator for BaseScheduledMutator {
         Ok(())
     }
 
+    fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>> {
+        &mut self.corpus
+    }
+
+}
+
+impl<InputT : Input> DefaultScheduledMutator<InputT> {
+    pub fn new(rand: Box<dyn Rand>) -> Self {
+        DefaultScheduledMutator {
+            rand: rand,
+            corpus: None,
+            mutations: vec![]
+        }
+    }
+
+    pub fn new_all(rand: Box<dyn Rand>, corpus: Option<Box<dyn Corpus>>, mutations: Vec<MutationFunction<Self, InputT>>) -> Self {
+        DefaultScheduledMutator {
+            rand: rand,
+            corpus: corpus,
+            mutations: mutations
+        }
+    }
+}
+
+pub fn mutation_bitflip<MutatorT: Mutator<BytesInput>>(mutator: &mut MutatorT, input: &mut BytesInput) -> Result<(), AflError> {
+    let bit = mutator.rand().below(input.bytes.len() as u64) as usize;
+    input.bytes[bit >> 3] ^= (128 >> (bit & 7)) as u8;
+    Ok(())
+}
+
+pub struct HavocMutator<ScheduledMutatorT: ScheduledMutator<BytesInput>> {
+    scheduled: ScheduledMutatorT,
+}
+
+impl<ScheduledMutatorT: ScheduledMutator<BytesInput>> Mutator<BytesInput> for HavocMutator<ScheduledMutatorT> {
+
+    fn rand(&mut self) -> &mut Box<dyn Rand> {
+        self.scheduled.rand()
+    }
+
+    fn mutate(&mut self, input: &mut BytesInput, stage_idx: i32) -> Result<(), AflError> {
+        self.scheduled.mutate(input, stage_idx)
+    }
+
+    fn corpus(&mut self) -> &mut Option<Box<dyn Corpus>> {
+        self.scheduled.corpus()
+    }
+
+}
+
+impl<ScheduledMutatorT: ScheduledMutator<BytesInput>> HavocMutator<ScheduledMutatorT> {
+    pub fn new(mut scheduled: ScheduledMutatorT) -> Self {
+        scheduled.add_mutation(mutation_bitflip);
+        HavocMutator {
+            scheduled: scheduled
+        }
+    }
+}
+
+impl HavocMutator<DefaultScheduledMutator<BytesInput>> {
+    pub fn new_default(rand: Box<dyn Rand>) -> Self {
+        let mut scheduled = DefaultScheduledMutator::new(rand);
+        scheduled.add_mutation(mutation_bitflip);
+        HavocMutator {
+            scheduled: scheduled
+        }
+    }
 }