use core::{any::type_name, marker::PhantomData};
use libafl_bolts::{rands::Rand, Named};
use crate::{
corpus::{Corpus, CorpusId, Testcase},
fuzzer::Evaluator,
inputs::Input,
mark_feature_time,
mutators::{MultiMutator, MutationResult, Mutator},
stages::{ExecutionCountRestartHelper, RetryRestartHelper, Stage},
start_timer,
state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, UsesState},
Error, HasMetadata, HasNamedMetadata,
};
#[cfg(feature = "introspection")]
use crate::{monitors::PerfFeature, state::HasClientPerfMonitor};
#[allow(unused_variables)]
pub trait MutatedTransformPost<S>: Sized {
#[inline]
fn post_exec(self, state: &mut S, new_corpus_idx: Option<CorpusId>) -> Result<(), Error> {
Ok(())
}
}
impl<S> MutatedTransformPost<S> for () {}
pub trait MutatedTransform<I, S>: Sized
where
I: Input,
{
type Post: MutatedTransformPost<S>;
fn try_transform_from(base: &mut Testcase<I>, state: &S) -> Result<Self, Error>;
fn try_transform_into(self, state: &S) -> Result<(I, Self::Post), Error>;
}
impl<I, S> MutatedTransform<I, S> for I
where
I: Input + Clone,
S: HasCorpus<Input = I>,
{
type Post = ();
#[inline]
fn try_transform_from(base: &mut Testcase<I>, state: &S) -> Result<Self, Error> {
state.corpus().load_input_into(base)?;
Ok(base.input().as_ref().unwrap().clone())
}
#[inline]
fn try_transform_into(self, _state: &S) -> Result<(I, Self::Post), Error> {
Ok((self, ()))
}
}
pub trait MutationalStage<E, EM, I, M, Z>: Stage<E, EM, Z>
where
E: UsesState<State = Self::State>,
M: Mutator<I, Self::State>,
EM: UsesState<State = Self::State>,
Z: Evaluator<E, EM, State = Self::State>,
Self::State: HasCorpus,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
fn mutator(&self) -> &M;
fn mutator_mut(&mut self) -> &mut M;
fn iterations(&self, state: &mut Z::State) -> Result<u64, Error>;
fn execs_since_progress_start(&mut self, state: &mut Z::State) -> Result<u64, Error>;
#[allow(clippy::cast_possible_wrap)] fn perform_mutational(
&mut self,
fuzzer: &mut Z,
executor: &mut E,
state: &mut Z::State,
manager: &mut EM,
) -> Result<(), Error> {
start_timer!(state);
let num = self.iterations(state)?;
let mut testcase = state.current_testcase_mut()?;
let Ok(input) = I::try_transform_from(&mut testcase, state) else {
return Ok(());
};
drop(testcase);
mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
for _ in 0..num {
let mut input = input.clone();
start_timer!(state);
let mutated = self.mutator_mut().mutate(state, &mut input)?;
mark_feature_time!(state, PerfFeature::Mutate);
if mutated == MutationResult::Skipped {
continue;
}
let (untransformed, post) = input.try_transform_into(state)?;
let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, untransformed)?;
start_timer!(state);
self.mutator_mut().post_exec(state, corpus_idx)?;
post.post_exec(state, corpus_idx)?;
mark_feature_time!(state, PerfFeature::MutatePostExec);
}
Ok(())
}
}
pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: u64 = 128;
#[derive(Clone, Debug)]
pub struct StdMutationalStage<E, EM, I, M, Z> {
mutator: M,
max_iterations: u64,
restart_helper: ExecutionCountRestartHelper,
#[allow(clippy::type_complexity)]
phantom: PhantomData<(E, EM, I, Z)>,
}
impl<E, EM, I, M, Z> MutationalStage<E, EM, I, M, Z> for StdMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: Mutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand + HasExecutions + HasMetadata,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
#[inline]
fn mutator(&self) -> &M {
&self.mutator
}
#[inline]
fn mutator_mut(&mut self) -> &mut M {
&mut self.mutator
}
fn iterations(&self, state: &mut Z::State) -> Result<u64, Error> {
Ok(1 + state.rand_mut().below(self.max_iterations))
}
fn execs_since_progress_start(&mut self, state: &mut <Z>::State) -> Result<u64, Error> {
self.restart_helper.execs_since_progress_start(state)
}
}
impl<E, EM, I, M, Z> UsesState for StdMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: Mutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
type State = Z::State;
}
impl<E, EM, I, M, Z> Stage<E, EM, Z> for StdMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: Mutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand + HasMetadata + HasExecutions,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
#[inline]
#[allow(clippy::let_and_return)]
fn perform(
&mut self,
fuzzer: &mut Z,
executor: &mut E,
state: &mut Z::State,
manager: &mut EM,
) -> Result<(), Error> {
let ret = self.perform_mutational(fuzzer, executor, state, manager);
#[cfg(feature = "introspection")]
state.introspection_monitor_mut().finish_stage();
ret
}
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
Ok(true)
}
fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
Ok(())
}
}
impl<E, EM, M, Z> StdMutationalStage<E, EM, Z::Input, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: Mutator<Z::Input, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
pub fn new(mutator: M) -> Self {
Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS)
}
pub fn with_max_iterations(mutator: M, max_iterations: u64) -> Self {
Self::transforming_with_max_iterations(mutator, max_iterations)
}
}
impl<E, EM, I, M, Z> StdMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: Mutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
pub fn transforming(mutator: M) -> Self {
Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS)
}
pub fn transforming_with_max_iterations(mutator: M, max_iterations: u64) -> Self {
Self {
mutator,
max_iterations,
restart_helper: ExecutionCountRestartHelper::default(),
phantom: PhantomData,
}
}
}
#[derive(Clone, Debug)]
pub struct MultiMutationalStage<E, EM, I, M, Z> {
mutator: M,
#[allow(clippy::type_complexity)]
phantom: PhantomData<(E, EM, I, Z)>,
}
impl<E, EM, I, M, Z> UsesState for MultiMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: MultiMutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
type State = Z::State;
}
impl<E, EM, I, M, Z> Named for MultiMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: MultiMutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
fn name(&self) -> &str {
type_name::<Self>()
}
}
impl<E, EM, I, M, Z> Stage<E, EM, Z> for MultiMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: MultiMutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand + HasNamedMetadata,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
#[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, 3)
}
#[inline]
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
RetryRestartHelper::clear_restart_progress(state, self)
}
#[inline]
#[allow(clippy::let_and_return)]
#[allow(clippy::cast_possible_wrap)]
fn perform(
&mut self,
fuzzer: &mut Z,
executor: &mut E,
state: &mut Z::State,
manager: &mut EM,
) -> Result<(), Error> {
let mut testcase = state.current_testcase_mut()?;
let Ok(input) = I::try_transform_from(&mut testcase, state) else {
return Ok(());
};
drop(testcase);
let generated = self.mutator.multi_mutate(state, &input, None)?;
for new_input in generated {
let (untransformed, post) = new_input.try_transform_into(state)?;
let (_, corpus_idx) = fuzzer.evaluate_input(state, executor, manager, untransformed)?;
self.mutator.multi_post_exec(state, corpus_idx)?;
post.post_exec(state, corpus_idx)?;
}
Ok(())
}
}
impl<E, EM, M, Z> MultiMutationalStage<E, EM, Z::Input, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: MultiMutator<Z::Input, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
pub fn new(mutator: M) -> Self {
Self::transforming(mutator)
}
}
impl<E, EM, I, M, Z> MultiMutationalStage<E, EM, I, M, Z>
where
E: UsesState<State = Z::State>,
EM: UsesState<State = Z::State>,
M: MultiMutator<I, Z::State>,
Z: Evaluator<E, EM>,
Z::State: HasCorpus + HasRand,
{
pub fn transforming(mutator: M) -> Self {
Self {
mutator,
phantom: PhantomData,
}
}
}