diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index d93b381065..5cb590ed2a 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -158,6 +158,7 @@ pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: u64 = 128; #[derive(Clone, Debug)] pub struct StdMutationalStage { mutator: M, + max_iterations: u64, #[allow(clippy::type_complexity)] phantom: PhantomData<(E, EM, I, Z)>, } @@ -185,7 +186,7 @@ where /// Gets the number of iterations as a random number fn iterations(&self, state: &mut Z::State, _corpus_idx: CorpusId) -> Result { - Ok(1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS)) + Ok(1 + state.rand_mut().below(self.max_iterations)) } } @@ -238,7 +239,12 @@ where { /// Creates a new default mutational stage pub fn new(mutator: M) -> Self { - Self::transforming(mutator) + Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS) + } + + /// Creates a new mutational stage with the given max iterations + pub fn with_max_iterations(mutator: M, max_iterations: u64) -> Self { + Self::transforming_with_max_iterations(mutator, max_iterations) } } @@ -250,10 +256,16 @@ where Z: Evaluator, Z::State: HasClientPerfMonitor + HasCorpus + HasRand, { - /// Creates a new transforming mutational stage + /// Creates a new transforming mutational stage with the default max iterations pub fn transforming(mutator: M) -> Self { + Self::transforming_with_max_iterations(mutator, DEFAULT_MUTATIONAL_MAX_ITERATIONS) + } + + /// Creates a new transforming mutational stage with the given max iterations + pub fn transforming_with_max_iterations(mutator: M, max_iterations: u64) -> Self { Self { mutator, + max_iterations, phantom: PhantomData, } }