SYS-OSS/FRET-LibAFL
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Restartable trait (#3004)

Browse Source 
* separate into Restartable

* typo

* fixes, blanket impl

* aa

* lol

* aa

* fixer

* lol
This commit is contained in:
Dongjia "toka" Zhang 2025-02-19 15:06:31 +01:00 committed by GitHub
parent 0e4c6722f0
commit 61568a9be1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
22 changed files with 322 additions and 150 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
fuzzers/forkserver/libafl-fuzz/src/stages/mutational_stage.rs
View File

@ -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<SM, P> {
pub enum SupportedMutationalStages<P, SM> {
StdMutational(SM, PhantomData<P>),
PowerMutational(P, PhantomData<SM>),
}
impl<S, SM, P> MutationalStage<S> for SupportedMutationalStages<SM, P>
impl<P, S, SM> MutationalStage<S> for SupportedMutationalStages<P, SM>
where
SM: MutationalStage<S>,
P: MutationalStage<S, Mutator = SM::Mutator>,
SM: MutationalStage<S>,
{
type Mutator = SM::Mutator;
/// The mutator, added to this stage
@ -44,10 +44,10 @@ where
}
}
impl<SM, P> Named for SupportedMutationalStages<SM, P>
impl<P, SM> Named for SupportedMutationalStages<P, SM>
where
SM: Named,
P: Named,
SM: Named,
{
fn name(&self) -> &Cow<'static, str> {
match self {
@ -57,10 +57,10 @@ where
}
}
impl<E, EM, S, SM, P, Z> Stage<E, EM, S, Z> for SupportedMutationalStages<SM, P>
impl<E, EM, P, S, SM, Z> Stage<E, EM, S, Z> for SupportedMutationalStages<P, SM>
where
SM: Stage<E, EM, S, Z>,
P: Stage<E, EM, S, Z>,
SM: Stage<E, EM, S, Z>,
{
#[inline]
fn perform(
@ -75,7 +75,13 @@ where
Self::PowerMutational(p, _) => p.perform(fuzzer, executor, state, manager),
}
}
}
impl<P, S, SM> Restartable<S> for SupportedMutationalStages<P, SM>
where
P: Restartable<S>,
SM: Restartable<S>,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
match self {
Self::StdMutational(m, _) => m.should_restart(state),

4
libafl/src/stages/afl_stats.rs
View File

@ -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<C, E, EM, I, O, S, Z> Restartable<S> for AflStatsStage<C, E, EM, I, O, S, Z> {
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
Ok(true)
}

7
libafl/src/stages/calibrate.rs
View File

@ -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<C, E, I, O, OT, S> Restartable<S> for CalibrationStage<C, E, I, O, OT, S>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// Calibration stage disallow restarts
// If a testcase that causes crash/timeout in the queue, we need to remove it from the queue immediately.

7
libafl/src/stages/colorization.rs
View File

@ -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<C, E, EM, I, O, S, Z> Restartable<S> for ColorizationStage<C, E, EM, I, O, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// This is a deterministic stage
// Once it failed, then don't retry,

13
libafl/src/stages/concolic.rs
View File

@ -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<EM, I, TE, S, Z> Restartable<S> for ConcolicTracingStage<'_, EM, I, TE, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// This is a deterministic stage
// Once it failed, then don't retry,
@ -418,7 +423,13 @@ where
}
Ok(())
}
}
#[cfg(feature = "concolic_mutation")]
impl<I, S, Z> Restartable<S> for SimpleConcolicMutationalStage<I, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// This is a deterministic stage

6
libafl/src/stages/dump.rs
View File

