Introduce StateAwareMappingMutator

crates/libafl/src/mutators/havoc_mutations.rs

@@ -6,7 +6,7 @@ use libafl_bolts::{
 };
 
 use crate::mutators::{
-    mapping::{ToMappingMutator, ToOptionalMutator},
+    mapping::{ToMappingMutator, ToOptionalMutator, ToStateAwareMappingMutator},
     mutations::{
         BitFlipMutator, ByteAddMutator, ByteDecMutator, ByteFlipMutator, ByteIncMutator,
         ByteInterestingMutator, ByteNegMutator, ByteRandMutator, BytesCopyMutator,
@@ -66,6 +66,12 @@ pub type MappedHavocMutationsType<F1, F2, I, O> = map_tuple_list_type!(
     ToMappingMutator<F1>
 );
 
+/// Tuple type of the mutations that compose the Havoc mutator for state-aware-mapped input types
+pub type StateAwareMappedHavocMutationsType<F1, F2, I, O> = map_tuple_list_type!(
+    merge_tuple_list_type!(HavocMutationsNoCrossoverType, MappedHavocCrossoverType<F2,I, O>),
+    ToStateAwareMappingMutator<F1>
+);
+
 /// Tuple type of the mutations that compose the Havoc mutator for mapped input types, for optional byte array input parts
 pub type OptionMappedHavocMutationsType<F1, F2, I, O> = map_tuple_list_type!(
     map_tuple_list_type!(
@@ -165,6 +171,27 @@ where
         .map(ToMappingMutator::new(current_input_mapper))
 }
 
+/// Get the mutations that compose the Havoc mutator for state-aware-mapped input types
+///
+/// Check the example fuzzer for details on how to use this.
+/// Check the docs of [`crate::mutators::mapping::StateAwareMappingMutator`] for how mapping works internally.
+#[must_use]
+pub fn state_aware_mapped_havoc_mutations<'a, F1, F2, IO1, IO2, II, O, S>(
+    current_input_mapper: F1,
+    input_from_corpus_mapper: F2,
+) -> StateAwareMappedHavocMutationsType<F1, F2, IO1, O>
+where
+    F1: Clone + FnMut(&'a mut IO1, &'a mut S) -> (Option<&'a mut II>, &'a mut S),
+    F2: Clone + Fn(&IO2) -> &O,
+    II: 'a,
+    IO1: 'a,
+    S: 'a,
+{
+    havoc_mutations_no_crossover()
+        .merge(havoc_crossover_with_corpus_mapper(input_from_corpus_mapper))
+        .map(ToStateAwareMappingMutator::new(current_input_mapper))
+}
+
 /// Get the mutations that compose the Havoc mutator for mapped input types, for optional input parts
 ///
 /// Check the example fuzzer for details on how to use this.

crates/libafl/src/mutators/mapping.rs

@@ -5,13 +5,16 @@ use libafl_bolts::{Named, tuples::MappingFunctor};
 
 use crate::{
     Error,
+    corpus::CorpusId,
     mutators::{MutationResult, Mutator},
 };
 
 /// Mapping [`Mutator`] using a function returning a reference.
 ///
 /// Allows using [`Mutator`]s for a certain type on (parts of) other input types that can be mapped to this type.
 ///
+/// For a more flexible alternative, which allows access to `state`, see [`StateAwareMappingMutator`].
+///
 /// # Example
 #[cfg_attr(feature = "std", doc = " ```")]
 #[cfg_attr(not(feature = "std"), doc = " ```ignore")]
@@ -74,11 +77,7 @@ where
         self.inner.mutate(state, (self.mapper)(input))
     }
     #[inline]
-    fn post_exec(
-        &mut self,
-        state: &mut S,
-        new_corpus_id: Option<crate::corpus::CorpusId>,
-    ) -> Result<(), Error> {
+    fn post_exec(&mut self, state: &mut S, new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
         self.inner.post_exec(state, new_corpus_id)
     }
 }
@@ -89,7 +88,7 @@ impl<M, F> Named for MappingMutator<M, F> {
     }
 }
 
-/// Mapper to use to map a [`tuple_list`] of [`Mutator`]s using [`ToMappingMutator`]s.
+/// Mapper to use to map a [`tuple_list`] of [`Mutator`]s using [`MappingMutator`]s.
 ///
 /// See the explanation of [`MappingMutator`] for details.
 ///
@@ -211,11 +210,7 @@ where
         }
     }
     #[inline]
-    fn post_exec(
-        &mut self,
-        state: &mut S,
-        new_corpus_id: Option<crate::corpus::CorpusId>,
-    ) -> Result<(), Error> {
+    fn post_exec(&mut self, state: &mut S, new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
         self.inner.post_exec(state, new_corpus_id)
     }
 }
@@ -272,3 +267,183 @@ where
         OptionalMutator::new(from)
     }
 }
