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
This commit is contained in:
Dominik Maier 2024-03-04 22:48:38 +01:00 committed by GitHub
parent f3c37db2b7
commit b7efe8eb7d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 227 additions and 23 deletions

View File

@ -33,11 +33,12 @@ pub use multi::*;
#[cfg(feature = "nautilus")] #[cfg(feature = "nautilus")]
pub mod 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")] #[cfg(feature = "nautilus")]
pub use nautilus::*; pub use nautilus::*;
use tuple_list::NonEmptyTuple;
use crate::{corpus::CorpusId, Error}; use crate::{corpus::CorpusId, Error};
@ -135,7 +136,7 @@ pub trait MultiMutator<I, S>: Named {
} }
/// A `Tuple` of `Mutators` that can execute multiple `Mutators` in a row. /// A `Tuple` of `Mutators` that can execute multiple `Mutators` in a row.
pub trait MutatorsTuple<I, S>: HasConstLen { pub trait MutatorsTuple<I, S>: HasLen {
/// Runs the `mutate` function on all `Mutators` in this `Tuple`. /// Runs the `mutate` function on all `Mutators` in this `Tuple`.
fn mutate_all( fn mutate_all(
&mut self, &mut self,
@ -289,6 +290,142 @@ where
} }
} }
impl<Head, Tail, I, S> IntoVec<Box<dyn Mutator<I, S>>> for (Head, Tail)
where
Head: Mutator<I, S> + 'static,
Tail: IntoVec<Box<dyn Mutator<I, S>>>,
{
fn into_vec(self) -> Vec<Box<dyn Mutator<I, S>>> {
let (head, tail) = self.uncons();
let mut ret = tail.into_vec();
ret.insert(0, Box::new(head));
ret
}
}
impl<Tail, I, S> MutatorsTuple<I, S> for (Tail,)
where
Tail: MutatorsTuple<I, S>,
{
fn mutate_all(
&mut self,
state: &mut S,
input: &mut I,
stage_idx: i32,
) -> Result<MutationResult, Error> {
self.0.mutate_all(state, input, stage_idx)
}
fn post_exec_all(
&mut self,
state: &mut S,
stage_idx: i32,
corpus_idx: Option<CorpusId>,
) -> 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<MutationResult, Error> {
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<CorpusId>,
) -> Result<(), Error> {
self.0
.get_and_post_exec(index, state, stage_idx, corpus_idx)
}
fn names(&self) -> Vec<&str> {
self.0.names()
}
}
impl<Tail, I, S> IntoVec<Box<dyn Mutator<I, S>>> for (Tail,)
where
Tail: IntoVec<Box<dyn Mutator<I, S>>>,
{
fn into_vec(self) -> Vec<Box<dyn Mutator<I, S>>> {
self.0.into_vec()
}
}
impl<I, S> MutatorsTuple<I, S> for Vec<Box<dyn Mutator<I, S>>> {
fn mutate_all(
&mut self,
state: &mut S,
input: &mut I,
stage_idx: i32,
) -> Result<MutationResult, Error> {
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<CorpusId>,
) -> 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<MutationResult, Error> {
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<CorpusId>,
) -> 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<I, S> IntoVec<Box<dyn Mutator<I, S>>> for Vec<Box<dyn Mutator<I, S>>> {
fn into_vec(self) -> Vec<Box<dyn Mutator<I, S>>> {
self
}
}
/// `Mutator` Python bindings /// `Mutator` Python bindings
#[cfg(feature = "python")] #[cfg(feature = "python")]
#[allow(missing_docs)] #[allow(missing_docs)]
@ -463,3 +600,30 @@ pub mod pybind {
Ok(()) 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<BytesInput, InMemoryCorpus<BytesInput>, StdRand, InMemoryCorpus<BytesInput>>;
#[test]
fn test_tuple_into_vec() {
let mutators = havoc_mutations::<BytesInput>();
let names_before = MutatorsTuple::<BytesInput, TestStdStateType>::names(&mutators);
let mutators = havoc_mutations::<BytesInput>();
let mutators_vec: Vec<Box<dyn Mutator<BytesInput, TestStdStateType>>> = mutators.into_vec();
assert_eq!(names_before, mutators_vec.names());
}
}

View File

@ -392,7 +392,7 @@ where
write!( write!(
f, f,
"StdMOptMutator with {} mutations for Input type {}", "StdMOptMutator with {} mutations for Input type {}",
MT::LEN, self.mutations.len(),
core::any::type_name::<I>() core::any::type_name::<I>()
) )
} }
@ -547,7 +547,7 @@ where
) -> Result<Self, Error> { ) -> Result<Self, Error> {
if !state.has_metadata::<MOpt>() { if !state.has_metadata::<MOpt>() {
let rand_seed = state.rand_mut().next(); let rand_seed = state.rand_mut().next();
state.add_metadata::<MOpt>(MOpt::new(MT::LEN, swarm_num, rand_seed)?); state.add_metadata::<MOpt>(MOpt::new(mutations.len(), swarm_num, rand_seed)?);
} }
Ok(Self { Ok(Self {
name: format!("StdMOptMutator[{}]", mutations.names().join(",")), name: format!("StdMOptMutator[{}]", mutations.names().join(",")),

View File

@ -135,7 +135,7 @@ where
write!( write!(
f, f,
"StdScheduledMutator with {} mutations for Input type {}", "StdScheduledMutator with {} mutations for Input type {}",
MT::LEN, self.mutations.len(),
core::any::type_name::<I>() core::any::type_name::<I>()
) )
} }
@ -197,8 +197,8 @@ where
/// Get the next mutation to apply /// Get the next mutation to apply
fn schedule(&self, state: &mut S, _: &I) -> MutationId { fn schedule(&self, state: &mut S, _: &I) -> MutationId {
debug_assert!(MT::LEN != 0); debug_assert!(self.mutations.len() != 0);
state.rand_mut().below(MT::LEN as u64).into() state.rand_mut().below(self.mutations.len() as u64).into()
} }
} }

View File

@ -100,7 +100,7 @@ where
write!( write!(
f, f,
"TuneableScheduledMutator with {} mutations for Input type {}", "TuneableScheduledMutator with {} mutations for Input type {}",
MT::LEN, self.mutations.len(),
core::any::type_name::<I>() core::any::type_name::<I>()
) )
} }
@ -186,7 +186,7 @@ where
/// Get the next mutation to apply /// Get the next mutation to apply
fn schedule(&self, state: &mut S, _: &I) -> MutationId { 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. // Assumption: we can not reach this code path without previously adding this metadatum.
let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap(); let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap();
@ -199,7 +199,7 @@ where
metadata.next_id = 0.into(); metadata.next_id = 0.into();
} }
debug_assert!( debug_assert!(
MT::LEN > ret.0, self.mutations.len() > ret.0,
"TuneableScheduler: next vec may not contain id larger than number of mutations!" "TuneableScheduler: next vec may not contain id larger than number of mutations!"
); );
return ret; return ret;
@ -214,7 +214,7 @@ where
let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap(); let metadata = TuneableScheduledMutatorMetadata::get_mut(state).unwrap();
debug_assert_eq!( debug_assert_eq!(
MT::LEN, self.mutations.len(),
metadata.mutation_probabilities_cumulative.len(), metadata.mutation_probabilities_cumulative.len(),
"TuneableScheduler: mutation probabilities do not match with number of mutations" "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. // 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()
} }
} }