@ -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<CB1, EM, I, S, Z> Restartable<S>
for DumpToDiskStage<CB1, fn(&Testcase<I>, &CorpusId) -> String, EM, I, S, Z>
{
#[inline]
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
// Not executing the target, so restart safety is not needed

31
libafl/src/stages/generalization.rs
View File

@ -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<C, EM, I, O, OT, S, Z> Named for GeneralizationStage<C, EM, I, O, OT, S, Z>
}
}
impl<C, EM, I, O, OT, S, Z> Restartable<S> for GeneralizationStage<C, EM, I, O, OT, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// TODO: We need to be able to resume better if something crashes or times out
RetryCountRestartHelper::should_restart::<S>(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::<S>(state, &self.name)
}
}
impl<C, E, EM, O, S, Z> Stage<E, EM, S, Z>
for GeneralizationStage<C, EM, BytesInput, O, E::Observers, S, Z>
where
@ -76,18 +93,6 @@ where
+ HasCurrentCorpusId
+ MaybeHasClientPerfMonitor,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// TODO: We need to be able to resume better if something crashes or times out
RetryCountRestartHelper::should_restart::<S>(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::<S>(state, &self.name)
}
#[inline]
#[expect(clippy::too_many_lines)]
fn perform(

9
libafl/src/stages/generation.rs
View File

@ -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<G, I, S, Z> Restartable<S> for GenStage<G, I, S, Z> {
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
// It's a random generation stage
// so you can restart for whatever times you want

39
libafl/src/stages/logics.rs
View File

@ -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<CB, E, EM, ST, S, Z> Restartable<S> for WhileStage<CB, E, EM, ST, S, Z>
where
S: HasNestedStage,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
NestedStageRetryCountRestartHelper::should_restart(state, self)
}
@ -112,7 +117,12 @@ where
}
Ok(())
}
}
impl<CB, E, EM, ST, S, Z> Restartable<S> for IfStage<CB, E, EM, ST, S, Z>
where
S: HasNestedStage,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
NestedStageRetryCountRestartHelper::should_restart(state, self)
}
@ -188,7 +198,12 @@ where
Ok(())
}
}
impl<CB, E, EM, ST1, ST2, S, Z> Restartable<S> for IfElseStage<CB, E, EM, ST1, ST2, S, Z>
where
S: HasNestedStage,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
NestedStageRetryCountRestartHelper::should_restart(state, self)
}
@ -238,7 +253,12 @@ where
Ok(())
}
}
}
impl<E, EM, ST, S, Z> Restartable<S> for OptionalStage<E, EM, ST, S, Z>
where
S: HasNestedStage,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
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<S> Restartable<S> for ResumeSucceededStage<S>
where
S: HasMetadata,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
TestProgress::should_restart(state, self)
}
@ -395,7 +420,12 @@ mod test {
}
Ok(())
}
}
impl<S> Restartable<S> for ResumeFailedStage<S>
where
S: HasMetadata,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
TestProgress::should_restart(state, self)
}
@ -546,7 +576,9 @@ mod test {
) -> Result<(), Error> {
panic!("Test failed; panic stage should never be executed.");
}
}
impl<S> Restartable<S> for PanicStage<S> {
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
Ok(true)
}
@ -555,7 +587,6 @@ mod test {
Ok(())
}
}
#[test]
fn check_resumability_if_else_if() {
let once = RefCell::new(true);

81
libafl/src/stages/mod.rs
View File

@ -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<E, EM, S, Z> {
/// 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<bool, Error>;
/// 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<E, EM, S, Z> {
) -> Result<(), Error>;
}
/// Restartable trait takes care of stage restart.
pub trait Restartable<S> {
/// 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<bool, Error>;
/// 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<E, EM, S, Z> {
/// Performs all `Stages` in this tuple.
@ -145,7 +149,7 @@ where
impl<Head, Tail, E, EM, S, Z> StagesTuple<E, EM, S, Z> for (Head, Tail)
where
Head: Stage<E, EM, S, Z>,
Head: Stage<E, EM, S, Z> + Restartable<S>,
Tail: StagesTuple<E, EM, S, Z> + 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<E, EM, S, Z> IntoVec<Box<dyn Stage<E, EM, S, Z>>> for Vec<Box<dyn Stage<E,
}
}
impl<E, EM, S, Z> StagesTuple<E, EM, S, Z> for Vec<Box<dyn Stage<E, EM, S, Z>>>
trait RestartableStage<E, EM, S, Z>: Stage<E, EM, S, Z> + Restartable<S> {
fn perform_restartable(
&mut self,
fuzzer: &mut Z,
executor: &mut E,
state: &mut S,
manager: &mut EM,
) -> Result<(), Error>;
}
impl<E, EM, S, ST, Z> RestartableStage<E, EM, S, Z> for ST
where
ST: Stage<E, EM, S, Z> + Restartable<S>,
{
/// 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<E, EM, S, Z> StagesTuple<E, EM, S, Z> for Vec<Box<dyn RestartableStage<E, EM, S, Z>>>
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<CB, E, EM, S, Z> Restartable<S> for ClosureStage<CB, E, EM, Z>
where
S: HasNamedMetadata + HasCurrentCorpusId,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// There's no restart safety in the content of the closure.

34
libafl/src/stages/mutational.rs
View File

@ -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<E, EM, I1, I2, M, S, Z> Restartable<S> for StdMutationalStage<E, EM, I1, I2, M, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
RetryCountRestartHelper::should_restart(state, &self.name, 3)
}
@ -316,17 +321,6 @@ where
S: HasRand + HasNamedMetadata + HasCurrentTestcase<I> + HasCurrentCorpusId,
Z: Evaluator<E, EM, I, S>,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// 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<E, EM, I, M, S, Z> Restartable<S> for MultiMutationalStage<E, EM, I, M, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// 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<E, EM, I, M, S, Z> MultiMutationalStage<E, EM, I, M, S, Z> {
/// Creates a new [`MultiMutationalStage`]
pub fn new(mutator: M) -> Self {

7
libafl/src/stages/power.rs
View File

@ -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<E, F, EM, I, M, S, Z> Restartable<S> for PowerMutationalStage<E, F, EM, I, M, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// Make sure we don't get stuck crashing on a single testcase
RetryCountRestartHelper::should_restart(state, &self.name, 3)

29
libafl/src/stages/push/mod.rs
View File

@ -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<CS, EM, I, OT, PS, Z> Named for PushStageAdapter<CS, EM, I, OT, PS, Z> {
}
}
impl<CS, EM, I, OT, PS, S, Z> Restartable<S> for PushStageAdapter<CS, EM, I, OT, PS, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// 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<CS, E, EM, I, OT, PS, S, Z> Stage<E, EM, S, Z> for PushStageAdapter<CS, EM, I, OT, PS, Z>
where
CS: Scheduler<I, S>,
@ -263,17 +279,6 @@ where
+ EvaluatorObservers<E, EM, I, OT>
+ HasScheduler<I, S>,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// 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,

