From b7efe8eb7d95a1e7d39dd785ff1261c8be030af5 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Mon, 4 Mar 2024 22:48:38 +0100 Subject: [PATCH] Implement MutatorTuple for Vecs to allow Dynamic Mutator Choices (#1893) * Implement MutatorTuple for Vecs to allow Dynamic Mutator Choices * fix test * clippy * Move into_vec to extra trait * fix no_std * more nostd * no_std --- libafl/src/mutators/mod.rs | 170 +++++++++++++++++++++++++++- libafl/src/mutators/mopt_mutator.rs | 4 +- libafl/src/mutators/scheduled.rs | 6 +- libafl/src/mutators/tuneable.rs | 10 +- libafl_bolts/src/lib.rs | 8 ++ libafl_bolts/src/os/unix_signals.rs | 10 +- libafl_bolts/src/tuples.rs | 42 ++++++- 7 files changed, 227 insertions(+), 23 deletions(-) diff --git a/libafl/src/mutators/mod.rs b/libafl/src/mutators/mod.rs index 7232b0e800..cfe503b86f 100644 --- a/libafl/src/mutators/mod.rs +++ b/libafl/src/mutators/mod.rs @@ -33,11 +33,12 @@ pub use multi::*; #[cfg(feature = "nautilus")] pub mod nautilus; -use alloc::vec::Vec; +use alloc::{boxed::Box, vec::Vec}; -use libafl_bolts::{tuples::HasConstLen, Named}; +use libafl_bolts::{tuples::IntoVec, HasLen, Named}; #[cfg(feature = "nautilus")] pub use nautilus::*; +use tuple_list::NonEmptyTuple; use crate::{corpus::CorpusId, Error}; @@ -135,7 +136,7 @@ pub trait MultiMutator: Named { } /// A `Tuple` of `Mutators` that can execute multiple `Mutators` in a row. -pub trait MutatorsTuple: HasConstLen { +pub trait MutatorsTuple: HasLen { /// Runs the `mutate` function on all `Mutators` in this `Tuple`. fn mutate_all( &mut self, @@ -289,6 +290,142 @@ where } } +impl IntoVec>> for (Head, Tail) +where + Head: Mutator + 'static, + Tail: IntoVec>>, +{ + fn into_vec(self) -> Vec>> { + let (head, tail) = self.uncons(); + let mut ret = tail.into_vec(); + ret.insert(0, Box::new(head)); + ret + } +} + +impl MutatorsTuple for (Tail,) +where + Tail: MutatorsTuple, +{ + fn mutate_all( + &mut self, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + self.0.mutate_all(state, input, stage_idx) + } + + fn post_exec_all( + &mut self, + state: &mut S, + stage_idx: i32, + corpus_idx: Option, + ) -> Result<(), Error> { + self.0.post_exec_all(state, stage_idx, corpus_idx) + } + + fn get_and_mutate( + &mut self, + index: MutationId, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + self.0.get_and_mutate(index, state, input, stage_idx) + } + + fn get_and_post_exec( + &mut self, + index: usize, + state: &mut S, + stage_idx: i32, + corpus_idx: Option, + ) -> Result<(), Error> { + self.0 + .get_and_post_exec(index, state, stage_idx, corpus_idx) + } + + fn names(&self) -> Vec<&str> { + self.0.names() + } +} + +impl IntoVec>> for (Tail,) +where + Tail: IntoVec>>, +{ + fn into_vec(self) -> Vec>> { + self.0.into_vec() + } +} + +impl MutatorsTuple for Vec>> { + fn mutate_all( + &mut self, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + self.iter_mut() + .try_fold(MutationResult::Skipped, |ret, mutator| { + if mutator.mutate(state, input, stage_idx)? == MutationResult::Mutated { + Ok(MutationResult::Mutated) + } else { + Ok(ret) + } + }) + } + + fn post_exec_all( + &mut self, + state: &mut S, + stage_idx: i32, + corpus_idx: Option, + ) -> Result<(), Error> { + for mutator in self.iter_mut() { + mutator.post_exec(state, stage_idx, corpus_idx)?; + } + Ok(()) + } + + fn get_and_mutate( + &mut self, + index: MutationId, + state: &mut S, + input: &mut I, + stage_idx: i32, + ) -> Result { + let mutator = self + .get_mut(index.0) + .ok_or_else(|| Error::key_not_found("Mutator with id {index:?} not found."))?; + mutator.mutate(state, input, stage_idx) + } + + fn get_and_post_exec( + &mut self, + index: usize, + state: &mut S, + stage_idx: i32, + corpus_idx: Option, + ) -> Result<(), Error> { + let mutator = self + .get_mut(index) + .ok_or_else(|| Error::key_not_found("Mutator with id {index:?} not found."))?; + mutator.post_exec(state, stage_idx, corpus_idx) + } + + fn names(&self) -> Vec<&str> { + self.iter().map(|x| x.name()).collect() + } +} + +impl IntoVec>> for Vec>> { + fn into_vec(self) -> Vec>> { + self + } +} + /// `Mutator` Python bindings #[cfg(feature = "python")] #[allow(missing_docs)] @@ -463,3 +600,30 @@ pub mod pybind { Ok(()) } } + +#[cfg(test)] +mod tests { + use alloc::{boxed::Box, vec::Vec}; + + use libafl_bolts::{rands::StdRand, tuples::IntoVec}; + + use crate::{ + corpus::InMemoryCorpus, + inputs::BytesInput, + mutators::{havoc_mutations, Mutator, MutatorsTuple}, + state::StdState, + }; + + type TestStdStateType = + StdState, StdRand, InMemoryCorpus>; + + #[test] + fn test_tuple_into_vec() { + let mutators = havoc_mutations::(); + let names_before = MutatorsTuple::::names(&mutators); + + let mutators = havoc_mutations::(); + let mutators_vec: Vec>> = mutators.into_vec(); + assert_eq!(names_before, mutators_vec.names()); + } +} diff --git a/libafl/src/mutators/mopt_mutator.rs b/libafl/src/mutators/mopt_mutator.rs index 274699d6c2..9165ec65e6 100644 --- a/libafl/src/mutators/mopt_mutator.rs +++ b/libafl/src/mutators/mopt_mutator.rs @@ -392,7 +392,7 @@ where write!( f, "StdMOptMutator with {} mutations for Input type {}", - MT::LEN, + self.mutations.len(), core::any::type_name::() ) } @@ -547,7 +547,7 @@ where ) -> Result { if !state.has_metadata::() { let rand_seed = state.rand_mut().next(); - state.add_metadata::(MOpt::new(MT::LEN, swarm_num, rand_seed)?); + state.add_metadata::(MOpt::new(mutations.len(), swarm_num, rand_seed)?); } Ok(Self { name: format!("StdMOptMutator[{}]", mutations.names().join(",")), diff --git a/libafl/src/mutators/scheduled.rs b/libafl/src/mutators/scheduled.rs index 4bbff8dd19..75e780d690 100644 --- a/libafl/src/mutators/scheduled.rs +++ b/libafl/src/mutators/scheduled.rs @@ -135,7 +135,7 @@ where write!( f, "StdScheduledMutator with {} mutations for Input type {}", - MT::LEN, + self.mutations.len(), core::any::type_name::() ) } @@ -197,8 +197,8 @@ where /// Get the next mutation to apply fn schedule(&self, state: &mut S, _: &I) -> MutationId { - debug_assert!(MT::LEN != 0); - state.rand_mut().below(MT::LEN as u64).into() + debug_assert!(self.mutations.len() != 0); + state.rand_mut().below(self.mutations.len() as u64).into() } } diff --git a/libafl/src/mutators/tuneable.rs b/libafl/src/mutators/tuneable.rs index f417c9a314..0b72c69f34 100644 --- a/libafl/src/mutators/tuneable.rs +++ b/libafl/src/mutators/tuneable.rs @@ -100,7 +100,7 @@ where write!( f, "TuneableScheduledMutator with {} mutations for Input type {}", - MT::LEN, + self.mutations.len(), core::any::type_name::() ) } @@ -186,7 +186,7 @@ where /// Get the next mutation to apply fn schedule(&self, state: &mut S, _: &I) -> MutationId { - debug_assert!(MT::LEN != 0); + debug_assert!(self.mutations.len() != 0); // Assumption: we can not reach this code path without previously adding this metadatum. let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap(); @@ -199,7 +199,7 @@ where metadata.next_id = 0.into(); } debug_assert!( - MT::LEN > ret.0, + self.mutations.len() > ret.0, "TuneableScheduler: next vec may not contain id larger than number of mutations!" ); return ret; @@ -214,7 +214,7 @@ where let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap(); debug_assert_eq!( - MT::LEN, + self.mutations.len(), metadata.mutation_probabilities_cumulative.len(), "TuneableScheduler: mutation probabilities do not match with number of mutations" ); @@ -230,7 +230,7 @@ where } // fall back to random if no entries in either vec, the scheduling is not tuned. - state.rand_mut().below(MT::LEN as u64).into() + state.rand_mut().below(self.mutations.len() as u64).into() } } diff --git a/libafl_bolts/src/lib.rs b/libafl_bolts/src/lib.rs index 6daf6d9f4a..f19143f7eb 100644 --- a/libafl_bolts/src/lib.rs +++ b/libafl_bolts/src/lib.rs @@ -737,6 +737,14 @@ pub trait HasLen { } } +#[cfg(feature = "alloc")] +impl HasLen for Vec { + #[inline] + fn len(&self) -> usize { + Vec::::len(self) + } +} + /// Has a ref count pub trait HasRefCnt { /// The ref count diff --git a/libafl_bolts/src/os/unix_signals.rs b/libafl_bolts/src/os/unix_signals.rs index a105303937..81893fc13e 100644 --- a/libafl_bolts/src/os/unix_signals.rs +++ b/libafl_bolts/src/os/unix_signals.rs @@ -174,11 +174,11 @@ pub struct arm_neon_state64 { #[repr(C)] #[allow(clippy::pub_underscore_fields)] pub struct mcontext64 { - /// _STRUCT_ARM_EXCEPTION_STATE64 + /// `_STRUCT_ARM_EXCEPTION_STATE64` pub __es: arm_exception_state64, - /// _STRUCT_ARM_THREAD_STATE64 + /// `_STRUCT_ARM_THREAD_STATE64` pub __ss: arm_thread_state64, - /// _STRUCT_ARM_NEON_STATE64 + /// `_STRUCT_ARM_NEON_STATE64` pub __ns: arm_neon_state64, } @@ -189,7 +189,7 @@ pub struct mcontext64 { /// __darwin_size_t ss_size; /* signal stack length */ /// int ss_flags; /* SA_DISABLE and/or SA_ONSTACK */ /// }; -/// ```` +/// ``` #[cfg(all(target_vendor = "apple", target_arch = "aarch64"))] #[derive(Debug)] #[allow(non_camel_case_types)] @@ -199,7 +199,7 @@ pub struct sigaltstack { pub ss_sp: *mut c_void, /// signal stack length pub ss_size: libc::size_t, - /// SA_DISABLE and/or SA_ONSTACK + /// `SA_DISABLE` and/or `SA_ONSTACK` pub ss_flags: c_int, } diff --git a/libafl_bolts/src/tuples.rs b/libafl_bolts/src/tuples.rs index c1c5c1061d..65ae01c4b8 100644 --- a/libafl_bolts/src/tuples.rs +++ b/libafl_bolts/src/tuples.rs @@ -1,5 +1,7 @@ //! Compiletime lists/tuples used throughout the `LibAFL` universe +#[cfg(feature = "alloc")] +use alloc::vec::Vec; #[rustversion::not(nightly)] use core::any::type_name; use core::{ @@ -88,6 +90,22 @@ where } } +/// Create a [`Vec`] from a tuple list or similar +/// (We need this trait since we cannot implement `Into` for foreign types) +#[cfg(feature = "alloc")] +pub trait IntoVec { + /// Convert this into a [`Vec`]. + fn into_vec(self) -> Vec; +} + +#[cfg(feature = "alloc")] +impl IntoVec for () { + #[inline] + fn into_vec(self) -> Vec { + Vec::new() + } +} + /// Gets the length of the element pub trait HasConstLen { /// The length as constant `usize` @@ -105,16 +123,30 @@ where const LEN: usize = 1 + Tail::LEN; } -impl HasLen for C +impl HasLen for (Head, Tail) where - C: HasConstLen, + Tail: HasLen, { + #[inline] fn len(&self) -> usize { - Self::LEN + self.1.len() + 1 } +} - fn is_empty(&self) -> bool { - Self::LEN != 0 +impl HasLen for (Tail,) +where + Tail: HasLen, +{ + #[inline] + fn len(&self) -> usize { + self.0.len() + } +} + +impl HasLen for () { + #[inline] + fn len(&self) -> usize { + 0 } }