View File

@ -737,6 +737,14 @@ pub trait HasLen {
} }
} }
#[cfg(feature = "alloc")]
impl<T> HasLen for Vec<T> {
#[inline]
fn len(&self) -> usize {
Vec::<T>::len(self)
}
}
/// Has a ref count /// Has a ref count
pub trait HasRefCnt { pub trait HasRefCnt {
/// The ref count /// The ref count

View File

@ -174,11 +174,11 @@ pub struct arm_neon_state64 {
#[repr(C)] #[repr(C)]
#[allow(clippy::pub_underscore_fields)] #[allow(clippy::pub_underscore_fields)]
pub struct mcontext64 { pub struct mcontext64 {
/// _STRUCT_ARM_EXCEPTION_STATE64 /// `_STRUCT_ARM_EXCEPTION_STATE64`
pub __es: arm_exception_state64, pub __es: arm_exception_state64,
/// _STRUCT_ARM_THREAD_STATE64 /// `_STRUCT_ARM_THREAD_STATE64`
pub __ss: arm_thread_state64, pub __ss: arm_thread_state64,
/// _STRUCT_ARM_NEON_STATE64 /// `_STRUCT_ARM_NEON_STATE64`
pub __ns: arm_neon_state64, pub __ns: arm_neon_state64,
} }
@ -189,7 +189,7 @@ pub struct mcontext64 {
/// __darwin_size_t ss_size; /* signal stack length */ /// __darwin_size_t ss_size; /* signal stack length */
/// int ss_flags; /* SA_DISABLE and/or SA_ONSTACK */ /// int ss_flags; /* SA_DISABLE and/or SA_ONSTACK */
/// }; /// };
/// ```` /// ```
#[cfg(all(target_vendor = "apple", target_arch = "aarch64"))] #[cfg(all(target_vendor = "apple", target_arch = "aarch64"))]
#[derive(Debug)] #[derive(Debug)]
#[allow(non_camel_case_types)] #[allow(non_camel_case_types)]
@ -199,7 +199,7 @@ pub struct sigaltstack {
pub ss_sp: *mut c_void, pub ss_sp: *mut c_void,
/// signal stack length /// signal stack length
pub ss_size: libc::size_t, pub ss_size: libc::size_t,
/// SA_DISABLE and/or SA_ONSTACK /// `SA_DISABLE` and/or `SA_ONSTACK`
pub ss_flags: c_int, pub ss_flags: c_int,
} }

View File

@ -1,5 +1,7 @@
//! Compiletime lists/tuples used throughout the `LibAFL` universe //! Compiletime lists/tuples used throughout the `LibAFL` universe
#[cfg(feature = "alloc")]
use alloc::vec::Vec;
#[rustversion::not(nightly)] #[rustversion::not(nightly)]
use core::any::type_name; use core::any::type_name;
use core::{ 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<T> {
/// Convert this into a [`Vec`].
fn into_vec(self) -> Vec<T>;
}
#[cfg(feature = "alloc")]
impl<T> IntoVec<T> for () {
#[inline]
fn into_vec(self) -> Vec<T> {
Vec::new()
}
}
/// Gets the length of the element /// Gets the length of the element
pub trait HasConstLen { pub trait HasConstLen {
/// The length as constant `usize` /// The length as constant `usize`
@ -105,16 +123,30 @@ where
const LEN: usize = 1 + Tail::LEN; const LEN: usize = 1 + Tail::LEN;
} }
impl<C> HasLen for C impl<Head, Tail> HasLen for (Head, Tail)
where where
C: HasConstLen, Tail: HasLen,
{ {
#[inline]
fn len(&self) -> usize { fn len(&self) -> usize {
Self::LEN self.1.len() + 1
} }
}
fn is_empty(&self) -> bool { impl<Tail> HasLen for (Tail,)
Self::LEN != 0 where
Tail: HasLen,
{
#[inline]
fn len(&self) -> usize {
self.0.len()
}
}
impl HasLen for () {
#[inline]
fn len(&self) -> usize {
0
} }
} }