diff --git a/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs b/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs index 873e26dea1..2cf5d2c7a0 100644 --- a/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs +++ b/fuzzers/structure_aware/baby_fuzzer_custom_input/src/main.rs @@ -15,6 +15,7 @@ use libafl::{ executors::{inprocess::InProcessExecutor, ExitKind}, feedbacks::{CrashFeedback, MaxMapFeedback}, fuzzer::{Fuzzer, StdFuzzer}, + inputs::MutVecInput, monitors::SimpleMonitor, mutators::scheduled::StdScheduledMutator, observers::StdMapObserver, @@ -138,14 +139,17 @@ pub fn main() { #[cfg(feature = "simple_interface")] let (mapped_mutators, optional_mapped_mutators) = { // Creating mutators that will operate on input.byte_array - let mapped_mutators = - mapped_havoc_mutations(CustomInput::byte_array_mut, CustomInput::byte_array); + let mapped_mutators = mapped_havoc_mutations::, &[u8]>( + CustomInput::byte_array_mut, + CustomInput::byte_array, + ); // Creating mutators that will operate on input.optional_byte_array - let optional_mapped_mutators = optional_mapped_havoc_mutations( - CustomInput::optional_byte_array_mut, - CustomInput::optional_byte_array, - ); + let optional_mapped_mutators = + optional_mapped_havoc_mutations::<_, Option>, Option<&[u8]>>( + CustomInput::optional_byte_array_mut, + CustomInput::optional_byte_array, + ); (mapped_mutators, optional_mapped_mutators) }; diff --git a/libafl/src/mutators/havoc_mutations.rs b/libafl/src/mutators/havoc_mutations.rs index 0e278e0475..1ddf945c42 100644 --- a/libafl/src/mutators/havoc_mutations.rs +++ b/libafl/src/mutators/havoc_mutations.rs @@ -1,21 +1,28 @@ //! [`crate::mutators::Mutator`] collection equivalent to AFL++'s havoc mutations -use libafl_bolts::tuples::{Map, Merge}; +use libafl_bolts::{ + merge_tuple_list_type, + tuples::{Map, Merge}, +}; use tuple_list::{tuple_list, tuple_list_type}; -use crate::mutators::{ - mapping::{ - MappedInputFunctionMappingMutator, OptionMappingMutator, - ToMappedInputFunctionMappingMutatorMapper, ToOptionMappingMutatorMapper, - }, - mutations::{ - BitFlipMutator, ByteAddMutator, ByteDecMutator, ByteFlipMutator, ByteIncMutator, - ByteInterestingMutator, ByteNegMutator, ByteRandMutator, BytesCopyMutator, - BytesDeleteMutator, BytesExpandMutator, BytesInsertCopyMutator, BytesInsertMutator, - BytesRandInsertMutator, BytesRandSetMutator, BytesSetMutator, BytesSwapMutator, - CrossoverInsertMutator, CrossoverReplaceMutator, DwordAddMutator, DwordInterestingMutator, - MappedCrossoverInsertMutator, MappedCrossoverReplaceMutator, QwordAddMutator, - WordAddMutator, WordInterestingMutator, +use crate::{ + inputs::MappedInput, + mutators::{ + mapping::{ + MappedInputFunctionMappingMutator, OptionMappingMutator, + ToMappedInputFunctionMappingMutatorMapper, ToOptionMappingMutatorMapper, + }, + mutations::{ + BitFlipMutator, ByteAddMutator, ByteDecMutator, ByteFlipMutator, ByteIncMutator, + ByteInterestingMutator, ByteNegMutator, ByteRandMutator, BytesCopyMutator, + BytesDeleteMutator, BytesExpandMutator, BytesInsertCopyMutator, BytesInsertMutator, + BytesRandInsertMutator, BytesRandSetMutator, BytesSetMutator, BytesSwapMutator, + CrossoverInsertMutator, CrossoverReplaceMutator, DwordAddMutator, + DwordInterestingMutator, MappedCrossoverInsertMutator, MappedCrossoverReplaceMutator, + QwordAddMutator, WordAddMutator, WordInterestingMutator, + }, + IntoOptionBytes, }, }; @@ -52,110 +59,83 @@ pub type HavocMutationsNoCrossoverType = tuple_list_type!( pub type HavocCrossoverType = tuple_list_type!(CrossoverInsertMutator, CrossoverReplaceMutator); /// Tuple type of the mutations that compose the Havoc mutator's crossover mutations for mapped input types -pub type MappedHavocCrossoverType = tuple_list_type!( - MappedCrossoverInsertMutator, - MappedCrossoverReplaceMutator, +pub type MappedHavocCrossoverType = tuple_list_type!( + MappedCrossoverInsertMutator, + MappedCrossoverReplaceMutator, ); /// Tuple type of the mutations that compose the Havoc mutator -pub type HavocMutationsType = tuple_list_type!( - BitFlipMutator, - ByteFlipMutator, - ByteIncMutator, - ByteDecMutator, - ByteNegMutator, - ByteRandMutator, - ByteAddMutator, - WordAddMutator, - DwordAddMutator, - QwordAddMutator, - ByteInterestingMutator, - WordInterestingMutator, - DwordInterestingMutator, - BytesDeleteMutator, - BytesDeleteMutator, - BytesDeleteMutator, - BytesDeleteMutator, - BytesExpandMutator, - BytesInsertMutator, - BytesRandInsertMutator, - BytesSetMutator, - BytesRandSetMutator, - BytesCopyMutator, - BytesInsertCopyMutator, - BytesSwapMutator, - CrossoverInsertMutator, - CrossoverReplaceMutator, -); +pub type HavocMutationsType = + merge_tuple_list_type!(HavocMutationsNoCrossoverType, HavocCrossoverType); /// Tuple type of the mutations that compose the Havoc mutator for mapped input types -pub type MappedHavocMutationsType = tuple_list_type!( - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, +pub type MappedHavocMutationsType = tuple_list_type!( + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, ); /// Tuple type of the mutations that compose the Havoc mutator for mapped input types, for optional byte array input parts -pub type OptionMappedHavocMutationsType = tuple_list_type!( - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, - MappedInputFunctionMappingMutator, F1, II>, +pub type OptionMappedHavocMutationsType = tuple_list_type!( + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, + MappedInputFunctionMappingMutator, IO, II1>, MappedInputFunctionMappingMutator< - OptionMappingMutator>, - F1, - II, + OptionMappingMutator>, + IO, + II1, >, MappedInputFunctionMappingMutator< - OptionMappingMutator>, - F1, - II, - >, + OptionMappingMutator>, + IO, + II1, + > ); /// Get the mutations that compose the Havoc mutator (only applied to single inputs) @@ -200,28 +180,30 @@ pub fn havoc_crossover() -> HavocCrossoverType { } /// Get the mutations that compose the Havoc mutator's crossover strategy with custom corpus extraction logic -pub fn havoc_crossover_with_corpus_mapper( - input_mapper: F, -) -> MappedHavocCrossoverType +#[must_use] +pub fn havoc_crossover_with_corpus_mapper( + input_mapper: fn(&IO) -> II::Type<'_>, +) -> MappedHavocCrossoverType where - F: Clone + Fn(IO) -> O, + II: IntoOptionBytes, { tuple_list!( - MappedCrossoverInsertMutator::new(input_mapper.clone()), - MappedCrossoverReplaceMutator::new(input_mapper.clone()), + MappedCrossoverInsertMutator::new(input_mapper), + MappedCrossoverReplaceMutator::new(input_mapper), ) } /// Get the mutations that compose the Havoc mutator's crossover strategy with custom corpus extraction logic -pub fn havoc_crossover_with_corpus_mapper_optional( - input_mapper: F, -) -> MappedHavocCrossoverType +#[must_use] +pub fn havoc_crossover_with_corpus_mapper_optional( + input_mapper: fn(&IO) -> II::Type<'_>, +) -> MappedHavocCrossoverType where - F: Clone, + II: IntoOptionBytes, { tuple_list!( - MappedCrossoverInsertMutator::new(input_mapper.clone()), - MappedCrossoverReplaceMutator::new(input_mapper.clone()), + MappedCrossoverInsertMutator::new(input_mapper), + MappedCrossoverReplaceMutator::new(input_mapper), ) } @@ -235,13 +217,13 @@ pub fn havoc_mutations() -> HavocMutationsType { /// /// Check the example fuzzer for details on how to use this. #[must_use] -pub fn mapped_havoc_mutations( - current_input_mapper: F1, - input_from_corpus_mapper: F2, -) -> MappedHavocMutationsType +pub fn mapped_havoc_mutations( + current_input_mapper: fn(&mut IO) -> II1::Type<'_>, + input_from_corpus_mapper: fn(&IO) -> II2::Type<'_>, +) -> MappedHavocMutationsType where - F1: Clone + FnMut(IO1) -> II, - F2: Clone + Fn(IO2) -> O, + II1: MappedInput, + II2: IntoOptionBytes, { havoc_mutations_no_crossover() .merge(havoc_crossover_with_corpus_mapper(input_from_corpus_mapper)) @@ -254,13 +236,13 @@ where /// /// Check the example fuzzer for details on how to use this. #[must_use] -pub fn optional_mapped_havoc_mutations( - current_input_mapper: F1, - input_from_corpus_mapper: F2, -) -> OptionMappedHavocMutationsType +pub fn optional_mapped_havoc_mutations( + current_input_mapper: fn(&mut IO) -> II1::Type<'_>, + input_from_corpus_mapper: fn(&IO) -> II2::Type<'_>, +) -> OptionMappedHavocMutationsType where - F1: Clone + FnMut(IO1) -> II, - F2: Clone + Fn(IO2) -> O, + II1: MappedInput, + II2: IntoOptionBytes, { havoc_mutations_no_crossover() .merge(havoc_crossover_with_corpus_mapper_optional( @@ -271,3 +253,68 @@ where current_input_mapper, )) } + +#[cfg(test)] +mod tests { + use std::string::{String, ToString}; + + use libafl_bolts::rands::StdRand; + use serde::{Deserialize, Serialize}; + + use super::{mapped_havoc_mutations, MappedHavocMutationsType}; + use crate::{ + corpus::{Corpus, CorpusId, InMemoryCorpus}, + inputs::{Input, MutVecInput}, + mutators::{DefaultMutators, MutationResult, StdScheduledMutator, Vec}, + prelude::Mutator as _, + state::StdState, + }; + + #[test] + fn test_default_mutators_custom_implementation() { + #[derive(Debug, Deserialize, Serialize, SerdeAny, Clone)] + struct CustomInput { + vec: Vec, + } + + impl CustomInput { + fn vec_mut(&mut self) -> MutVecInput<'_> { + (&mut self.vec).into() + } + fn vec(&self) -> &[u8] { + &self.vec + } + } + impl DefaultMutators for CustomInput { + type Type = MappedHavocMutationsType, &'static [u8]>; + + fn default_mutators() -> Self::Type { + mapped_havoc_mutations(Self::vec_mut, Self::vec) + } + } + + impl Input for CustomInput { + fn generate_name(&self, _id: Option) -> String { + "CustomInput".to_string() + } + } + let mut input = CustomInput { + vec: vec![0x1, 0x2, 0x3], + }; + let mutations = CustomInput::default_mutators(); + let mut scheduler = StdScheduledMutator::new(mutations); + let mut corpus = InMemoryCorpus::new(); + corpus.add(input.clone().into()).unwrap(); + let mut state = StdState::new( + StdRand::new(), + corpus, + InMemoryCorpus::new(), + &mut (), + &mut (), + ) + .unwrap(); + + let res = scheduler.mutate(&mut state, &mut input).unwrap(); + assert_eq!(res, MutationResult::Mutated); + } +} diff --git a/libafl/src/mutators/mapping.rs b/libafl/src/mutators/mapping.rs index 27dc7990d5..f6aa890be0 100644 --- a/libafl/src/mutators/mapping.rs +++ b/libafl/src/mutators/mapping.rs @@ -1,6 +1,5 @@ //! Allowing mixing and matching between [`Mutator`] and [`crate::inputs::Input`] types. use alloc::borrow::Cow; -use core::marker::PhantomData; use libafl_bolts::{tuples::MappingFunctor, Named}; @@ -58,6 +57,7 @@ pub struct FunctionMappingMutator { impl FunctionMappingMutator { /// Creates a new [`FunctionMappingMutator`] + #[must_use] pub fn new(mapper: F, inner: M) -> Self where M: Named, @@ -137,6 +137,7 @@ pub struct ToFunctionMappingMutatorMapper { impl ToFunctionMappingMutatorMapper { /// Creates a new [`ToFunctionMappingMutatorMapper`] + #[must_use] pub fn new(mapper: F) -> Self { Self { mapper } } @@ -188,16 +189,22 @@ where /// assert_eq!(input, (vec![2],)); /// ``` #[derive(Debug)] -pub struct MappedInputFunctionMappingMutator { - mapper: F, +pub struct MappedInputFunctionMappingMutator +where + II: MappedInput, +{ + mapper: for<'a> fn(&'a mut IO) -> II::Type<'a>, inner: M, name: Cow<'static, str>, - phantom: PhantomData, } -impl MappedInputFunctionMappingMutator { +impl MappedInputFunctionMappingMutator +where + II: MappedInput, +{ /// Creates a new [`MappedInputFunctionMappingMutator`] - pub fn new(mapper: F, inner: M) -> Self + #[must_use] + pub fn new(mapper: for<'a> fn(&'a mut IO) -> II::Type<'a>, inner: M) -> Self where M: Named, { @@ -210,16 +217,14 @@ impl MappedInputFunctionMappingMutator { mapper, inner, name, - phantom: PhantomData, } } } -impl Mutator for MappedInputFunctionMappingMutator +impl Mutator for MappedInputFunctionMappingMutator where for<'a> M: Mutator, S>, - for<'a> II: MappedInput + 'a, - for<'a> F: FnMut(&'a mut IO) -> II::Type<'a>, + II: MappedInput, { fn mutate(&mut self, state: &mut S, input: &mut IO) -> Result { let mapped = &mut (self.mapper)(input); @@ -227,7 +232,10 @@ where } } -impl Named for MappedInputFunctionMappingMutator { +impl Named for MappedInputFunctionMappingMutator +where + II: MappedInput, +{ fn name(&self) -> &Cow<'static, str> { &self.name } @@ -271,33 +279,33 @@ impl Named for MappedInputFunctionMappingMutator { /// assert_eq!(input, (vec![2],)); /// ``` #[derive(Debug)] -pub struct ToMappedInputFunctionMappingMutatorMapper { - mapper: F, - phantom: PhantomData, +pub struct ToMappedInputFunctionMappingMutatorMapper +where + II: MappedInput, +{ + mapper: for<'a> fn(&'a mut IO) -> II::Type<'a>, } -impl ToMappedInputFunctionMappingMutatorMapper { +impl ToMappedInputFunctionMappingMutatorMapper +where + II: MappedInput, +{ /// Creates a new [`ToMappedInputFunctionMappingMutatorMapper`] - pub fn new(mapper: F) -> Self - where - F: FnMut(IO) -> II, - { - Self { - mapper, - phantom: PhantomData, - } + #[must_use] + pub fn new(mapper: for<'a> fn(&'a mut IO) -> II::Type<'a>) -> Self { + Self { mapper } } } -impl MappingFunctor for ToMappedInputFunctionMappingMutatorMapper +impl MappingFunctor for ToMappedInputFunctionMappingMutatorMapper where - F: Clone, M: Named, + II: MappedInput, { - type Output = MappedInputFunctionMappingMutator; + type Output = MappedInputFunctionMappingMutator; fn apply(&mut self, from: M) -> Self::Output { - MappedInputFunctionMappingMutator::new(self.mapper.clone(), from) + MappedInputFunctionMappingMutator::new(self.mapper, from) } } @@ -339,6 +347,7 @@ pub struct OptionMappingMutator { impl OptionMappingMutator { /// Creates a new [`OptionMappingMutator`] + #[must_use] pub fn new(inner: M) -> Self where M: Named, diff --git a/libafl/src/mutators/mod.rs b/libafl/src/mutators/mod.rs index 73ab635a1d..5219ae4e23 100644 --- a/libafl/src/mutators/mod.rs +++ b/libafl/src/mutators/mod.rs @@ -397,3 +397,12 @@ impl Named for NopMutator { &Cow::Borrowed("NopMutator") } } + +/// Extensions of [`crate::inputs::Input`]s that have default mutators +pub trait DefaultMutators { + /// The resulting mutator list type + type Type; + /// Get the default mutators for this type + #[must_use] + fn default_mutators() -> Self::Type; +} diff --git a/libafl/src/mutators/mutations.rs b/libafl/src/mutators/mutations.rs index 4c014fca9b..4cb0f56a85 100644 --- a/libafl/src/mutators/mutations.rs +++ b/libafl/src/mutators/mutations.rs @@ -6,7 +6,6 @@ use alloc::{ }; use core::{ cmp::min, - marker::PhantomData, mem::size_of, num::{NonZero, NonZeroUsize}, ops::Range, @@ -1279,9 +1278,12 @@ impl CrossoverReplaceMutator { } } -trait IntoOptionBytes { +/// Anything that can be mapped to `Option<&[u8]>`, ensuring that the lifetime of the input and output are the same +pub trait IntoOptionBytes { + /// Should be Self, ensure that the lifetimes match type Type<'b>; + /// Run the mapping fn into_option_bytes<'a>(self) -> Option<&'a [u8]> where Self: 'a; @@ -1311,28 +1313,32 @@ impl IntoOptionBytes for Option<&[u8]> { /// Crossover insert mutation for inputs mapped to a bytes vector #[derive(Debug)] -pub struct MappedCrossoverInsertMutator { - input_mapper: F, - phantom: PhantomData, +pub struct MappedCrossoverInsertMutator +where + II: IntoOptionBytes, +{ + input_mapper: for<'a> fn(&'a IO) -> II::Type<'a>, } -impl MappedCrossoverInsertMutator { +impl MappedCrossoverInsertMutator +where + II: IntoOptionBytes, +{ /// Creates a new [`MappedCrossoverInsertMutator`] - pub fn new(input_mapper: F) -> Self { - Self { - input_mapper, - phantom: PhantomData, - } + #[must_use] + pub fn new(input_mapper: fn(&IO) -> II::Type<'_>) -> Self { + Self { input_mapper } } } -impl Mutator for MappedCrossoverInsertMutator +impl Mutator for MappedCrossoverInsertMutator where + II: IntoOptionBytes, S: HasCorpus + HasMaxSize + HasRand, + S::Corpus: Corpus, I: HasMutatorBytes, - for<'a> O: IntoOptionBytes, - for<'a> O::Type<'a>: IntoOptionBytes, - for<'a> F: Fn(&'a ::Input) -> ::Type<'a>, + for<'a> II: IntoOptionBytes, + for<'a> II::Type<'a>: IntoOptionBytes, { fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { let size = input.bytes().len(); @@ -1392,7 +1398,10 @@ where } } -impl Named for MappedCrossoverInsertMutator { +impl Named for MappedCrossoverInsertMutator +where + II: IntoOptionBytes, +{ fn name(&self) -> &Cow<'static, str> { static NAME: Cow<'static, str> = Cow::Borrowed("MappedCrossoverInsertMutator"); &NAME @@ -1401,28 +1410,32 @@ impl Named for MappedCrossoverInsertMutator { /// Crossover replace mutation for inputs mapped to a bytes vector #[derive(Debug)] -pub struct MappedCrossoverReplaceMutator { - input_mapper: F, - phantom: PhantomData, +pub struct MappedCrossoverReplaceMutator +where + II: IntoOptionBytes, +{ + input_mapper: for<'a> fn(&'a IO) -> II::Type<'a>, } -impl MappedCrossoverReplaceMutator { +impl MappedCrossoverReplaceMutator +where + II: IntoOptionBytes, +{ /// Creates a new [`MappedCrossoverReplaceMutator`] - pub fn new(input_mapper: F) -> Self { - Self { - input_mapper, - phantom: PhantomData, - } + #[must_use] + pub fn new(input_mapper: fn(&IO) -> II::Type<'_>) -> Self { + Self { input_mapper } } } -impl Mutator for MappedCrossoverReplaceMutator +impl Mutator for MappedCrossoverReplaceMutator where + II: IntoOptionBytes, S: HasCorpus + HasMaxSize + HasRand, + S::Corpus: Corpus, I: HasMutatorBytes, - O: IntoOptionBytes, - for<'a> O::Type<'a>: IntoOptionBytes, - for<'a> F: Fn(&'a ::Input) -> ::Type<'a>, + for<'a> II: IntoOptionBytes, + for<'a> II::Type<'a>: IntoOptionBytes, { fn mutate(&mut self, state: &mut S, input: &mut I) -> Result { let size = input.bytes().len(); @@ -1479,7 +1492,10 @@ where } } -impl Named for MappedCrossoverReplaceMutator { +impl Named for MappedCrossoverReplaceMutator +where + II: IntoOptionBytes, +{ fn name(&self) -> &Cow<'static, str> { static NAME: Cow<'static, str> = Cow::Borrowed("MappedCrossoverReplaceMutator"); &NAME diff --git a/libafl_bolts/src/tuples.rs b/libafl_bolts/src/tuples.rs index 05280f3298..2aaac8715e 100644 --- a/libafl_bolts/src/tuples.rs +++ b/libafl_bolts/src/tuples.rs @@ -865,6 +865,37 @@ impl PlusOne for (Head, Tail) where */ +/// Merge two lists of types created by the [`tuple_list_type`] macro +/// +/// ```rust +/// use libafl_bolts::{merge_tuple_list_type, tuples::{Merge, tuple_list, tuple_list_type}}; +/// #[derive(PartialEq, Debug)] +/// struct T1; +/// #[derive(PartialEq, Debug)] +/// struct T2; +/// #[derive(PartialEq, Debug)] +/// struct T3; +/// #[derive(PartialEq, Debug)] +/// struct T4; +/// #[derive(PartialEq, Debug)] +/// struct T5; +/// +/// type List1 = tuple_list_type!(T1, T2); +/// let list1: List1 = tuple_list!(T1, T2); +/// type List2 = tuple_list_type!(T3, T4, T5); +/// let list2: List2 = tuple_list!(T3, T4, T5); +/// type Combined = merge_tuple_list_type!(List1, List2); +/// let combined: Combined = list1.merge(list2); +/// let manual: Combined = tuple_list!(T1, T2, T3, T4, T5); +/// assert_eq!(combined, manual); +/// ``` +#[macro_export] +macro_rules! merge_tuple_list_type { + ($Type1:ty, $Type2:ty) => { + <$Type1 as $crate::tuples::Merge<$Type2>>::MergeResult + }; +} + #[cfg(test)] mod test { use tuple_list::{tuple_list, tuple_list_type}; diff --git a/libafl_derive/Cargo.toml b/libafl_derive/Cargo.toml index 2cc461b66c..d98e7b08e2 100644 --- a/libafl_derive/Cargo.toml +++ b/libafl_derive/Cargo.toml @@ -25,6 +25,7 @@ proc-macro = true syn = { version = "2.0.77", features = ["full", "extra-traits"] } quote = "1.0.37" proc-macro2 = "1.0.86" +proc-macro-crate = "3.2" [lints] workspace = true diff --git a/libafl_derive/src/lib.rs b/libafl_derive/src/lib.rs index 354bbc9ca4..8a3139b72c 100644 --- a/libafl_derive/src/lib.rs +++ b/libafl_derive/src/lib.rs @@ -42,8 +42,17 @@ )] use proc_macro::TokenStream; -use quote::quote; -use syn::{parse_macro_input, Data::Struct, DeriveInput, Field, Fields::Named, Type}; +use proc_macro_crate::{crate_name, FoundCrate}; +use quote::{format_ident, quote}; +use syn::{ + parse_macro_input, + punctuated::Punctuated, + token::Comma, + Data::Struct, + DeriveInput, Error, Field, + Fields::{Named, Unit, Unnamed}, + GenericArgument, Ident, PathArguments, PathSegment, Type, +}; /// Derive macro to implement `SerdeAny`, to use a type in a `SerdeAnyMap` #[proc_macro_derive(SerdeAny)] @@ -144,3 +153,139 @@ fn libafl_display_field_by_type(it: &Field) -> proc_macro2::TokenStream { write!(f, #fmt, self.#ident)?; } } + +/// TODO +#[proc_macro_derive(HasHavocMutators)] +pub fn derive_has_mutator_bytes(input: TokenStream) -> TokenStream { + let input_ast = parse_macro_input!(input as DeriveInput); + + let struct_name = input_ast.ident.clone(); + + let fields = match extract_fields(input_ast) { + Ok(f) => f, + Err(e) => return e.into_compile_error().into(), + }; + + let (getter_methods, mutator_merge_call) = match create_functions_on_fields(&fields) { + Ok(e) => e, + Err(e) => return e.into_compile_error().into(), + }; + + // required to be able to use it from within libafl — used for testing + let libafl_source = match crate_name("libafl").expect("Could not figure out current crate") { + FoundCrate::Itself => quote! { crate }, + FoundCrate::Name(_) => quote! { libafl }, + }; + + // Generate the impl block + let expanded = quote! { + use #libafl_source::{inputs::MutVecInput, mutators::{Mutator, mapped_havoc_mutations}}; + use libafl_bolts::tuples::{Merge, NamedTuple, tuple_list}; + + impl #struct_name { + #getter_methods + } + + impl HasHavocMutators for #struct_name { + fn havoc_mutators() -> MT { + #mutator_merge_call + } + } + }; + + TokenStream::from(expanded) +} + +fn extract_fields(ast: DeriveInput) -> Result, Error> { + match &ast.data { + Struct(data_struct) => match &data_struct.fields { + Named(fields_named) => Ok(fields_named.named.clone()), + Unnamed(fields_unnamed) => Ok(fields_unnamed.unnamed.clone()), + Unit => Err(Error::new_spanned( + ast, + "HasHavocMutators can not be derived for unit structs", + )), + }, + _ => Err(Error::new_spanned( + ast, + "HasHavocMutators can only be derived for structs", + )), + } +} + +fn create_functions_on_fields( + fields: &Punctuated, +) -> Result<(proc_macro2::TokenStream, proc_macro2::TokenStream), Error> { + let functions_res = fields.iter().map(|field| match field.ty.clone() { + Type::Path(type_path) => { + let segment = type_path.path.segments.last().unwrap(); + if let Some(tokens) = create_functions_on_type(segment, field.ident.as_ref().unwrap()) { + return Ok(tokens); + } + + Err(Error::new_spanned( + segment.ident.clone(), + "HasHavocMutators does not support struct parts of this type", + )) + } + _ => Err(Error::new_spanned( + field, + "HasHavocMutators can only be derived for structs", + )), + }); + + // check if any fields could not be parsed into functions, combine the errors and return them + if let Some(errors) = functions_res + .clone() + .filter(Result::is_err) + .map(Result::unwrap_err) + .reduce(|mut acc, e| { + acc.combine(e); + acc + }) + { + return Err(errors); + } + + Ok(functions_res.map(Result::unwrap).fold( + (quote! {}, quote! { tuple_list!() }), + |(acc1, acc2), (e1, e2)| { + ( + quote! { + #acc1 + #e1 + }, + quote! { #acc2.merge(#e2) }, + ) + }, + )) +} + +fn create_functions_on_type( + segment: &PathSegment, + field_name: &Ident, +) -> Option<(proc_macro2::TokenStream, proc_macro2::TokenStream)> { + if segment.ident == "Vec" { + if let PathArguments::AngleBracketed(args) = &segment.arguments { + if let Some(GenericArgument::Type(Type::Path(arg_type))) = args.args.first() { + let arg_ident = &arg_type.path.segments.last().unwrap().ident; + if arg_ident == "u8" { + let mutable_method_name = format_ident!("{}_mut", field_name); + let immutable_method_name = field_name; + return Some(( + quote! { + pub fn #mutable_method_name(&mut self) -> MutVecInput<'_> { + (&mut self.#field_name).into() + } + pub fn #immutable_method_name(&self) -> &[u8] { + &self.#field_name + } + }, + quote! { mapped_havoc_mutations(Self::#mutable_method_name, Self::#immutable_method_name) }, + )); + } + } + } + } + None +}