diff --git a/fuzzers/forkserver/libafl-fuzz/src/stages/mutational_stage.rs b/fuzzers/forkserver/libafl-fuzz/src/stages/mutational_stage.rs index 7e91d65027..009dbe6a47 100644 --- a/fuzzers/forkserver/libafl-fuzz/src/stages/mutational_stage.rs +++ b/fuzzers/forkserver/libafl-fuzz/src/stages/mutational_stage.rs @@ -1,21 +1,21 @@ use std::{borrow::Cow, marker::PhantomData}; use libafl::{ - stages::{MutationalStage, Stage}, + stages::{MutationalStage, Restartable, Stage}, Error, }; use libafl_bolts::Named; #[derive(Debug)] -pub enum SupportedMutationalStages { +pub enum SupportedMutationalStages { StdMutational(SM, PhantomData

), PowerMutational(P, PhantomData), } -impl MutationalStage for SupportedMutationalStages +impl MutationalStage for SupportedMutationalStages where - SM: MutationalStage, P: MutationalStage, + SM: MutationalStage, { type Mutator = SM::Mutator; /// The mutator, added to this stage @@ -44,10 +44,10 @@ where } } -impl Named for SupportedMutationalStages +impl Named for SupportedMutationalStages where - SM: Named, P: Named, + SM: Named, { fn name(&self) -> &Cow<'static, str> { match self { @@ -57,10 +57,10 @@ where } } -impl Stage for SupportedMutationalStages +impl Stage for SupportedMutationalStages where - SM: Stage, P: Stage, + SM: Stage, { #[inline] fn perform( @@ -75,7 +75,13 @@ where Self::PowerMutational(p, _) => p.perform(fuzzer, executor, state, manager), } } +} +impl Restartable for SupportedMutationalStages +where + P: Restartable, + SM: Restartable, +{ fn should_restart(&mut self, state: &mut S) -> Result { match self { Self::StdMutational(m, _) => m.should_restart(state), diff --git a/libafl/src/stages/afl_stats.rs b/libafl/src/stages/afl_stats.rs index acec81e61d..2efe98e67b 100644 --- a/libafl/src/stages/afl_stats.rs +++ b/libafl/src/stages/afl_stats.rs @@ -30,7 +30,7 @@ use crate::{ mutators::Tokens, observers::MapObserver, schedulers::{minimizer::IsFavoredMetadata, HasQueueCycles}, - stages::{calibrate::UnstableEntriesMetadata, Stage}, + stages::{calibrate::UnstableEntriesMetadata, Restartable, Stage}, state::{HasCorpus, HasExecutions, HasImported, HasStartTime, Stoppable}, std::string::ToString, Error, HasMetadata, HasNamedMetadata, HasScheduler, @@ -432,7 +432,9 @@ where Ok(()) } +} +impl Restartable for AflStatsStage { fn should_restart(&mut self, _state: &mut S) -> Result { Ok(true) } diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs index cfd163fb4e..2203d8b556 100644 --- a/libafl/src/stages/calibrate.rs +++ b/libafl/src/stages/calibrate.rs @@ -22,7 +22,7 @@ use crate::{ monitors::stats::{AggregatorOps, UserStats, UserStatsValue}, observers::{MapObserver, ObserversTuple}, schedulers::powersched::SchedulerMetadata, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, state::{HasCorpus, HasCurrentTestcase, HasExecutions}, Error, HasMetadata, HasNamedMetadata, }; @@ -364,7 +364,12 @@ where Ok(()) } +} +impl Restartable for CalibrationStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { // Calibration stage disallow restarts // If a testcase that causes crash/timeout in the queue, we need to remove it from the queue immediately. diff --git a/libafl/src/stages/colorization.rs b/libafl/src/stages/colorization.rs index 8f3f57072c..4af91aaf5e 100644 --- a/libafl/src/stages/colorization.rs +++ b/libafl/src/stages/colorization.rs @@ -22,7 +22,7 @@ use crate::{ mutators::mutations::buffer_copy, nonzero, observers::ObserversTuple, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, state::{HasCorpus, HasCurrentTestcase, HasRand}, Error, HasMetadata, HasNamedMetadata, }; @@ -98,7 +98,12 @@ where Ok(()) } +} +impl Restartable for ColorizationStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { // This is a deterministic stage // Once it failed, then don't retry, diff --git a/libafl/src/stages/concolic.rs b/libafl/src/stages/concolic.rs index dd8659f46f..a1175c7173 100644 --- a/libafl/src/stages/concolic.rs +++ b/libafl/src/stages/concolic.rs @@ -18,7 +18,7 @@ use crate::{ corpus::HasCurrentCorpusId, executors::{Executor, HasObservers}, observers::{concolic::ConcolicObserver, ObserversTuple}, - stages::{RetryCountRestartHelper, Stage, TracingStage}, + stages::{Restartable, RetryCountRestartHelper, Stage, TracingStage}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, MaybeHasClientPerfMonitor}, Error, HasMetadata, HasNamedMetadata, }; @@ -76,7 +76,12 @@ where } Ok(()) } +} +impl Restartable for ConcolicTracingStage<'_, EM, I, TE, S, Z> +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { // This is a deterministic stage // Once it failed, then don't retry, @@ -418,7 +423,13 @@ where } Ok(()) } +} +#[cfg(feature = "concolic_mutation")] +impl Restartable for SimpleConcolicMutationalStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ #[inline] fn should_restart(&mut self, state: &mut S) -> Result { // This is a deterministic stage diff --git a/libafl/src/stages/dump.rs b/libafl/src/stages/dump.rs index 97e043c89b..62ff6163be 100644 --- a/libafl/src/stages/dump.rs +++ b/libafl/src/stages/dump.rs @@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize}; use crate::{ corpus::{Corpus, CorpusId, Testcase}, inputs::Input, - stages::Stage, + stages::{Restartable, Stage}, state::{HasCorpus, HasRand, HasSolutions}, Error, HasMetadata, }; @@ -60,7 +60,11 @@ where ) -> Result<(), Error> { self.dump_state_to_disk(state) } +} +impl Restartable + for DumpToDiskStage, &CorpusId) -> String, EM, I, S, Z> +{ #[inline] fn should_restart(&mut self, _state: &mut S) -> Result { // Not executing the target, so restart safety is not needed diff --git a/libafl/src/stages/generalization.rs b/libafl/src/stages/generalization.rs index 6dd1b11ccd..d1d6a289fe 100644 --- a/libafl/src/stages/generalization.rs +++ b/libafl/src/stages/generalization.rs @@ -23,7 +23,7 @@ use crate::{ mark_feature_time, observers::{CanTrack, MapObserver, ObserversTuple}, require_novelties_tracking, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, start_timer, state::{HasCorpus, HasExecutions, MaybeHasClientPerfMonitor}, Error, HasMetadata, HasNamedMetadata, @@ -62,6 +62,23 @@ impl Named for GeneralizationStage } } +impl Restartable for GeneralizationStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ + #[inline] + fn should_restart(&mut self, state: &mut S) -> Result { + // TODO: We need to be able to resume better if something crashes or times out + RetryCountRestartHelper::should_restart::(state, &self.name, 3) + } + + #[inline] + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { + // TODO: We need to be able to resume better if something crashes or times out + RetryCountRestartHelper::clear_progress::(state, &self.name) + } +} + impl Stage for GeneralizationStage where @@ -76,18 +93,6 @@ where + HasCurrentCorpusId + MaybeHasClientPerfMonitor, { - #[inline] - fn should_restart(&mut self, state: &mut S) -> Result { - // TODO: We need to be able to resume better if something crashes or times out - RetryCountRestartHelper::should_restart::(state, &self.name, 3) - } - - #[inline] - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { - // TODO: We need to be able to resume better if something crashes or times out - RetryCountRestartHelper::clear_progress::(state, &self.name) - } - #[inline] #[expect(clippy::too_many_lines)] fn perform( diff --git a/libafl/src/stages/generation.rs b/libafl/src/stages/generation.rs index faf56656f4..ac312f48f1 100644 --- a/libafl/src/stages/generation.rs +++ b/libafl/src/stages/generation.rs @@ -6,7 +6,12 @@ use core::marker::PhantomData; -use crate::{generators::Generator, stages::Stage, state::HasRand, Error, Evaluator}; +use crate::{ + generators::Generator, + stages::{Restartable, Stage}, + state::HasRand, + Error, Evaluator, +}; /// A [`Stage`] that generates a single input via a [`Generator`] and evaluates /// it using the fuzzer, possibly adding it to the corpus. @@ -40,7 +45,9 @@ where fuzzer.evaluate_filtered(state, executor, manager, &input)?; Ok(()) } +} +impl Restartable for GenStage { fn should_restart(&mut self, _state: &mut S) -> Result { // It's a random generation stage // so you can restart for whatever times you want diff --git a/libafl/src/stages/logics.rs b/libafl/src/stages/logics.rs index bad0196e0a..00adc546f1 100644 --- a/libafl/src/stages/logics.rs +++ b/libafl/src/stages/logics.rs @@ -3,7 +3,7 @@ use core::marker::PhantomData; use crate::{ - stages::{Stage, StageId, StagesTuple}, + stages::{Restartable, Stage, StageId, StagesTuple}, state::HasNestedStage, Error, }; @@ -59,7 +59,12 @@ where Ok(()) } +} +impl Restartable for WhileStage +where + S: HasNestedStage, +{ fn should_restart(&mut self, state: &mut S) -> Result { NestedStageRetryCountRestartHelper::should_restart(state, self) } @@ -112,7 +117,12 @@ where } Ok(()) } +} +impl Restartable for IfStage +where + S: HasNestedStage, +{ fn should_restart(&mut self, state: &mut S) -> Result { NestedStageRetryCountRestartHelper::should_restart(state, self) } @@ -188,7 +198,12 @@ where Ok(()) } +} +impl Restartable for IfElseStage +where + S: HasNestedStage, +{ fn should_restart(&mut self, state: &mut S) -> Result { NestedStageRetryCountRestartHelper::should_restart(state, self) } @@ -238,7 +253,12 @@ where Ok(()) } } +} +impl Restartable for OptionalStage +where + S: HasNestedStage, +{ fn should_restart(&mut self, state: &mut S) -> Result { NestedStageRetryCountRestartHelper::should_restart(state, self) } @@ -295,8 +315,8 @@ mod test { events::NopEventManager, executors::test::NopExecutor, stages::{ - ClosureStage, CorpusId, HasCurrentCorpusId, IfElseStage, IfStage, Stage, StagesTuple, - WhileStage, + ClosureStage, CorpusId, HasCurrentCorpusId, IfElseStage, IfStage, Restartable, Stage, + StagesTuple, WhileStage, }, state::{HasCurrentStageId, StdState}, HasMetadata, NopFuzzer, @@ -362,7 +382,12 @@ mod test { ); Ok(()) } + } + impl Restartable for ResumeSucceededStage + where + S: HasMetadata, + { fn should_restart(&mut self, state: &mut S) -> Result { TestProgress::should_restart(state, self) } @@ -395,7 +420,12 @@ mod test { } Ok(()) } + } + impl Restartable for ResumeFailedStage + where + S: HasMetadata, + { fn should_restart(&mut self, state: &mut S) -> Result { TestProgress::should_restart(state, self) } @@ -546,7 +576,9 @@ mod test { ) -> Result<(), Error> { panic!("Test failed; panic stage should never be executed."); } + } + impl Restartable for PanicStage { fn should_restart(&mut self, _state: &mut S) -> Result { Ok(true) } @@ -555,7 +587,6 @@ mod test { Ok(()) } } - #[test] fn check_resumability_if_else_if() { let once = RefCell::new(true); diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 6c77aa518e..16a3c578dc 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -87,20 +87,11 @@ pub mod verify_timeouts; /// A stage is one step in the fuzzing process. /// Multiple stages will be scheduled one by one for each input. pub trait Stage { - /// This method will be called before every call to [`Stage::perform`]. - /// Initialize the restart tracking for this stage, _if it is not yet initialized_. - /// On restart, this will be called again. - /// As long as [`Stage::clear_progress`], all subsequent calls happen on restart. - /// Returns `true`, if the stage's [`Stage::perform`] method should run, else `false`. - fn should_restart(&mut self, state: &mut S) -> Result; - - /// Clear the current status tracking of the associated stage - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error>; - /// Run the stage. /// - /// Before a call to perform, [`Stage::should_restart`] will be (must be!) called. - /// After returning (so non-target crash or timeout in a restarting case), [`Stage::clear_progress`] gets called. + /// If you want this stage to restart, then + /// Before a call to perform, [`Restartable::should_restart`] will be (must be!) called. + /// After returning (so non-target crash or timeout in a restarting case), [`Restartable::clear_progress`] gets called. fn perform( &mut self, fuzzer: &mut Z, @@ -110,6 +101,19 @@ pub trait Stage { ) -> Result<(), Error>; } +/// Restartable trait takes care of stage restart. +pub trait Restartable { + /// This method will be called before every call to [`Stage::perform`]. + /// Initialize the restart tracking for this stage, _if it is not yet initialized_. + /// On restart, this will be called again. + /// As long as [`Restartable::clear_progress`], all subsequent calls happen on restart. + /// Returns `true`, if the stage's [`Stage::perform`] method should run, else `false`. + fn should_restart(&mut self, state: &mut S) -> Result; + + /// Clear the current status tracking of the associated stage + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error>; +} + /// A tuple holding all `Stages` used for fuzzing. pub trait StagesTuple { /// Performs all `Stages` in this tuple. @@ -145,7 +149,7 @@ where impl StagesTuple for (Head, Tail) where - Head: Stage, + Head: Stage + Restartable, Tail: StagesTuple + HasConstLen, S: HasCurrentStageId + Stoppable, EM: SendExiting, @@ -169,10 +173,7 @@ where let stage = &mut self.0; - if stage.should_restart(state)? { - stage.perform(fuzzer, executor, state, manager)?; - } - stage.clear_progress(state)?; + stage.perform_restartable(fuzzer, executor, state, manager)?; state.clear_stage_id()?; } @@ -185,10 +186,7 @@ where let stage = &mut self.0; - if stage.should_restart(state)? { - stage.perform(fuzzer, executor, state, manager)?; - } - stage.clear_progress(state)?; + stage.perform_restartable(fuzzer, executor, state, manager)?; state.clear_stage_id()?; } @@ -240,7 +238,36 @@ impl IntoVec>> for Vec StagesTuple for Vec>> +trait RestartableStage: Stage + Restartable { + fn perform_restartable( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + ) -> Result<(), Error>; +} + +impl RestartableStage for ST +where + ST: Stage + Restartable, +{ + /// Run the stage, calling [`Stage::should_restart`] and [`Stage::clear_progress`] appropriately + fn perform_restartable( + &mut self, + fuzzer: &mut Z, + executor: &mut E, + state: &mut S, + manager: &mut EM, + ) -> Result<(), Error> { + if self.should_restart(state)? { + self.perform(fuzzer, executor, state, manager)?; + } + self.clear_progress(state) + } +} + +impl StagesTuple for Vec>> where EM: SendExiting, S: HasCurrentStageId + Stoppable, @@ -261,10 +288,7 @@ where manager.on_shutdown()?; return Err(Error::shutting_down()); } - if stage.should_restart(state)? { - stage.perform(fuzzer, executor, state, manager)?; - } - stage.clear_progress(state) + stage.perform_restartable(fuzzer, executor, state, manager) }) } } @@ -301,7 +325,12 @@ where ) -> Result<(), Error> { (self.closure)(fuzzer, executor, state, manager) } +} +impl Restartable for ClosureStage +where + S: HasNamedMetadata + HasCurrentCorpusId, +{ #[inline] fn should_restart(&mut self, state: &mut S) -> Result { // There's no restart safety in the content of the closure. diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index 4e630ed7ca..5f44a916b4 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -18,7 +18,7 @@ use crate::{ mark_feature_time, mutators::{MultiMutator, MutationResult, Mutator}, nonzero, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, start_timer, state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, MaybeHasClientPerfMonitor}, Error, HasMetadata, HasNamedMetadata, @@ -170,7 +170,12 @@ where ret } +} +impl Restartable for StdMutationalStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { RetryCountRestartHelper::should_restart(state, &self.name, 3) } @@ -316,17 +321,6 @@ where S: HasRand + HasNamedMetadata + HasCurrentTestcase + HasCurrentCorpusId, Z: Evaluator, { - #[inline] - fn should_restart(&mut self, state: &mut S) -> Result { - // Make sure we don't get stuck crashing on a single testcase - RetryCountRestartHelper::should_restart(state, &self.name, 3) - } - - #[inline] - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { - RetryCountRestartHelper::clear_progress(state, &self.name) - } - #[inline] fn perform( &mut self, @@ -354,6 +348,22 @@ where } } +impl Restartable for MultiMutationalStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ + #[inline] + fn should_restart(&mut self, state: &mut S) -> Result { + // Make sure we don't get stuck crashing on a single testcase + RetryCountRestartHelper::should_restart(state, &self.name, 3) + } + + #[inline] + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { + RetryCountRestartHelper::clear_progress(state, &self.name) + } +} + impl MultiMutationalStage { /// Creates a new [`MultiMutationalStage`] pub fn new(mutator: M) -> Self { diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs index fb4ea30913..e19cfe0a9e 100644 --- a/libafl/src/stages/power.rs +++ b/libafl/src/stages/power.rs @@ -19,7 +19,7 @@ use crate::{ schedulers::{testcase_score::CorpusPowerTestcaseScore, TestcaseScore}, stages::{ mutational::{MutatedTransform, MutatedTransformPost}, - MutationalStage, RetryCountRestartHelper, Stage, + MutationalStage, Restartable, RetryCountRestartHelper, Stage, }, start_timer, state::{HasCurrentTestcase, HasExecutions, HasRand, MaybeHasClientPerfMonitor}, @@ -101,7 +101,12 @@ where let ret = self.perform_mutational(fuzzer, executor, state, manager); ret } +} +impl Restartable for PowerMutationalStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { // Make sure we don't get stuck crashing on a single testcase RetryCountRestartHelper::should_restart(state, &self.name, 3) diff --git a/libafl/src/stages/push/mod.rs b/libafl/src/stages/push/mod.rs index 5ad8b6115b..92855ee6e7 100644 --- a/libafl/src/stages/push/mod.rs +++ b/libafl/src/stages/push/mod.rs @@ -28,7 +28,7 @@ use crate::{ executors::{Executor, ExitKind, HasObservers}, observers::ObserversTuple, schedulers::Scheduler, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, state::{HasCorpus, HasExecutions, HasLastReportTime, HasRand}, Error, EvaluatorObservers, ExecutesInput, ExecutionProcessor, HasMetadata, HasScheduler, }; @@ -244,6 +244,22 @@ impl Named for PushStageAdapter { } } +impl Restartable for PushStageAdapter +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ + #[inline] + fn should_restart(&mut self, state: &mut S) -> Result { + // TODO: Proper restart handling - call post_exec at the right time, etc... + RetryCountRestartHelper::no_retry(state, &self.name) + } + + #[inline] + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { + RetryCountRestartHelper::clear_progress(state, &self.name) + } +} + impl Stage for PushStageAdapter where CS: Scheduler, @@ -263,17 +279,6 @@ where + EvaluatorObservers + HasScheduler, { - #[inline] - fn should_restart(&mut self, state: &mut S) -> Result { - // TODO: Proper restart handling - call post_exec at the right time, etc... - RetryCountRestartHelper::no_retry(state, &self.name) - } - - #[inline] - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { - RetryCountRestartHelper::clear_progress(state, &self.name) - } - fn perform( &mut self, fuzzer: &mut Z, diff --git a/libafl/src/stages/replay.rs b/libafl/src/stages/replay.rs index dcef475387..868e829d45 100644 --- a/libafl/src/stages/replay.rs +++ b/libafl/src/stages/replay.rs @@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize}; use crate::{ corpus::{Corpus, CorpusId}, - stages::Stage, + stages::{Restartable, Stage}, state::{HasCorpus, HasSolutions}, Error, Evaluator, }; @@ -114,15 +114,6 @@ where Z: Evaluator, I: Clone, { - fn should_restart(&mut self, _state: &mut S) -> Result { - Ok(true) - } - - fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> { - self.restart_helper.clear(); - Ok(()) - } - #[inline] fn perform( &mut self, @@ -169,3 +160,14 @@ where Ok(()) } } + +impl Restartable for ReplayStage { + fn should_restart(&mut self, _state: &mut S) -> Result { + Ok(true) + } + + fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> { + self.restart_helper.clear(); + Ok(()) + } +} diff --git a/libafl/src/stages/sync.rs b/libafl/src/stages/sync.rs index 4c52c50c34..af6eaf1c50 100644 --- a/libafl/src/stages/sync.rs +++ b/libafl/src/stages/sync.rs @@ -21,7 +21,7 @@ use crate::{ executors::{Executor, ExitKind, HasObservers}, fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor}, inputs::{Input, InputConverter}, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, state::{ HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, HasSolutions, MaybeHasClientPerfMonitor, Stoppable, @@ -85,18 +85,6 @@ where + HasCurrentCorpusId + MaybeHasClientPerfMonitor, { - #[inline] - fn should_restart(&mut self, state: &mut S) -> Result { - // TODO: Needs proper crash handling for when an imported testcase crashes - // For now, Make sure we don't get stuck crashing on this testcase - RetryCountRestartHelper::no_retry(state, &self.name) - } - - #[inline] - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { - RetryCountRestartHelper::clear_progress(state, &self.name) - } - #[inline] fn perform( &mut self, @@ -158,6 +146,23 @@ where } } +impl Restartable for SyncFromDiskStage +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ + #[inline] + fn should_restart(&mut self, state: &mut S) -> Result { + // TODO: Needs proper crash handling for when an imported testcase crashes + // For now, Make sure we don't get stuck crashing on this testcase + RetryCountRestartHelper::no_retry(state, &self.name) + } + + #[inline] + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { + RetryCountRestartHelper::clear_progress(state, &self.name) + } +} + impl SyncFromDiskStage { /// Creates a new [`SyncFromDiskStage`] #[must_use] @@ -249,18 +254,6 @@ where SP: ShMemProvider, Z: EvaluatorObservers + ExecutionProcessor, { - #[inline] - fn should_restart(&mut self, _state: &mut S) -> Result { - // No restart handling needed - does not execute the target. - Ok(true) - } - - #[inline] - fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> { - // Not needed - does not execute the target. - Ok(()) - } - #[inline] fn perform( &mut self, @@ -320,6 +313,20 @@ where } } +impl Restartable for SyncFromBrokerStage { + #[inline] + fn should_restart(&mut self, _state: &mut S) -> Result { + // No restart handling needed - does not execute the target. + Ok(true) + } + + #[inline] + fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> { + // Not needed - does not execute the target. + Ok(()) + } +} + impl SyncFromBrokerStage { /// Creates a new [`SyncFromBrokerStage`] #[must_use] diff --git a/libafl/src/stages/time_tracker.rs b/libafl/src/stages/time_tracker.rs index 6e4f57792e..1e6d75df3f 100644 --- a/libafl/src/stages/time_tracker.rs +++ b/libafl/src/stages/time_tracker.rs @@ -3,7 +3,10 @@ use std::{marker::PhantomData, time::Duration}; use libafl_bolts::{current_time, Error}; -use crate::{stages::Stage, HasMetadata}; +use crate::{ + stages::{Restartable, Stage}, + HasMetadata, +}; /// Track an inner Stage's execution time #[derive(Debug)] pub struct TimeTrackingStageWrapper { @@ -48,7 +51,12 @@ where } Ok(()) } +} +impl Restartable for TimeTrackingStageWrapper +where + ST: Restartable, +{ fn should_restart(&mut self, state: &mut S) -> Result { self.inner.should_restart(state) } diff --git a/libafl/src/stages/tmin.rs b/libafl/src/stages/tmin.rs index 26749c7840..20667f950a 100644 --- a/libafl/src/stages/tmin.rs +++ b/libafl/src/stages/tmin.rs @@ -30,7 +30,7 @@ use crate::{ schedulers::RemovableScheduler, stages::{ mutational::{MutatedTransform, MutatedTransformPost}, - ExecutionCountRestartHelper, Stage, + ExecutionCountRestartHelper, Restartable, Stage, }, start_timer, state::{ @@ -82,14 +82,6 @@ where M: Mutator, I: Input + Hash + HasLen, { - fn should_restart(&mut self, state: &mut S) -> Result { - self.restart_helper.should_restart(state, &self.name) - } - - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { - self.restart_helper.clear_progress::(state, &self.name) - } - fn perform( &mut self, fuzzer: &mut Z, @@ -106,6 +98,19 @@ where } } +impl Restartable for StdTMinMutationalStage +where + S: HasNamedMetadata + HasExecutions, +{ + fn should_restart(&mut self, state: &mut S) -> Result { + self.restart_helper.should_restart(state, &self.name) + } + + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { + self.restart_helper.clear_progress::(state, &self.name) + } +} + impl FeedbackFactory for StdTMinMutationalStage where diff --git a/libafl/src/stages/tracing.rs b/libafl/src/stages/tracing.rs index 1952a40179..7430358ba4 100644 --- a/libafl/src/stages/tracing.rs +++ b/libafl/src/stages/tracing.rs @@ -16,7 +16,7 @@ use crate::{ inputs::Input, mark_feature_time, observers::ObserversTuple, - stages::{RetryCountRestartHelper, Stage}, + stages::{Restartable, RetryCountRestartHelper, Stage}, start_timer, state::{HasCorpus, HasCurrentTestcase, HasExecutions, MaybeHasClientPerfMonitor}, Error, HasNamedMetadata, @@ -93,7 +93,12 @@ where ) -> Result<(), Error> { self.trace(fuzzer, state, manager) } +} +impl Restartable for TracingStage +where + S: HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { RetryCountRestartHelper::no_retry(state, &self.name) } @@ -174,14 +179,6 @@ where + HasCurrentCorpusId + MaybeHasClientPerfMonitor, { - fn should_restart(&mut self, state: &mut S) -> Result { - RetryCountRestartHelper::no_retry(state, &self.name) - } - - fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { - RetryCountRestartHelper::clear_progress(state, &self.name) - } - #[inline] fn perform( &mut self, @@ -219,6 +216,19 @@ where } } +impl Restartable for ShadowTracingStage +where + S: HasNamedMetadata + HasCurrentCorpusId, +{ + fn should_restart(&mut self, state: &mut S) -> Result { + RetryCountRestartHelper::no_retry(state, &self.name) + } + + fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> { + RetryCountRestartHelper::clear_progress(state, &self.name) + } +} + impl ShadowTracingStage where E: Executor + HasObservers, diff --git a/libafl/src/stages/tuneable.rs b/libafl/src/stages/tuneable.rs index d2535e11dc..92764d8352 100644 --- a/libafl/src/stages/tuneable.rs +++ b/libafl/src/stages/tuneable.rs @@ -14,7 +14,7 @@ use crate::{ nonzero, stages::{ mutational::{MutatedTransform, MutatedTransformPost, DEFAULT_MUTATIONAL_MAX_ITERATIONS}, - ExecutionCountRestartHelper, MutationalStage, Stage, + ExecutionCountRestartHelper, MutationalStage, Restartable, Stage, }, start_timer, state::{HasCurrentTestcase, HasExecutions, HasRand, MaybeHasClientPerfMonitor}, @@ -217,7 +217,12 @@ where ret } +} +impl Restartable for TuneableMutationalStage +where + S: HasNamedMetadata + HasExecutions, +{ fn should_restart(&mut self, state: &mut S) -> Result { self.restart_helper.should_restart(state, &self.name) } diff --git a/libafl/src/stages/unicode.rs b/libafl/src/stages/unicode.rs index a75b30a36f..79ff778164 100644 --- a/libafl/src/stages/unicode.rs +++ b/libafl/src/stages/unicode.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; use crate::{ inputs::{BytesInput, HasTargetBytes}, - stages::Stage, + stages::{Restartable, Stage}, state::{HasCorpus, HasCurrentTestcase}, HasMetadata, }; @@ -112,6 +112,18 @@ impl Stage for UnicodeIdentificationStage + HasCurrentTestcase, { + fn perform( + &mut self, + _fuzzer: &mut Z, + _executor: &mut E, + state: &mut S, + _manager: &mut EM, + ) -> Result<(), Error> { + UnicodeIdentificationStage::identify_unicode_in_current_testcase(state) + } +} + +impl Restartable for UnicodeIdentificationStage { #[inline] fn should_restart(&mut self, _state: &mut S) -> Result { // Stage does not run the target. No reset helper needed. @@ -123,14 +135,4 @@ where // Stage does not run the target. No reset helper needed. Ok(()) } - - fn perform( - &mut self, - _fuzzer: &mut Z, - _executor: &mut E, - state: &mut S, - _manager: &mut EM, - ) -> Result<(), Error> { - UnicodeIdentificationStage::identify_unicode_in_current_testcase(state) - } } diff --git a/libafl/src/stages/verify_timeouts.rs b/libafl/src/stages/verify_timeouts.rs index 668df1f3f9..bad419487e 100644 --- a/libafl/src/stages/verify_timeouts.rs +++ b/libafl/src/stages/verify_timeouts.rs @@ -13,7 +13,7 @@ use crate::{ executors::{Executor, HasObservers, HasTimeout}, inputs::BytesInput, observers::ObserversTuple, - stages::Stage, + stages::{Restartable, Stage}, Evaluator, HasMetadata, }; @@ -111,6 +111,9 @@ where *res = TimeoutsToVerify::::new(); Ok(()) } +} + +impl Restartable for VerifyTimeoutsStage { fn should_restart(&mut self, _state: &mut S) -> Result { Ok(true) } diff --git a/libafl_targets/src/cmps/stages/aflpptracing.rs b/libafl_targets/src/cmps/stages/aflpptracing.rs index 54309d8a25..71782a46f1 100644 --- a/libafl_targets/src/cmps/stages/aflpptracing.rs +++ b/libafl_targets/src/cmps/stages/aflpptracing.rs @@ -6,7 +6,7 @@ use libafl::{ executors::{Executor, HasObservers}, inputs::BytesInput, observers::ObserversTuple, - stages::{colorization::TaintMetadata, RetryCountRestartHelper, Stage}, + stages::{colorization::TaintMetadata, Restartable, RetryCountRestartHelper, Stage}, state::{HasCorpus, HasCurrentTestcase}, Error, HasMetadata, HasNamedMetadata, }; @@ -111,7 +111,12 @@ where Ok(()) } +} +impl Restartable for AFLppCmplogTracingStage<'_, EM, TE, S, Z> +where + S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId, +{ fn should_restart(&mut self, state: &mut S) -> Result { // Tracing stage is always deterministic // don't restart