+
+/// Mapping [`Mutator`] using a function returning a reference.
+///
+/// Allows using [`Mutator`]s for a certain type on (parts of) other input types that can be mapped to this type.
+///
+/// Provides access to the state. If [`Option::None`] is returned from the mapping function, the mutator is returning [`MutationResult::Skipped`].
+///
+/// # Example
+#[cfg_attr(feature = "std", doc = " ```")]
+#[cfg_attr(not(feature = "std"), doc = " ```ignore")]
+/// use std::vec::Vec;
+///
+/// use libafl::{
+///     mutators::{ByteIncMutator, MutationResult, Mutator, StateAwareMappingMutator},
+///     state::{HasRand, NopState},
+/// };
+/// use libafl_bolts::rands::Rand as _;
+///
+/// #[derive(Debug, PartialEq)]
+/// struct CustomInput(Vec<u8>);
+///
+/// impl CustomInput {
+///     pub fn possibly_vec<'a, S: HasRand>(
+///         &'a mut self,
+///         state: &'a mut S,
+///     ) -> (Option<&'a mut Vec<u8>>, &'a mut S) {
+///         // we have access to the state
+///         if state.rand_mut().coinflip(0.5) {
+///             // If this input cannot be mutated with the outer mutator, return None
+///             // e.g. because the input doesn't contain the field this mutator is supposed to mutate
+///             (None, state)
+///         } else {
+///             // else, return the type that the outer mutator can mutate
+///             (Some(&mut self.0), state)
+///         }
+///     }
+/// }
+///
+/// // construct a mutator that works on &mut Vec<u8> (since it impls `HasMutatorBytes`)
+/// let inner = ByteIncMutator::new();
+/// // construct a mutator that works on &mut CustomInput
+/// let mut outer = StateAwareMappingMutator::new(CustomInput::possibly_vec, inner);
+///
+/// let mut input = CustomInput(vec![1]);
+///
+/// let mut state: NopState<CustomInput> = NopState::new();
+/// let res = outer.mutate(&mut state, &mut input).unwrap();
+/// if res == MutationResult::Mutated {
+///     assert_eq!(input, CustomInput(vec![2],));
+/// } else {
+///     assert_eq!(input, CustomInput(vec![1],));
+/// }
+/// ```
+#[derive(Debug)]
+pub struct StateAwareMappingMutator<M, F> {
+    mapper: F,
+    inner: M,
+    name: Cow<'static, str>,
+}
+
+impl<M, F> StateAwareMappingMutator<M, F> {
+    /// Creates a new [`StateAwareMappingMutator`]
+    pub fn new(mapper: F, inner: M) -> Self
+    where
+        M: Named,
+    {
+        let name = Cow::Owned(format!("StateAwareMappingMutator<{}>", inner.name()));
+        Self {
+            mapper,
+            inner,
+            name,
+        }
+    }
+}
+
+impl<M, S, F, IO, II> Mutator<IO, S> for StateAwareMappingMutator<M, F>
+where
+    F: for<'a> FnMut(&'a mut IO, &'a mut S) -> (Option<&'a mut II>, &'a mut S),
+    M: Mutator<II, S>,
+{
+    fn mutate(&mut self, state: &mut S, input: &mut IO) -> Result<MutationResult, Error> {
+        let (mapped, state) = (self.mapper)(input, state);
+        match mapped {
+            Some(mapped) => self.inner.mutate(state, mapped),
+            None => Ok(MutationResult::Skipped),
+        }
+    }
+    #[inline]
+    fn post_exec(&mut self, state: &mut S, new_corpus_id: Option<CorpusId>) -> Result<(), Error> {
+        self.inner.post_exec(state, new_corpus_id)
+    }
+}
+
+impl<M, F> Named for StateAwareMappingMutator<M, F> {
+    fn name(&self) -> &Cow<'static, str> {
+        &self.name
+    }
+}
+
+/// Mapper to use to map a [`tuple_list`] of [`Mutator`]s using [`StateAwareMappingMutator`]s.
+///
+/// See the explanation of [`StateAwareMappingMutator`] for details.
+///
+/// # Example
+#[cfg_attr(feature = "std", doc = " ```")]
+#[cfg_attr(not(feature = "std"), doc = " ```ignore")]
+/// use std::vec::Vec;
+///
+/// use libafl::{
+///     mutators::{
+///         ByteIncMutator, MutationResult, MutatorsTuple, ToStateAwareMappingMutator,
+///     },
+///     state::{HasRand, NopState},
+/// };
+///  
+/// use libafl_bolts::{
+///     tuples::{tuple_list, Map},
+///     rands::Rand as _,
+/// };
+///  
+/// #[derive(Debug, PartialEq)]
+/// struct CustomInput(Vec<u8>);
+///  
+/// impl CustomInput {
+///     pub fn possibly_vec<'a, S: HasRand>(
+///         &'a mut self,
+///         state: &'a mut S,
+///     ) -> (Option<&'a mut Vec<u8>>, &'a mut S) {
+///         // we have access to the state
+///         if state.rand_mut().coinflip(0.5) {
+///             // If this input cannot be mutated with the outer mutator, return None
+///             // e.g. because the input doesn't contain the field this mutator is supposed to mutate
+///             (None, state)
+///         } else {
+///             // else, return the type that the outer mutator can mutate
+///             (Some(&mut self.0), state)
+///         }
+///     }
+/// }
+///  
+/// // construct a mutator that works on &mut Vec<u8> (since it impls `HasMutatorBytes`)
+/// let mutators = tuple_list!(ByteIncMutator::new(), ByteIncMutator::new());
+/// // construct a mutator that works on &mut CustomInput
+/// let mut mapped_mutators =
+///     mutators.map(ToStateAwareMappingMutator::new(CustomInput::possibly_vec));
+///  
+/// let mut input = CustomInput(vec![1]);
+///  
+/// let mut state: NopState<CustomInput> = NopState::new();
+/// let res = mapped_mutators.mutate_all(&mut state, &mut input).unwrap();
+/// if res == MutationResult::Mutated {
+///     // no way of knowing if either or both mutated
+///     assert!(input.0 == vec![2] || input.0 == vec![3]);
+/// } else {
+///     assert_eq!(input, CustomInput(vec![1],));
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToStateAwareMappingMutator<F> {
+    mapper: F,
+}
+
+impl<F> ToStateAwareMappingMutator<F> {
+    /// Creates a new [`ToStateAwareMappingMutator`]
+    pub fn new(mapper: F) -> Self {
+        Self { mapper }
+    }
+}
+
+impl<M, F> MappingFunctor<M> for ToStateAwareMappingMutator<F>
+where
+    F: Clone,
+    M: Named,
+{
+    type Output = StateAwareMappingMutator<M, F>;
+
+    fn apply(&mut self, from: M) -> Self::Output {
+        StateAwareMappingMutator::new(self.mapper.clone(), from)
+    }
+}

