diff --git a/src/inputs/bytes.rs b/src/inputs/bytes.rs index 04594c5f37..98bb1cfdc7 100644 --- a/src/inputs/bytes.rs +++ b/src/inputs/bytes.rs @@ -1,6 +1,11 @@ use crate::inputs::Input; use crate::AflError; +pub trait HasBytesVec { + fn bytes(&self) -> &Vec; + fn bytes_mut(&mut self) -> &mut Vec; +} + #[derive(Clone, Debug, Default)] pub struct BytesInput { pub bytes: Vec, @@ -17,6 +22,15 @@ impl Input for BytesInput { } } +impl HasBytesVec for BytesInput { + fn bytes(&self) -> &Vec { + &self.bytes + } + fn bytes_mut(&mut self) -> &mut Vec { + &mut self.bytes + } +} + impl BytesInput { pub fn new(bytes: Vec) -> Self { BytesInput { diff --git a/src/inputs/mod.rs b/src/inputs/mod.rs index af7d21d854..a305e4beef 100644 --- a/src/inputs/mod.rs +++ b/src/inputs/mod.rs @@ -1,5 +1,5 @@ pub mod bytes; -pub use bytes::BytesInput; +pub use bytes::{HasBytesVec, BytesInput}; use std::fs::File; use std::io::Read; diff --git a/src/mutators/scheduled.rs b/src/mutators/scheduled.rs index b79be05ba5..65888afe9c 100644 --- a/src/mutators/scheduled.rs +++ b/src/mutators/scheduled.rs @@ -1,9 +1,11 @@ use crate::mutators::Mutator; use crate::utils::Rand; use crate::corpus::Corpus; -use crate::inputs::{Input, BytesInput}; +use crate::inputs::{Input, HasBytesVec}; use crate::AflError; +use std::marker::PhantomData; + type MutationFunction = fn(&mut MutatorT, &mut InputT) -> Result<(), AflError>; pub trait ScheduledMutator: Mutator { @@ -94,23 +96,24 @@ impl DefaultScheduledMutator { } } -pub fn mutation_bitflip>(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; +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> { +pub struct HavocMutator, InputT: Input + HasBytesVec> { scheduled: ScheduledMutatorT, + _phantom: PhantomData } -impl> Mutator for HavocMutator { +impl, InputT: Input + HasBytesVec> Mutator for HavocMutator { fn rand(&mut self) -> &mut Box { self.scheduled.rand() } - fn mutate(&mut self, input: &mut BytesInput, stage_idx: i32) -> Result<(), AflError> { + fn mutate(&mut self, input: &mut InputT, stage_idx: i32) -> Result<(), AflError> { self.scheduled.mutate(input, stage_idx) } @@ -120,21 +123,23 @@ impl> Mutator for Ha } -impl> HavocMutator { +impl, InputT: Input + HasBytesVec> HavocMutator { pub fn new(mut scheduled: ScheduledMutatorT) -> Self { scheduled.add_mutation(mutation_bitflip); HavocMutator { - scheduled: scheduled + scheduled: scheduled, + _phantom: PhantomData } } } -impl HavocMutator> { +impl HavocMutator, InputT> { pub fn new_default(rand: Box) -> Self { let mut scheduled = DefaultScheduledMutator::new(rand); scheduled.add_mutation(mutation_bitflip); HavocMutator { - scheduled: scheduled + scheduled: scheduled, + _phantom: PhantomData } } } \ No newline at end of file