diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index 65888afe9c..a474512122 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -6,13 +6,27 @@ use crate::AflError; use std::marker::PhantomData; +/// The generic function type that identifies mutations type MutationFunction = fn(&mut MutatorT, &mut InputT) -> Result<(), AflError>; -pub trait ScheduledMutator: Mutator { +pub trait ComposedByMutations { + /// Get a mutation by index + fn mutation_by_idx(&self, index: usize) -> Result, AflError>; + + /// Get the number of mutations + fn mutations_count(&self) -> usize; + + /// Add a mutation + fn add_mutation(&mut self, mutation: MutationFunction); +} + +pub trait ScheduledMutator: Mutator + ComposedByMutations { + /// Computer the number of iterations used to apply stacked mutations fn iterations(&mut self, _input: &InputT) -> u64 { 1 << (1 + self.rand().below(7)) } + /// Get the next mutation to apply fn schedule(&mut self, _input: &InputT) -> Result, AflError> { let count = self.mutations_count() as u64; if count == 0 { @@ -25,12 +39,15 @@ pub trait ScheduledMutator: Mutator { self.mutation_by_idx(idx) } - fn mutation_by_idx(&self, index: usize) -> Result, AflError>; - - fn mutations_count(&self) -> usize; - - fn add_mutation(&mut self, mutation: MutationFunction); - + /// New default implementation for mutate + /// Implementations must forward mutate() to this method + fn scheduled_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)?; + } + Ok(()) + } } pub struct DefaultScheduledMutator { @@ -39,8 +56,7 @@ pub struct DefaultScheduledMutator { mutations: Vec> } -impl ScheduledMutator for DefaultScheduledMutator { - +impl ComposedByMutations for DefaultScheduledMutator { fn mutation_by_idx(&self, index: usize) -> Result, AflError> { if index >= self.mutations.len() { return Err(AflError::Unknown("oob".to_string())); @@ -55,30 +71,28 @@ impl ScheduledMutator for DefaultScheduledMutator) { self.mutations.push(mutation) } +} +impl ScheduledMutator for DefaultScheduledMutator { + // Just use the default methods } impl Mutator for DefaultScheduledMutator { - fn rand(&mut self) -> &mut Box { &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)?; - } - Ok(()) + self.scheduled_mutate(input, _stage_idx) } fn corpus(&mut self) -> &mut Option> { &mut self.corpus } - } impl DefaultScheduledMutator { + /// Create a new DefaultScheduledMutator instance without mutations and corpus pub fn new(rand: Box) -> Self { DefaultScheduledMutator { rand: rand, @@ -87,6 +101,7 @@ impl DefaultScheduledMutator { } } + /// Create a new DefaultScheduledMutator instance specifying mutations and corpus too pub fn new_all(rand: Box, corpus: Option>, mutations: Vec>) -> Self { DefaultScheduledMutator { rand: rand, @@ -96,19 +111,20 @@ impl DefaultScheduledMutator { } } +/// Bitflip mutation for inputs with a bytes vector pub fn mutation_bitflip, InputT: Input + HasBytesVec>(mutator: &mut MutatorT, input: &mut InputT) -> Result<(), AflError> { let bit = mutator.rand().below(input.bytes().len() as u64) as usize; input.bytes_mut()[bit >> 3] ^= (128 >> (bit & 7)) as u8; Ok(()) } -pub struct HavocMutator, InputT: Input + HasBytesVec> { +/// Schedule some selected byte level mutations given a ScheduledMutator type +pub struct HavocBytesMutator, InputT: Input + HasBytesVec> { scheduled: ScheduledMutatorT, _phantom: PhantomData } -impl, InputT: Input + HasBytesVec> Mutator for HavocMutator { - +impl, InputT: Input + HasBytesVec> Mutator for HavocBytesMutator { fn rand(&mut self) -> &mut Box { self.scheduled.rand() } @@ -120,24 +136,25 @@ impl, InputT: Input + HasBytesVec> M fn corpus(&mut self) -> &mut Option> { self.scheduled.corpus() } - } -impl, InputT: Input + HasBytesVec> HavocMutator { +impl, InputT: Input + HasBytesVec> HavocBytesMutator { + /// Create a new HavocBytesMutator instance given a ScheduledMutator to wrap pub fn new(mut scheduled: ScheduledMutatorT) -> Self { scheduled.add_mutation(mutation_bitflip); - HavocMutator { + HavocBytesMutator { scheduled: scheduled, _phantom: PhantomData } } } -impl HavocMutator, InputT> { +impl HavocBytesMutator, InputT> { + /// Create a new HavocBytesMutator instance wrapping DefaultScheduledMutator pub fn new_default(rand: Box) -> Self { let mut scheduled = DefaultScheduledMutator::new(rand); scheduled.add_mutation(mutation_bitflip); - HavocMutator { + HavocBytesMutator { scheduled: scheduled, _phantom: PhantomData }