crates/libafl/src/mutators/numeric.rs

@@ -4,13 +4,13 @@ use alloc::borrow::Cow;
 use core::marker::PhantomData;
 
 use libafl_bolts::{
-    Error, Named,
+    Error, Named, map_tuple_list_type, merge_tuple_list_type,
     rands::Rand,
     tuples::{Map as _, Merge},
 };
 use tuple_list::{tuple_list, tuple_list_type};
 
-use super::{MappingMutator, MutationResult, Mutator, ToMappingMutator};
+use super::{MutationResult, Mutator, ToMappingMutator, ToStateAwareMappingMutator};
 use crate::{
     corpus::Corpus,
     random_corpus_id_with_disabled,
@@ -80,14 +80,15 @@ pub fn int_mutators() -> IntMutatorsType {
 }
 
 /// Mapped mutators for integer-like inputs
-pub type MappedIntMutatorsType<F1, F2, I> = tuple_list_type!(
-    MappingMutator<BitFlipMutator,F1>,
-    MappingMutator<NegateMutator,F1>,
-    MappingMutator<IncMutator,F1>,
-    MappingMutator<DecMutator,F1>,
-    MappingMutator<TwosComplementMutator,F1>,
-    MappingMutator<RandMutator,F1>,
-    MappingMutator<MappedCrossoverMutator<F2, I>,F1>
+pub type MappedIntMutatorsType<F1, F2, I> = map_tuple_list_type!(
+    merge_tuple_list_type!(IntMutatorsNoCrossoverType, MappedIntMutatorsCrossoverType<F2, I>),
+    ToMappingMutator<F1>
+);
+
+/// State-aware Mapped mutators for integer-like inputs
+pub type StateAwareMappedIntMutatorsType<F1, F2, I> = map_tuple_list_type!(
+    merge_tuple_list_type!(IntMutatorsNoCrossoverType, MappedIntMutatorsCrossoverType<F2, I>),
+    ToStateAwareMappingMutator<F1>
 );
 
 /// Mapped mutators for integer-like inputs
@@ -104,6 +105,25 @@ where
         .merge(mapped_int_mutators_crossover(input_from_corpus_mapper))
         .map(ToMappingMutator::new(current_input_mapper))
 }
+
+/// State-awareMapped mutators for integer-like inputs
+///
+/// Modelled after the applicable mutators from [`super::havoc_mutations::havoc_mutations`]
+pub fn state_aware_mapped_int_mutators<'a, F1, F2, IO, II, S>(
+    current_input_mapper: F1,
+    input_from_corpus_mapper: F2,
+) -> StateAwareMappedIntMutatorsType<F1, F2, IO>
+where
+    F1: Clone + FnMut(&'a mut IO, &'a mut S) -> (Option<&'a mut II>, &'a mut S),
+    IO: 'a,
+    II: 'a,
+    S: 'a,
+{
+    int_mutators_no_crossover()
+        .merge(mapped_int_mutators_crossover(input_from_corpus_mapper))
+        .map(ToStateAwareMappingMutator::new(current_input_mapper))
+}
+
 /// Functionality required for Numeric Mutators (see [`int_mutators`])
 pub trait Numeric {
     /// Flip all bits of the number.