Tuneable stage with per-seed timeout (#1209)

* Tunable mutations with timeouts

* fix

* fmt

* Introspection fix

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<u64>,
+    fuzz_time: Option<Duration>,
 }
 
 impl_serdeany!(TuneableMutationalStageMetadata);
@@ -29,7 +33,7 @@ pub fn set_iters<S: HasMetadata>(state: &mut S, iters: u64) -> Result<(), Error>
     let metadata = state
         .metadata_map_mut()
         .get_mut::<TuneableMutationalStageMetadata>()
-        .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<S: HasMetadata>(state: &S) -> Result<Option<u64>, Error> {
     state
         .metadata_map()
         .get::<TuneableMutationalStageMetadata>()
-        .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<S: HasMetadata>(state: &mut S, fuzz_time: Duration) -> Result<(), Error> {
+    let metadata = state
+        .metadata_map_mut()
+        .get_mut::<TuneableMutationalStageMetadata>()
+        .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<S: HasMetadata>(state: &mut S) -> Result<Option<Duration>, Error> {
+    state
+        .metadata_map()
+        .get::<TuneableMutationalStageMetadata>()
+        .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<S: HasMetadata>(state: &mut S) -> Result<(), Error> {
     state
         .metadata_map_mut()
         .get_mut::<TuneableMutationalStageMetadata>()
-        .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<Z::Input, Z::State> + 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<S: HasMetadata>(state: &S) -> Result<Option<u64>, Error> {
         get_iters(state)
     }
+
+    /// Set the time to mutate a single input in this mutational stage
+    pub fn set_seed_fuzz_time<S: HasMetadata>(
+        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<S: HasMetadata>(state: &mut S) -> Result<Option<Duration>, Error> {
+        get_seed_fuzz_time(state)
+    }
 }
 
 impl<E, EM, I, M, Z> TuneableMutationalStage<E, EM, I, M, Z>