From cdd3d8ace0138d9b40c7d790d6e30529965aea08 Mon Sep 17 00:00:00 2001 From: Dominik Maier Date: Thu, 13 Apr 2023 15:30:47 +0200 Subject: [PATCH] Tuneable stage with per-seed timeout (#1209) * Tunable mutations with timeouts * fix * fmt * Introspection fix --- libafl/src/stages/tuneable.rs | 119 +++++++++++++++++++++++++++++++--- 1 file changed, 110 insertions(+), 9 deletions(-) diff --git a/libafl/src/stages/tuneable.rs b/libafl/src/stages/tuneable.rs index 01b9aa0cd6..b70534204e 100644 --- a/libafl/src/stages/tuneable.rs +++ b/libafl/src/stages/tuneable.rs @@ -1,18 +1,21 @@ //! A [`crate::stages::MutationalStage`] where the mutator iteration can be tuned at runtime -use core::marker::PhantomData; +use core::{marker::PhantomData, time::Duration}; use serde::{Deserialize, Serialize}; +#[cfg(feature = "introspection")] +use crate::monitors::PerfFeature; use crate::{ - bolts::rands::Rand, - corpus::CorpusId, - impl_serdeany, + bolts::{current_time, rands::Rand}, + corpus::{Corpus, CorpusId}, + impl_serdeany, mark_feature_time, mutators::Mutator, stages::{ - mutational::{MutatedTransform, DEFAULT_MUTATIONAL_MAX_ITERATIONS}, + mutational::{MutatedTransform, MutatedTransformPost, DEFAULT_MUTATIONAL_MAX_ITERATIONS}, MutationalStage, Stage, }, + start_timer, state::{HasClientPerfMonitor, HasCorpus, HasMetadata, HasRand, UsesState}, Error, Evaluator, }; @@ -20,6 +23,7 @@ use crate::{ #[derive(Default, Clone, Copy, Eq, PartialEq, Debug, Serialize, Deserialize)] struct TuneableMutationalStageMetadata { iters: Option, + fuzz_time: Option, } impl_serdeany!(TuneableMutationalStageMetadata); @@ -29,7 +33,7 @@ pub fn set_iters(state: &mut S, iters: u64) -> Result<(), Error> let metadata = state .metadata_map_mut() .get_mut::() - .ok_or_else(|| Error::illegal_state("TuneableMutationslStage not in use")); + .ok_or_else(|| Error::illegal_state("TuneableMutationalStage not in use")); metadata.map(|metadata| { metadata.iters = Some(iters); }) @@ -40,17 +44,40 @@ pub fn get_iters(state: &S) -> Result, Error> { state .metadata_map() .get::() - .ok_or_else(|| Error::illegal_state("TuneableMutationslStage not in use")) + .ok_or_else(|| Error::illegal_state("TuneableMutationalStage not in use")) .map(|metadata| metadata.iters) } +/// Set the time for a single seed to be used by this mutational stage +pub fn set_seed_fuzz_time(state: &mut S, fuzz_time: Duration) -> Result<(), Error> { + let metadata = state + .metadata_map_mut() + .get_mut::() + .ok_or_else(|| Error::illegal_state("TuneableMutationalStage not in use")); + metadata.map(|metadata| { + metadata.fuzz_time = Some(fuzz_time); + }) +} + +/// Get the time for a single seed to be used by this mutational stage +pub fn get_seed_fuzz_time(state: &mut S) -> Result, Error> { + state + .metadata_map() + .get::() + .ok_or_else(|| Error::illegal_state("TuneableMutationalStage not in use")) + .map(|metadata| metadata.fuzz_time) +} + /// Reset this to a normal, randomized, stage pub fn reset(state: &mut S) -> Result<(), Error> { state .metadata_map_mut() .get_mut::() - .ok_or_else(|| Error::illegal_state("TuneableMutationslStage not in use")) - .map(|metadata| metadata.iters = None) + .ok_or_else(|| Error::illegal_state("TuneableMutationalStage not in use")) + .map(|metadata| { + metadata.iters = None; + metadata.fuzz_time = None; + }) } /// A [`crate::stages::MutationalStage`] where the mutator iteration can be tuned at runtime @@ -69,6 +96,67 @@ where Z::State: HasClientPerfMonitor + HasCorpus + HasRand + HasMetadata, I: MutatedTransform + Clone, { + /// Runs this (mutational) stage for the given `testcase` + /// Exactly the same functionality as [`MutationalStage::perform_mutational`], but with added timeout support. + #[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely... + fn perform_mutational( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut Z::State, + manager: &mut EM, + corpus_idx: CorpusId, + ) -> Result<(), Error> { + let metadata: &TuneableMutationalStageMetadata = state.metadata()?; + + let fuzz_time = metadata.fuzz_time; + let iters = metadata.iters; + + let (start_time, iters) = if fuzz_time.is_some() { + (Some(current_time()), iters) + } else { + (None, Some(self.iterations(state, corpus_idx)?)) + }; + + start_timer!(state); + let mut testcase = state.corpus().get(corpus_idx)?.borrow_mut(); + let Ok(input) = I::try_transform_from(&mut testcase, state, corpus_idx) else { return Ok(()); }; + drop(testcase); + mark_feature_time!(state, PerfFeature::GetInputFromCorpus); + + let mut i = 0_usize; + loop { + if let Some(start_time) = start_time { + if current_time() - start_time >= fuzz_time.unwrap() { + break; + } + } + if let Some(iters) = iters { + if i >= iters as usize { + break; + } + } else { + i += 1; + } + + let mut input = input.clone(); + + start_timer!(state); + self.mutator_mut().mutate(state, &mut input, i as i32)?; + mark_feature_time!(state, PerfFeature::Mutate); + + // Time is measured directly the `evaluate_input` function + 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, i as i32, corpus_idx)?; + post.post_exec(state, i as i32, corpus_idx)?; + mark_feature_time!(state, PerfFeature::MutatePostExec); + } + Ok(()) + } + /// The mutator, added to this stage #[inline] fn mutator(&self) -> &M { @@ -158,6 +246,19 @@ impl TuneableMutationalStage<(), (), (), (), ()> { pub fn iters(state: &S) -> Result, Error> { get_iters(state) } + + /// Set the time to mutate a single input in this mutational stage + pub fn set_seed_fuzz_time( + state: &mut S, + fuzz_time: Duration, + ) -> Result<(), Error> { + set_seed_fuzz_time(state, fuzz_time) + } + + /// Set the time to mutate a single input in this mutational stage + pub fn seed_fuzz_time(state: &mut S) -> Result, Error> { + get_seed_fuzz_time(state) + } } impl TuneableMutationalStage