22
libafl/src/stages/replay.rs
View File

@ -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<E, EM, I, S>,
I: Clone,
{
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
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<I, S> Restartable<S> for ReplayStage<I> {
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
Ok(true)
}
fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
self.restart_helper.clear();
Ok(())
}
}

57
libafl/src/stages/sync.rs
View File

@ -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<bool, Error> {
// 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<CB, E, EM, I, S, Z> Restartable<S> for SyncFromDiskStage<CB, E, EM, I, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
#[inline]
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// 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<CB, E, EM, I, S, Z> SyncFromDiskStage<CB, E, EM, I, S, Z> {
/// Creates a new [`SyncFromDiskStage`]
#[must_use]
@ -249,18 +254,6 @@ where
SP: ShMemProvider<ShMem = SHM>,
Z: EvaluatorObservers<E, EM, I, S> + ExecutionProcessor<EM, I, E::Observers, S>,
{
#[inline]
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
// 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<I, IC, ICB, S, SHM, SP> Restartable<S> for SyncFromBrokerStage<I, IC, ICB, S, SHM, SP> {
#[inline]
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
// 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<I, IC, ICB, S, SHM, SP> SyncFromBrokerStage<I, IC, ICB, S, SHM, SP> {
/// Creates a new [`SyncFromBrokerStage`]
#[must_use]

10
libafl/src/stages/time_tracker.rs
View File

@ -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<T, S, ST> {
@ -48,7 +51,12 @@ where
}
Ok(())
}
}
impl<T, S, ST> Restartable<S> for TimeTrackingStageWrapper<T, S, ST>
where
ST: Restartable<S>,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
self.inner.should_restart(state)
}

23
libafl/src/stages/tmin.rs
View File

@ -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, S>,
I: Input + Hash + HasLen,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
self.restart_helper.should_restart(state, &self.name)
}
fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
self.restart_helper.clear_progress::<S>(state, &self.name)
}
fn perform(
&mut self,
fuzzer: &mut Z,
@ -106,6 +98,19 @@ where
}
}
impl<E, EM, F, FF, I, M, S, Z> Restartable<S> for StdTMinMutationalStage<E, EM, F, FF, I, M, S, Z>
where
S: HasNamedMetadata + HasExecutions,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
self.restart_helper.should_restart(state, &self.name)
}
fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
self.restart_helper.clear_progress::<S>(state, &self.name)
}
}
impl<E, EM, F, FF, I, M, S, Z> FeedbackFactory<F, E::Observers>
for StdTMinMutationalStage<E, EM, F, FF, I, M, S, Z>
where

28
libafl/src/stages/tracing.rs
View File

@ -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<EM, I, TE, S, Z> Restartable<S> for TracingStage<EM, I, TE, S, Z>
where
S: HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
RetryCountRestartHelper::no_retry(state, &self.name)
}
@ -174,14 +179,6 @@ where
+ HasCurrentCorpusId
+ MaybeHasClientPerfMonitor,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
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<E, EM, I, SOT, S, Z> Restartable<S> for ShadowTracingStage<E, EM, I, SOT, S, Z>
where
S: HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
RetryCountRestartHelper::no_retry(state, &self.name)
}
fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
RetryCountRestartHelper::clear_progress(state, &self.name)
}
}
impl<E, EM, I, SOT, S, Z> ShadowTracingStage<E, EM, I, SOT, S, Z>
where
E: Executor<EM, I, S, Z> + HasObservers,

7
libafl/src/stages/tuneable.rs
View File

@ -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<E, EM, I, M, S, Z> Restartable<S> for TuneableMutationalStage<E, EM, I, M, S, Z>
where
S: HasNamedMetadata + HasExecutions,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
self.restart_helper.should_restart(state, &self.name)
}

24
libafl/src/stages/unicode.rs
View File

@ -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<E, EM, S, Z> Stage<E, EM, S, Z> for UnicodeIdentificationStage<BytesInput,
where
S: HasCorpus<BytesInput> + HasCurrentTestcase<BytesInput>,
{
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<S> Restartable<S> for UnicodeIdentificationStage<BytesInput, S> {
#[inline]
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
// 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)
}
}

5
libafl/src/stages/verify_timeouts.rs
View File

@ -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::<I>::new();
Ok(())
}
}
impl<E, I, S> Restartable<S> for VerifyTimeoutsStage<E, I, S> {
fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
Ok(true)
}

7
libafl_targets/src/cmps/stages/aflpptracing.rs
View File

@ -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<EM, TE, S, Z> Restartable<S> for AFLppCmplogTracingStage<'_, EM, TE, S, Z>
where
S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
{
fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
// Tracing stage is always deterministic
// don't restart