From c221108916ab2f6547bddaa5a9b9b23db7bdf56d Mon Sep 17 00:00:00 2001 From: am009 Date: Thu, 28 Mar 2024 01:59:45 +0800 Subject: [PATCH] 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. --- libafl/src/stages/tmin.rs | 69 ++++++++++++++++++++++++++------------- 1 file changed, 47 insertions(+), 22 deletions(-) diff --git a/libafl/src/stages/tmin.rs b/libafl/src/stages/tmin.rs index 93851498f6..a4235cdc78 100644 --- a/libafl/src/stages/tmin.rs +++ b/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: +pub trait TMinMutationalStage: Stage + FeedbackFactory where Self::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize, @@ -40,12 +43,14 @@ where EM: EventFirer, F1: Feedback, F2: Feedback, - M: Mutator, + M: Mutator, OT: ObserversTuple, Z: ExecutionProcessor + ExecutesInput + HasFeedback + HasScheduler, + IP: MutatedTransformPost + Clone, + I: MutatedTransform + 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 { +pub struct StdTMinMutationalStage { /// The mutator(s) this stage uses mutator: M, /// The factory @@ -194,22 +210,24 @@ pub struct StdTMinMutationalStage { /// 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 UsesState - for StdTMinMutationalStage +impl UsesState + for StdTMinMutationalStage where CS: Scheduler, - M: Mutator, + M: Mutator, Z: ExecutionProcessor, CS::State: HasCorpus, + IP: MutatedTransformPost + Clone, + I: MutatedTransform + Clone, { type State = CS::State; } -impl Stage - for StdTMinMutationalStage +impl Stage + for StdTMinMutationalStage where CS: Scheduler + RemovableScheduler, CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasCorpus + HasMetadata, @@ -219,12 +237,14 @@ where F1: Feedback, F2: Feedback, FF: FeedbackFactory, - M: Mutator, + M: Mutator, OT: ObserversTuple, Z: ExecutionProcessor + ExecutesInput + HasFeedback + HasScheduler, + IP: MutatedTransformPost + Clone, + I: MutatedTransform + Clone, { fn perform( &mut self, @@ -250,8 +270,8 @@ where } } -impl FeedbackFactory - for StdTMinMutationalStage +impl FeedbackFactory + for StdTMinMutationalStage where F2: Feedback, FF: FeedbackFactory, @@ -262,8 +282,8 @@ where } } -impl TMinMutationalStage - for StdTMinMutationalStage +impl TMinMutationalStage + for StdTMinMutationalStage where CS: Scheduler + RemovableScheduler, E: HasObservers + Executor, @@ -272,13 +292,15 @@ where F2: Feedback, FF: FeedbackFactory, ::Input: HasLen + Hash, - M: Mutator, + M: Mutator, OT: ObserversTuple, CS::State: HasCorpus + HasSolutions + HasExecutions + HasMaxSize + HasMetadata, Z: ExecutionProcessor + ExecutesInput + HasFeedback + HasScheduler, + IP: MutatedTransformPost + Clone, + I: MutatedTransform + Clone, { /// The mutator, added to this stage #[inline] @@ -302,12 +324,15 @@ where } } -impl StdTMinMutationalStage +impl + StdTMinMutationalStage where CS: Scheduler, - M: Mutator, + M: Mutator, Z: ExecutionProcessor, CS::State: HasCorpus, + IP: MutatedTransformPost + Clone, + I: MutatedTransform + Clone, { /// Creates a new minimizing mutational stage that will minimize provided corpus entries pub fn new(mutator: M, factory: FF, runs: usize) -> Self {