Add MutatedTransform to the input type in TMinMutationalStage (#1251) (#1971)

* Support `MutatedTransform` in `TMinMutationalStage`.

* Run `MutatedTransformPost` for the replaced testcase.

* Add clone trait bound for `MutatedTransformPost`.

* Return an error instead of using unwrap.

Author: am009

libafl/src/stages/tmin.rs

@@ -1,7 +1,7 @@
 //! The [`TMinMutationalStage`] is a stage which will attempt to minimize corpus entries.
 
 use alloc::string::{String, ToString};
-use core::{fmt::Debug, hash::Hash, marker::PhantomData};
+use core::{borrow::BorrowMut, fmt::Debug, hash::Hash, marker::PhantomData};
 
 use ahash::RandomState;
 use libafl_bolts::{HasLen, Named};
@@ -16,7 +16,10 @@ use crate::{
     mutators::{MutationResult, Mutator},
     observers::{MapObserver, ObserversTuple},
     schedulers::{RemovableScheduler, Scheduler},
-    stages::{ExecutionCountRestartHelper, Stage},
+    stages::{
+        mutational::{MutatedTransform, MutatedTransformPost},
+        ExecutionCountRestartHelper, Stage,
+    },
     start_timer,
     state::{
         HasCorpus, HasCurrentTestcase, HasExecutions, HasMaxSize, HasMetadata, HasSolutions, State,
@@ -30,7 +33,7 @@ use crate::{monitors::PerfFeature, state::HasClientPerfMonitor};
 /// Mutational stage which minimizes corpus entries.
 ///
 /// You must provide at least one mutator that actually reduces size.
-pub trait TMinMutationalStage<CS, E, EM, F1, F2, M, OT, Z>:
+pub trait TMinMutationalStage<CS, E, EM, F1, F2, I, IP, M, OT, Z>:
     Stage<E, EM, Z> + FeedbackFactory<F2, CS::State, OT>
 where
     Self::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize,
@@ -40,12 +43,14 @@ where
     EM: EventFirer<State = Self::State>,
     F1: Feedback<Self::State>,
     F2: Feedback<Self::State>,
-    M: Mutator<Self::Input, Self::State>,
+    M: Mutator<I, Self::State>,
     OT: ObserversTuple<CS::State>,
     Z: ExecutionProcessor<OT, State = Self::State>
         + ExecutesInput<E, EM>
         + HasFeedback<Feedback = F1>
         + HasScheduler<Scheduler = CS>,
+    IP: MutatedTransformPost<Self::State> + Clone,
+    I: MutatedTransform<Self::Input, Self::State, Post = IP> + Clone,
 {
     /// The mutator registered for this stage
     fn mutator(&self) -> &M;
@@ -77,7 +82,10 @@ where
             - usize::try_from(self.execs_since_progress_start(state)?).unwrap();
 
         start_timer!(state);
+        let transformed = I::try_transform_from(state.current_testcase_mut()?.borrow_mut(), state)?;
         let mut base = state.current_input_cloned()?;
+        // potential post operation if base is replaced by a shorter input
+        let mut base_post = None;
         let base_hash = RandomState::with_seeds(0, 0, 0, 0).hash_one(&base);
         mark_feature_time!(state, PerfFeature::GetInputFromCorpus);
 
@@ -93,20 +101,21 @@ where
             }
 
             let mut next_i = i + 1;
-            let mut input = base.clone();
+            let mut input_transformed = transformed.clone();
 
-            let before_len = input.len();
+            let before_len = base.len();
 
             state.set_max_size(before_len);
 
             start_timer!(state);
-            let mutated = self.mutator_mut().mutate(state, &mut input)?;
+            let mutated = self.mutator_mut().mutate(state, &mut input_transformed)?;
             mark_feature_time!(state, PerfFeature::Mutate);
 
             if mutated == MutationResult::Skipped {
                 continue;
             }
 
+            let (input, post) = input_transformed.try_transform_into(state)?;
             let corpus_idx = if input.len() < before_len {
                 // run the input
                 let exit_kind = fuzzer.execute_input(state, executor, manager, &input)?;
@@ -134,6 +143,7 @@ where
                     if feedback.is_interesting(state, manager, &input, observers, &exit_kind)? {
                         // we found a reduced corpus entry! use the smaller base
                         base = input;
+                        base_post = Some(post.clone());
 
                         // do more runs! maybe we can minify further
                         next_i = 0;
@@ -149,6 +159,7 @@ where
 
             start_timer!(state);
             self.mutator_mut().post_exec(state, corpus_idx)?;
+            post.post_exec(state, corpus_idx)?;
             mark_feature_time!(state, PerfFeature::MutatePostExec);
 
             i = next_i;
@@ -171,6 +182,11 @@ where
             fuzzer
                 .scheduler_mut()
                 .on_replace(state, base_corpus_idx, &prev)?;
+            // perform the post operation for the new testcase, e.g. to update metadata.
+            // base_post should be updated along with the base (and is no longer None)
+            base_post
+                .ok_or_else(|| Error::empty_optional("Failed to get the MutatedTransformPost"))?
+                .post_exec(state, Some(base_corpus_idx))?;
         }
 
         state.set_max_size(orig_max_size);
@@ -184,7 +200,7 @@ where
 
 /// The default corpus entry minimising mutational stage
 #[derive(Clone, Debug)]
-pub struct StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z> {
+pub struct StdTMinMutationalStage<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z> {
     /// The mutator(s) this stage uses
     mutator: M,
     /// The factory
@@ -194,22 +210,24 @@ pub struct StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z> {
     /// The progress helper for this stage, keeping track of resumes after timeouts/crashes
     restart_helper: ExecutionCountRestartHelper,
     #[allow(clippy::type_complexity)]
-    phantom: PhantomData<(CS, E, EM, F1, F2, OT, Z)>,
+    phantom: PhantomData<(CS, E, EM, F1, F2, I, IP, OT, Z)>,
 }
 
-impl<CS, E, EM, F1, F2, FF, M, OT, Z> UsesState
-    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
+impl<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z> UsesState
+    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z>
 where
     CS: Scheduler,
-    M: Mutator<CS::Input, CS::State>,
+    M: Mutator<I, CS::State>,
     Z: ExecutionProcessor<OT, State = CS::State>,
     CS::State: HasCorpus,
+    IP: MutatedTransformPost<CS::State> + Clone,
+    I: MutatedTransform<CS::Input, CS::State, Post = IP> + Clone,
 {
     type State = CS::State;
 }
 
-impl<CS, E, EM, F1, F2, FF, M, OT, Z> Stage<E, EM, Z>
-    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
+impl<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z> Stage<E, EM, Z>
+    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z>
 where
     CS: Scheduler + RemovableScheduler,
     CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasCorpus + HasMetadata,
@@ -219,12 +237,14 @@ where
     F1: Feedback<CS::State>,
     F2: Feedback<CS::State>,
     FF: FeedbackFactory<F2, CS::State, OT>,
-    M: Mutator<CS::Input, CS::State>,
+    M: Mutator<I, CS::State>,
     OT: ObserversTuple<CS::State>,
     Z: ExecutionProcessor<OT, State = CS::State>
         + ExecutesInput<E, EM>
         + HasFeedback<Feedback = F1>
         + HasScheduler<Scheduler = CS>,
+    IP: MutatedTransformPost<CS::State> + Clone,
+    I: MutatedTransform<CS::Input, CS::State, Post = IP> + Clone,
 {
     fn perform(
         &mut self,
@@ -250,8 +270,8 @@ where
     }
 }
 
-impl<CS, E, EM, F1, F2, FF, M, OT, Z> FeedbackFactory<F2, Z::State, OT>
-    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
+impl<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z> FeedbackFactory<F2, Z::State, OT>
+    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z>
 where
     F2: Feedback<Z::State>,
     FF: FeedbackFactory<F2, Z::State, OT>,
@@ -262,8 +282,8 @@ where
     }
 }
 
-impl<CS, E, EM, F1, F2, FF, M, OT, Z> TMinMutationalStage<CS, E, EM, F1, F2, M, OT, Z>
-    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
+impl<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z> TMinMutationalStage<CS, E, EM, F1, F2, I, IP, M, OT, Z>
+    for StdTMinMutationalStage<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z>
 where
     CS: Scheduler + RemovableScheduler,
     E: HasObservers<Observers = OT, State = CS::State> + Executor<EM, Z>,
@@ -272,13 +292,15 @@ where
     F2: Feedback<CS::State>,
     FF: FeedbackFactory<F2, CS::State, OT>,
     <CS::State as UsesInput>::Input: HasLen + Hash,
-    M: Mutator<CS::Input, CS::State>,
+    M: Mutator<I, CS::State>,
     OT: ObserversTuple<CS::State>,
     CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasMetadata,
     Z: ExecutionProcessor<OT, State = CS::State>
         + ExecutesInput<E, EM>
         + HasFeedback<Feedback = F1>
         + HasScheduler<Scheduler = CS>,
+    IP: MutatedTransformPost<CS::State> + Clone,
+    I: MutatedTransform<CS::Input, CS::State, Post = IP> + Clone,
 {
     /// The mutator, added to this stage
     #[inline]
@@ -302,12 +324,15 @@ where
     }
 }
 
-impl<CS, E, EM, F1, F2, FF, M, OT, Z> StdTMinMutationalStage<CS, E, EM, F1, F2, FF, M, OT, Z>
+impl<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z>
+    StdTMinMutationalStage<CS, E, EM, F1, F2, FF, I, IP, M, OT, Z>
 where
     CS: Scheduler,
-    M: Mutator<CS::Input, CS::State>,
+    M: Mutator<I, CS::State>,
     Z: ExecutionProcessor<OT, State = CS::State>,
     CS::State: HasCorpus,
+    IP: MutatedTransformPost<CS::State> + Clone,
+    I: MutatedTransform<CS::Input, CS::State, Post = IP> + Clone,
 {
     /// Creates a new minimizing mutational stage that will minimize provided corpus entries
     pub fn new(mutator: M, factory: FF, runs: usize) -> Self {