From e3dd7cf0dcdae465ff938a80d843766be375c389 Mon Sep 17 00:00:00 2001 From: "Dongjia \"toka\" Zhang" Date: Thu, 20 Jun 2024 17:38:15 +0200 Subject: [PATCH] Don't restart in deterministic stages. Don't restart where there's no restart safety. Make stage names unique (#2331) * push * fuck * add * add * api * api * add multi machine to workspace * doc * api * api * add * more * fix * stats * rev * fix * fix * real fix * add * fmt * add * add * fix * a * add * revert workflow --------- Co-authored-by: Your Name --- .github/workflows/build_and_test.yml | 1 + Cargo.toml | 1 + fuzzers/baby_fuzzer_minimizing/src/main.rs | 2 +- .../fuzzbench_forkserver_cmplog/src/main.rs | 2 +- .../fuzzer/src/main.rs | 4 +- libafl/src/fuzzer/mod.rs | 3 +- libafl/src/stages/calibrate.rs | 44 ++-- libafl/src/stages/colorization.rs | 24 +- libafl/src/stages/concolic.rs | 82 +++++-- libafl/src/stages/dump.rs | 4 +- libafl/src/stages/generalization.rs | 26 ++- libafl/src/stages/generation.rs | 6 +- libafl/src/stages/logics.rs | 40 ++-- libafl/src/stages/mod.rs | 218 +++++++++++------- libafl/src/stages/mutational.rs | 82 ++++--- libafl/src/stages/power.rs | 44 ++-- libafl/src/stages/stats.rs | 4 +- libafl/src/stages/sync.rs | 20 +- libafl/src/stages/tmin.rs | 55 ++++- libafl/src/stages/tracing.rs | 79 ++++--- libafl/src/stages/tuneable.rs | 22 +- libafl/src/stages/unicode.rs | 4 +- libafl_bolts/src/lib.rs | 2 +- libafl_bolts/src/llmp.rs | 8 +- .../src/cmps/stages/aflpptracing.rs | 86 ++++--- utils/multi_machine_generator/src/graph.rs | 52 +++-- utils/multi_machine_generator/src/main.rs | 7 +- 27 files changed, 548 insertions(+), 374 deletions(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index fa0633be66..a61592b33c 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -277,6 +277,7 @@ jobs: - ubuntu - fuzzers-preflight strategy: + fail-fast: true matrix: os: [ ubuntu-latest ] fuzzer: diff --git a/Cargo.toml b/Cargo.toml index 03ca9953b4..4c2092d80f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ members = [ "utils/libafl_benches", "utils/gramatron/construct_automata", "utils/desyscall", + "utils/multi_machine_generator", ] default-members = [ "libafl", diff --git a/fuzzers/baby_fuzzer_minimizing/src/main.rs b/fuzzers/baby_fuzzer_minimizing/src/main.rs index 188e73c087..1f32b4ff08 100644 --- a/fuzzers/baby_fuzzer_minimizing/src/main.rs +++ b/fuzzers/baby_fuzzer_minimizing/src/main.rs @@ -125,7 +125,7 @@ pub fn main() -> Result<(), Error> { let mut stages = tuple_list!(StdTMinMutationalStage::new( minimizer, CrashFeedback::new(), - 1 << 10 + 1 << 10, )); let scheduler = QueueScheduler::new(); diff --git a/fuzzers/fuzzbench_forkserver_cmplog/src/main.rs b/fuzzers/fuzzbench_forkserver_cmplog/src/main.rs index 403296744d..79e1568e1b 100644 --- a/fuzzers/fuzzbench_forkserver_cmplog/src/main.rs +++ b/fuzzers/fuzzbench_forkserver_cmplog/src/main.rs @@ -368,7 +368,7 @@ fn fuzz( .build(tuple_list!(cmplog_observer)) .unwrap(); - let tracing = AFLppCmplogTracingStage::with_cmplog_observer(cmplog_executor, cmplog_ref); + let tracing = AFLppCmplogTracingStage::new(cmplog_executor, cmplog_ref); // Setup a randomic Input2State stage let rq = MultiMutationalStage::new(AFLppRedQueen::with_cmplog_options(true, true)); diff --git a/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs b/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs index f68bea3abf..cbe6e1d7fd 100644 --- a/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs +++ b/fuzzers/libfuzzer_stb_image_concolic/fuzzer/src/main.rs @@ -214,12 +214,12 @@ fn fuzz( // Create a concolic trace ConcolicTracingStage::new( TracingStage::new( - MyCommandConfigurator.into_executor(tuple_list!(concolic_observer)) + MyCommandConfigurator.into_executor(tuple_list!(concolic_observer)), ), concolic_ref, ), // Use the concolic trace for z3-based solving - SimpleConcolicMutationalStage::default(), + SimpleConcolicMutationalStage::new(), ); fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut restarting_mgr)?; diff --git a/libafl/src/fuzzer/mod.rs b/libafl/src/fuzzer/mod.rs index 433078b3fd..0934b3f89f 100644 --- a/libafl/src/fuzzer/mod.rs +++ b/libafl/src/fuzzer/mod.rs @@ -27,7 +27,7 @@ use crate::{ use crate::{monitors::PerfFeature, state::HasClientPerfMonitor}; /// Send a monitor update all 15 (or more) seconds -const STATS_TIMEOUT_DEFAULT: Duration = Duration::from_secs(1); +const STATS_TIMEOUT_DEFAULT: Duration = Duration::from_secs(15); /// Holds a scheduler pub trait HasScheduler: UsesState @@ -247,7 +247,6 @@ where let monitor_timeout = STATS_TIMEOUT_DEFAULT; for _ in 0..iters { - // log::info!("Starting another fuzz_loop"); manager.maybe_report_progress(state, monitor_timeout)?; ret = Some(self.fuzz_one(stages, executor, state, manager)?); } diff --git a/libafl/src/stages/calibrate.rs b/libafl/src/stages/calibrate.rs index ae9d43a2f3..0296c818f7 100644 --- a/libafl/src/stages/calibrate.rs +++ b/libafl/src/stages/calibrate.rs @@ -1,6 +1,9 @@ //! The calibration stage. The fuzzer measures the average exec time and the bitmap size. -use alloc::{borrow::Cow, vec::Vec}; +use alloc::{ + borrow::{Cow, ToOwned}, + vec::Vec, +}; use core::{fmt::Debug, marker::PhantomData, time::Duration}; use hashbrown::HashSet; @@ -17,7 +20,7 @@ use crate::{ monitors::{AggregatorOps, UserStats, UserStatsValue}, observers::{MapObserver, ObserversTuple}, schedulers::powersched::SchedulerMetadata, - stages::{ExecutionCountRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState}, Error, HasMetadata, HasNamedMetadata, }; @@ -75,7 +78,6 @@ pub struct CalibrationStage { stage_max: usize, /// If we should track stability track_stability: bool, - restart_helper: ExecutionCountRestartHelper, phantom: PhantomData<(E, O, OT)>, } @@ -125,8 +127,6 @@ where let mut iter = self.stage_max; // If we restarted after a timeout or crash, do less iterations. - iter -= usize::try_from(self.restart_helper.execs_since_progress_start(state)?)?; - let input = state.current_input_cloned()?; // Run once to get the initial calibration map @@ -350,14 +350,18 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - // TODO: Make sure this is the correct way / there may be a better way? - self.restart_helper.restart_progress_should_run(state) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // Calibration stage disallow restarts + // If a testcase that causes crash/timeout in the queue, we need to remove it from the queue immediately. + StdRestartHelper::no_retry(state, &self.name) + + // todo + // remove this guy from corpus queue } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { // TODO: Make sure this is the correct way / there may be a better way? - self.restart_helper.clear_restart_progress(state) + StdRestartHelper::clear_progress(state, &self.name) } } @@ -375,14 +379,16 @@ where where F: HasObserverHandle + Named, { + let map_name = map_feedback.name().clone(); Self { map_observer_handle: map_feedback.observer_handle().clone(), - map_name: map_feedback.name().clone(), + map_name: map_name.clone(), stage_max: CAL_STAGE_START, track_stability: true, - restart_helper: ExecutionCountRestartHelper::default(), phantom: PhantomData, - name: Cow::Borrowed(CALIBRATION_STAGE_NAME), + name: Cow::Owned( + CALIBRATION_STAGE_NAME.to_owned() + ":" + map_name.into_owned().as_str(), + ), } } @@ -392,15 +398,9 @@ where where F: HasObserverHandle + Named, { - Self { - map_observer_handle: map_feedback.observer_handle().clone(), - map_name: map_feedback.name().clone(), - stage_max: CAL_STAGE_START, - track_stability: false, - restart_helper: ExecutionCountRestartHelper::default(), - phantom: PhantomData, - name: Cow::Borrowed(CALIBRATION_STAGE_NAME), - } + let mut ret = Self::new(map_feedback); + ret.track_stability = false; + ret } } diff --git a/libafl/src/stages/colorization.rs b/libafl/src/stages/colorization.rs index a3f8844a69..2edab93288 100644 --- a/libafl/src/stages/colorization.rs +++ b/libafl/src/stages/colorization.rs @@ -1,5 +1,9 @@ //! The colorization stage from `colorization()` in afl++ -use alloc::{borrow::Cow, collections::binary_heap::BinaryHeap, vec::Vec}; +use alloc::{ + borrow::{Cow, ToOwned}, + collections::binary_heap::BinaryHeap, + vec::Vec, +}; use core::{cmp::Ordering, fmt::Debug, marker::PhantomData, ops::Range}; use libafl_bolts::{ @@ -15,7 +19,7 @@ use crate::{ inputs::HasMutatorBytes, mutators::mutations::buffer_copy, observers::{MapObserver, ObserversTuple}, - stages::{RetryRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, state::{HasCorpus, HasCurrentTestcase, HasRand, UsesState}, Error, HasMetadata, HasNamedMetadata, }; @@ -104,14 +108,15 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - // TODO this stage needs a proper resume - RetryRestartHelper::restart_progress_should_run(state, self, 3) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // This is a deterministic stage + // Once it failed, then don't retry, + // It will just fail again + StdRestartHelper::no_retry(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - // TODO this stage needs a proper resume - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } @@ -309,9 +314,10 @@ where #[must_use] /// Creates a new [`ColorizationStage`] pub fn new(map_observer: &C) -> Self { + let obs_name = map_observer.name().clone().into_owned(); Self { map_observer_handle: map_observer.handle(), - name: Cow::Borrowed(COLORIZATION_STAGE_NAME), + name: Cow::Owned(COLORIZATION_STAGE_NAME.to_owned() + ":" + obs_name.as_str()), phantom: PhantomData, } } diff --git a/libafl/src/stages/concolic.rs b/libafl/src/stages/concolic.rs index 515b615870..3d82a5a2ad 100644 --- a/libafl/src/stages/concolic.rs +++ b/libafl/src/stages/concolic.rs @@ -1,8 +1,7 @@ //! This module contains the `concolic` stages, which can trace a target using symbolic execution //! and use the results for fuzzer input and mutations. //! - -use alloc::borrow::Cow; +use alloc::borrow::{Cow, ToOwned}; #[cfg(feature = "concolic_mutation")] use alloc::{string::ToString, vec::Vec}; #[cfg(feature = "concolic_mutation")] @@ -20,7 +19,7 @@ use crate::state::HasClientPerfMonitor; use crate::{ executors::{Executor, HasObservers}, observers::concolic::ConcolicObserver, - stages::{RetryRestartHelper, Stage, TracingStage}, + stages::{Stage, StdRestartHelper, TracingStage}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState}, Error, HasMetadata, HasNamedMetadata, }; @@ -29,7 +28,6 @@ use crate::{ inputs::HasMutatorBytes, mark_feature_time, observers::concolic::{ConcolicMetadata, SymExpr, SymExprRef}, - stages::ExecutionCountRestartHelper, start_timer, state::State, Evaluator, @@ -38,6 +36,7 @@ use crate::{ /// Wraps a [`TracingStage`] to add concolic observing. #[derive(Clone, Debug)] pub struct ConcolicTracingStage<'a, EM, TE, Z> { + name: Cow<'static, str>, inner: TracingStage, observer_handle: Handle>, } @@ -49,10 +48,12 @@ where type State = TE::State; } +/// The name for concolic tracer +pub const CONCOLIC_TRACING_STAGE_NAME: &str = "concolictracing"; + impl Named for ConcolicTracingStage<'_, EM, TE, Z> { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed("ConcolicTracingStage"); - &NAME + &self.name } } @@ -83,12 +84,15 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - RetryRestartHelper::restart_progress_should_run(state, self, 3) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // This is a deterministic stage + // Once it failed, then don't retry, + // It will just fail again + StdRestartHelper::no_retry(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } @@ -99,9 +103,13 @@ impl<'a, EM, TE, Z> ConcolicTracingStage<'a, EM, TE, Z> { inner: TracingStage, observer_handle: Handle>, ) -> Self { + let observer_name = observer_handle.name().clone(); Self { inner, observer_handle, + name: Cow::Owned( + CONCOLIC_TRACING_STAGE_NAME.to_owned() + ":" + observer_name.into_owned().as_str(), + ), } } } @@ -351,10 +359,9 @@ fn generate_mutations(iter: impl Iterator) -> Vec< /// A mutational stage that uses Z3 to solve concolic constraints attached to the [`crate::corpus::Testcase`] by the [`ConcolicTracingStage`]. #[cfg(feature = "concolic_mutation")] -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Default)] pub struct SimpleConcolicMutationalStage { - /// The helper keeps track of progress for timeouting/restarting targets - restart_helper: ExecutionCountRestartHelper, + name: Cow<'static, str>, phantom: PhantomData, } @@ -366,6 +373,21 @@ where type State = Z::State; } +#[cfg(feature = "concolic_mutation")] +/// The unique id for this stage +static mut SIMPLE_CONCOLIC_MUTATIONAL_ID: usize = 0; + +#[cfg(feature = "concolic_mutation")] +/// The name for concolic mutation stage +pub const SIMPLE_CONCOLIC_MUTATIONAL_NAME: &str = "concolicmutation"; + +#[cfg(feature = "concolic_mutation")] +impl Named for SimpleConcolicMutationalStage { + fn name(&self) -> &Cow<'static, str> { + &self.name + } +} + #[cfg(feature = "concolic_mutation")] impl Stage for SimpleConcolicMutationalStage where @@ -373,7 +395,7 @@ where EM: UsesState, Z: Evaluator, Z::Input: HasMutatorBytes, - Self::State: State + HasExecutions + HasCorpus + HasMetadata, + Self::State: State + HasExecutions + HasCorpus + HasMetadata + HasNamedMetadata, { #[inline] fn perform( @@ -396,11 +418,8 @@ where mutations }); - let post_restart_skip_cnt = - usize::try_from(self.restart_helper.execs_since_progress_start(state)?)?; - if let Some(mutations) = mutations { - for mutation in mutations.into_iter().skip(post_restart_skip_cnt) { + for mutation in mutations { let mut input_copy = state.current_input_cloned()?; for (index, new_byte) in mutation { input_copy.bytes_mut()[index] = new_byte; @@ -413,21 +432,34 @@ where } #[inline] - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.restart_progress_should_run(state) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // This is a deterministic stage + // Once it failed, then don't retry, + // It will just fail again + StdRestartHelper::no_retry(state, &self.name) } #[inline] - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - self.restart_helper.clear_restart_progress(state) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } #[cfg(feature = "concolic_mutation")] -impl Default for SimpleConcolicMutationalStage { - fn default() -> Self { +impl SimpleConcolicMutationalStage { + #[must_use] + /// Construct this stage + pub fn new() -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = SIMPLE_CONCOLIC_MUTATIONAL_ID; + SIMPLE_CONCOLIC_MUTATIONAL_ID += 1; + ret + }; Self { - restart_helper: ExecutionCountRestartHelper::default(), + name: Cow::Owned( + SIMPLE_CONCOLIC_MUTATIONAL_NAME.to_owned() + ":" + stage_id.to_string().as_str(), + ), phantom: PhantomData, } } diff --git a/libafl/src/stages/dump.rs b/libafl/src/stages/dump.rs index a099df759b..5e80591348 100644 --- a/libafl/src/stages/dump.rs +++ b/libafl/src/stages/dump.rs @@ -115,13 +115,13 @@ where } #[inline] - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { + fn should_restart(&mut self, _state: &mut Self::State) -> Result { // Not executing the target, so restart safety is not needed Ok(true) } #[inline] - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { // Not executing the target, so restart safety is not needed Ok(()) } diff --git a/libafl/src/stages/generalization.rs b/libafl/src/stages/generalization.rs index 44240dd1f0..d1d2741519 100644 --- a/libafl/src/stages/generalization.rs +++ b/libafl/src/stages/generalization.rs @@ -1,6 +1,9 @@ //! The tracing stage can trace the target and enrich a [`crate::corpus::Testcase`] with metadata, for example for `CmpLog`. -use alloc::{borrow::Cow, vec::Vec}; +use alloc::{ + borrow::{Cow, ToOwned}, + vec::Vec, +}; use core::{fmt::Debug, marker::PhantomData}; use libafl_bolts::{ @@ -16,7 +19,7 @@ use crate::{ mark_feature_time, observers::{CanTrack, MapObserver, ObserversTuple}, require_novelties_tracking, - stages::{RetryRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, start_timer, state::{HasCorpus, HasExecutions, UsesState}, Error, HasMetadata, HasNamedMetadata, @@ -40,9 +43,13 @@ fn find_next_char(list: &[Option], mut idx: usize, ch: u8) -> usize { idx } +/// The name for generalization stage +pub static GENERALIZATION_STAGE_NAME: &str = "generalization"; + /// A stage that runs a tracer executor #[derive(Clone, Debug)] pub struct GeneralizationStage { + name: Cow<'static, str>, map_observer_handle: Handle, #[allow(clippy::type_complexity)] phantom: PhantomData<(EM, O, OT, Z)>, @@ -50,8 +57,7 @@ pub struct GeneralizationStage { impl Named for GeneralizationStage { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed("GeneralizationStage"); - &NAME + &self.name } } @@ -320,15 +326,15 @@ where } #[inline] - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { + fn should_restart(&mut self, state: &mut Self::State) -> Result { // TODO: We need to be able to resume better if something crashes or times out - RetryRestartHelper::restart_progress_should_run(state, self, 3) + StdRestartHelper::should_restart(state, &self.name, 3) } #[inline] - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { // TODO: We need to be able to resume better if something crashes or times out - RetryRestartHelper::clear_restart_progress(state, self) + StdRestartHelper::clear_progress(state, &self.name) } } @@ -345,7 +351,11 @@ where #[must_use] pub fn new(map_observer: &C) -> Self { require_novelties_tracking!("GeneralizationStage", C); + let name = map_observer.name().clone(); Self { + name: Cow::Owned( + GENERALIZATION_STAGE_NAME.to_owned() + ":" + name.into_owned().as_str(), + ), map_observer_handle: map_observer.handle(), phantom: PhantomData, } diff --git a/libafl/src/stages/generation.rs b/libafl/src/stages/generation.rs index 46f7200c10..189b013ecc 100644 --- a/libafl/src/stages/generation.rs +++ b/libafl/src/stages/generation.rs @@ -54,11 +54,13 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { + fn should_restart(&mut self, _state: &mut Self::State) -> Result { + // It's a random generation stage + // so you can restart for whatever times you want Ok(true) } - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { Ok(()) } } diff --git a/libafl/src/stages/logics.rs b/libafl/src/stages/logics.rs index cb1ff665f5..ab0dc8847f 100644 --- a/libafl/src/stages/logics.rs +++ b/libafl/src/stages/logics.rs @@ -10,10 +10,10 @@ use crate::{ /// Progress for nested stages. This merely enters/exits the inner stage's scope. #[derive(Debug)] -pub struct NestedStageRestartHelper; +pub struct NestedStageStdRestartHelper; -impl NestedStageRestartHelper { - fn restart_progress_should_run(state: &mut S, _stage: &ST) -> Result +impl NestedStageStdRestartHelper { + fn should_restart(state: &mut S, _stage: &ST) -> Result where S: HasNestedStageStatus, { @@ -21,7 +21,7 @@ impl NestedStageRestartHelper { Ok(true) } - fn clear_restart_progress(state: &mut S, _stage: &ST) -> Result<(), Error> + fn clear_progress(state: &mut S, _stage: &ST) -> Result<(), Error> where S: HasNestedStageStatus, { @@ -70,12 +70,12 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - NestedStageRestartHelper::restart_progress_should_run(state, self) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + NestedStageStdRestartHelper::should_restart(state, self) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - NestedStageRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + NestedStageStdRestartHelper::clear_progress(state, self) } } @@ -134,12 +134,12 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - NestedStageRestartHelper::restart_progress_should_run(state, self) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + NestedStageStdRestartHelper::should_restart(state, self) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - NestedStageRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + NestedStageStdRestartHelper::clear_progress(state, self) } } @@ -219,12 +219,12 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - NestedStageRestartHelper::restart_progress_should_run(state, self) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + NestedStageStdRestartHelper::should_restart(state, self) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - NestedStageRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + NestedStageStdRestartHelper::clear_progress(state, self) } } @@ -280,12 +280,12 @@ where } } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - NestedStageRestartHelper::restart_progress_should_run(state, self) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + NestedStageStdRestartHelper::should_restart(state, self) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - NestedStageRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + NestedStageStdRestartHelper::clear_progress(state, self) } } diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 670b6c22bd..71d0c11dbf 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -4,7 +4,12 @@ A well-known [`Stage`], for example, is the mutational stage, running multiple [ Other stages may enrich [`crate::corpus::Testcase`]s with metadata. */ -use alloc::{borrow::Cow, boxed::Box, vec::Vec}; +use alloc::{ + borrow::{Cow, ToOwned}, + boxed::Box, + string::ToString, + vec::Vec, +}; use core::{fmt, marker::PhantomData}; pub use calibrate::CalibrationStage; @@ -86,18 +91,19 @@ where /// 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_restart_progress`], all subsequent calls happen on restart. + /// 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 restart_progress_should_run(&mut self, state: &mut Self::State) -> Result; + fn should_restart(&mut self, state: &mut Self::State) -> Result; /// Clear the current status tracking of the associated stage - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error>; + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error>; /// Run the stage. /// - /// Before a call to perform, [`Stage::restart_progress_should_run`] will be (must be!) called. - /// After returning (so non-target crash or timeout in a restarting case), [`Stage::clear_restart_progress`] gets called. + /// 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. /// A call to [`Stage::perform_restartable`] will do these things implicitly. + /// DON'T call this function directly except from `preform_restartable` !! fn perform( &mut self, fuzzer: &mut Z, @@ -106,7 +112,7 @@ where manager: &mut EM, ) -> Result<(), Error>; - /// Run the stage, calling [`Stage::restart_progress_should_run`] and [`Stage::clear_restart_progress`] appropriately + /// Run the stage, calling [`Stage::should_restart`] and [`Stage::clear_progress`] appropriately fn perform_restartable( &mut self, fuzzer: &mut Z, @@ -114,10 +120,10 @@ where state: &mut Self::State, manager: &mut EM, ) -> Result<(), Error> { - if self.restart_progress_should_run(state)? { + if self.should_restart(state)? { self.perform(fuzzer, executor, state, manager)?; } - self.clear_restart_progress(state) + self.clear_progress(state) } } @@ -283,9 +289,14 @@ where } } +static mut CLOSURE_STAGE_ID: usize = 0; +/// The name for closure stage +pub static CLOSURE_STAGE_NAME: &str = "closure"; + /// A [`Stage`] that will call a closure #[derive(Debug)] pub struct ClosureStage { + name: Cow<'static, str>, closure: CB, phantom: PhantomData<(E, EM, Z)>, } @@ -299,8 +310,7 @@ where impl Named for ClosureStage { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed(""); - &NAME + &self.name } } @@ -323,14 +333,15 @@ where } #[inline] - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - // Make sure we don't get stuck crashing on a single closure - RetryRestartHelper::restart_progress_should_run(state, self, 3) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // There's no restart safety in the content of the closure. + // don't restart + StdRestartHelper::no_retry(state, &self.name) } #[inline] - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } @@ -339,28 +350,25 @@ impl ClosureStage { /// Create a new [`ClosureStage`] #[must_use] pub fn new(closure: CB) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = CLOSURE_STAGE_ID; + CLOSURE_STAGE_ID += 1; + ret + }; Self { + name: Cow::Owned(CLOSURE_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_ref()), closure, phantom: PhantomData, } } } -impl From for ClosureStage -where - CB: FnMut(&mut Z, &mut E, &mut ::State, &mut EM) -> Result<(), Error>, - E: UsesState, -{ - #[must_use] - fn from(closure: CB) -> Self { - Self::new(closure) - } -} - /// Allows us to use a [`push::PushStage`] as a normal [`Stage`] #[allow(clippy::type_complexity)] #[derive(Debug)] pub struct PushStageAdapter { + name: Cow<'static, str>, push_stage: PS, phantom: PhantomData<(CS, EM, OT, Z)>, } @@ -370,12 +378,25 @@ impl PushStageAdapter { /// to be used as a normal [`Stage`] #[must_use] pub fn new(push_stage: PS) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = PUSH_STAGE_ADAPTER_ID; + PUSH_STAGE_ADAPTER_ID += 1; + ret + }; Self { + name: Cow::Owned( + PUSH_STAGE_ADAPTER_NAME.to_owned() + ":" + stage_id.to_string().as_str(), + ), push_stage, phantom: PhantomData, } } } +/// The unique counter for this stage +static mut PUSH_STAGE_ADAPTER_ID: usize = 0; +/// The name for push stage adapter +pub static PUSH_STAGE_ADAPTER_NAME: &str = "pushstageadapter"; impl UsesState for PushStageAdapter where @@ -384,11 +405,23 @@ where type State = CS::State; } +impl Named for PushStageAdapter { + #[must_use] + fn name(&self) -> &Cow<'static, str> { + &self.name + } +} + impl Stage for PushStageAdapter where CS: Scheduler, - Self::State: - HasExecutions + HasMetadata + HasRand + HasCorpus + HasLastReportTime + HasCurrentCorpusId, + Self::State: HasExecutions + + HasRand + + HasCorpus + + HasLastReportTime + + HasCurrentCorpusId + + HasNamedMetadata + + HasMetadata, E: Executor + HasObservers, EM: EventFirer + EventRestarter @@ -446,48 +479,51 @@ where } #[inline] - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { + fn should_restart(&mut self, state: &mut Self::State) -> Result { // TODO: Proper restart handling - call post_exec at the right time, etc... - Ok(true) + StdRestartHelper::no_retry(state, &self.name) } #[inline] - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { - Ok(()) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } /// Progress which permits a fixed amount of resumes per round of fuzzing. If this amount is ever /// exceeded, the input will no longer be executed by this stage. #[derive(Clone, Deserialize, Serialize, Debug)] -pub struct RetryRestartHelper { +pub struct StdRestartHelper { tries_remaining: Option, skipped: HashSet, } -impl_serdeany!(RetryRestartHelper); +impl_serdeany!(StdRestartHelper); + +impl StdRestartHelper { + /// Don't allow restart + pub fn no_retry(state: &mut S, name: &str) -> Result + where + S: HasNamedMetadata + HasCurrentCorpusId, + { + Self::should_restart(state, name, 1) + } -impl RetryRestartHelper { /// Initializes (or counts down in) the progress helper, giving it the amount of max retries /// /// Returns `true` if the stage should run - pub fn restart_progress_should_run( - state: &mut S, - stage: &ST, - max_retries: usize, - ) -> Result + pub fn should_restart(state: &mut S, name: &str, max_retries: usize) -> Result where S: HasNamedMetadata + HasCurrentCorpusId, - ST: Named, { let corpus_id = state.current_corpus_id()?.ok_or_else(|| { Error::illegal_state( - "No current_corpus_id set in State, but called RetryRestartHelper::should_skip", + "No current_corpus_id set in State, but called StdRestartHelper::should_skip", ) })?; let initial_tries_remaining = max_retries + 1; - let metadata = state.named_metadata_or_insert_with(stage.name(), || Self { + let metadata = state.named_metadata_or_insert_with(name, || Self { tries_remaining: Some(initial_tries_remaining), skipped: HashSet::new(), }); @@ -515,14 +551,11 @@ impl RetryRestartHelper { } /// Clears the progress - pub fn clear_restart_progress(state: &mut S, stage: &ST) -> Result<(), Error> + pub fn clear_progress(state: &mut S, name: &str) -> Result<(), Error> where S: HasNamedMetadata, - ST: Named, { - state - .named_metadata_mut::(stage.name())? - .tries_remaining = None; + state.named_metadata_mut::(name)?.tries_remaining = None; Ok(()) } } @@ -597,15 +630,15 @@ impl ExecutionCountRestartHelper { } /// The execs done since start of this [`Stage`]/helper - pub fn execs_since_progress_start(&mut self, state: &mut S) -> Result + pub fn execs_since_progress_start(&mut self, state: &mut S, name: &str) -> Result where - S: HasMetadata + HasExecutions, + S: HasNamedMetadata + HasExecutions, { let started_at_execs = if let Some(started_at_execs) = self.started_at_execs { started_at_execs } else { state - .metadata::() + .named_metadata::(name) .map(|x| { self.started_at_execs = Some(x.started_at_execs); x.started_at_execs @@ -620,26 +653,27 @@ impl ExecutionCountRestartHelper { } /// Initialize progress for the stage this wrapper wraps. - pub fn restart_progress_should_run(&mut self, state: &mut S) -> Result + pub fn should_restart(&mut self, state: &mut S, name: &str) -> Result where - S: HasMetadata + HasExecutions, + S: HasNamedMetadata + HasExecutions, { let executions = *state.executions(); - let metadata = state.metadata_or_insert_with(|| ExecutionCountRestartHelperMetadata { - started_at_execs: executions, - }); + let metadata = + state.named_metadata_or_insert_with(name, || ExecutionCountRestartHelperMetadata { + started_at_execs: executions, + }); self.started_at_execs = Some(metadata.started_at_execs); Ok(true) } /// Clear progress for the stage this wrapper wraps. - pub fn clear_restart_progress(&mut self, state: &mut S) -> Result<(), Error> + pub fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> where S: HasMetadata, { self.started_at_execs = None; let _metadata = state.remove_metadata::(); - debug_assert!(_metadata.is_some(), "Called clear_restart_progress, but restart_progress_should_run was not called before (or did mutational stages get nested?)"); + debug_assert!(_metadata.is_some(), "Called clear_progress, but should_restart was not called before (or did mutational stages get nested?)"); Ok(()) } } @@ -655,7 +689,7 @@ pub mod test { use crate::{ corpus::{Corpus, HasCurrentCorpusId, Testcase}, inputs::NopInput, - stages::{RetryRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, state::{test::test_std_state, HasCorpus, State, UsesState}, HasMetadata, }; @@ -674,7 +708,7 @@ pub mod test { impl TestProgress { #[allow(clippy::unnecessary_wraps)] - fn restart_progress_should_run(state: &mut S, _stage: &ST) -> Result + fn should_restart(state: &mut S, _stage: &ST) -> Result where S: HasMetadata, { @@ -690,7 +724,7 @@ pub mod test { Ok(true) } - fn clear_restart_progress(state: &mut S, _stage: &ST) -> Result<(), Error> + fn clear_progress(state: &mut S, _stage: &ST) -> Result<(), Error> where S: HasMetadata, { @@ -727,12 +761,12 @@ pub mod test { Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - TestProgress::restart_progress_should_run(state, self) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + TestProgress::should_restart(state, self) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - TestProgress::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + TestProgress::clear_progress(state, self) } } @@ -742,7 +776,7 @@ pub mod test { // No concurrency per testcase #[cfg(any(not(feature = "serdeany_autoreg"), miri))] unsafe { - RetryRestartHelper::register(); + StdRestartHelper::register(); } struct StageWithOneTry; @@ -763,44 +797,58 @@ pub mod test { for _ in 0..10 { // used normally, no retries means we never skip - assert!(RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 1 + assert!(StdRestartHelper::should_restart( + &mut state, + stage.name(), + 1 )?); - RetryRestartHelper::clear_restart_progress(&mut state, &stage)?; + StdRestartHelper::clear_progress(&mut state, stage.name())?; } for _ in 0..10 { // used normally, only one retry means we never skip - assert!(RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 2 + assert!(StdRestartHelper::should_restart( + &mut state, + stage.name(), + 2 )?); - assert!(RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 2 + assert!(StdRestartHelper::should_restart( + &mut state, + stage.name(), + 2 )?); - RetryRestartHelper::clear_restart_progress(&mut state, &stage)?; + StdRestartHelper::clear_progress(&mut state, stage.name())?; } - assert!(RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 2 + assert!(StdRestartHelper::should_restart( + &mut state, + stage.name(), + 2 )?); // task failed, let's resume // we still have one more try! - assert!(RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 2 + assert!(StdRestartHelper::should_restart( + &mut state, + stage.name(), + 2 )?); // task failed, let's resume // out of retries, so now we skip - assert!(!RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 2 + assert!(!StdRestartHelper::should_restart( + &mut state, + stage.name(), + 2 )?); - RetryRestartHelper::clear_restart_progress(&mut state, &stage)?; + StdRestartHelper::clear_progress(&mut state, stage.name())?; // we previously exhausted this testcase's retries, so we skip - assert!(!RetryRestartHelper::restart_progress_should_run( - &mut state, &stage, 2 + assert!(!StdRestartHelper::should_restart( + &mut state, + stage.name(), + 2 )?); - RetryRestartHelper::clear_restart_progress(&mut state, &stage)?; + StdRestartHelper::clear_progress(&mut state, stage.name())?; Ok(()) } diff --git a/libafl/src/stages/mutational.rs b/libafl/src/stages/mutational.rs index 856e413296..f66257d5d6 100644 --- a/libafl/src/stages/mutational.rs +++ b/libafl/src/stages/mutational.rs @@ -1,7 +1,10 @@ //| The [`MutationalStage`] is the default stage used during fuzzing. //! For the current input, it will perform a range of random mutations, and then run them in the executor. -use alloc::borrow::Cow; +use alloc::{ + borrow::{Cow, ToOwned}, + string::ToString, +}; use core::marker::PhantomData; use libafl_bolts::{rands::Rand, Named}; @@ -12,7 +15,7 @@ use crate::{ inputs::Input, mark_feature_time, mutators::{MultiMutator, MutationResult, Mutator}, - stages::{ExecutionCountRestartHelper, RetryRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, start_timer, state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, UsesState}, Error, HasMetadata, HasNamedMetadata, @@ -94,9 +97,6 @@ where /// Gets the number of iterations this mutator should run for. fn iterations(&self, state: &mut Self::State) -> Result; - /// Gets the number of executions this mutator already did since it got first called in this fuzz round. - fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result; - /// Runs this (mutational) stage for the given testcase #[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely... fn perform_mutational( @@ -155,12 +155,12 @@ pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: usize = 128; /// The default mutational stage #[derive(Clone, Debug)] pub struct StdMutationalStage { + /// The name + name: Cow<'static, str>, /// The mutator(s) to use mutator: M, /// The maximum amount of iterations we should do each round max_iterations: usize, - /// The progress helper for this mutational stage - restart_helper: ExecutionCountRestartHelper, #[allow(clippy::type_complexity)] phantom: PhantomData<(E, EM, I, Z)>, } @@ -171,7 +171,7 @@ where EM: UsesState, M: Mutator, Z: Evaluator, - Self::State: HasCorpus + HasRand + HasExecutions + HasMetadata, + Self::State: HasCorpus + HasRand + HasExecutions + HasMetadata + HasNamedMetadata, I: MutatedTransform + Clone, { /// The mutator, added to this stage @@ -190,12 +190,13 @@ where fn iterations(&self, state: &mut Self::State) -> Result { Ok(1 + state.rand_mut().below(self.max_iterations)) } - - fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.execs_since_progress_start(state) - } } +/// The unique id for mutational stage +static mut MUTATIONAL_STAGE_ID: usize = 0; +/// The name for mutational stage +pub static MUTATIONAL_STAGE_NAME: &str = "mutational"; + impl UsesState for StdMutationalStage where Z: UsesState, @@ -203,13 +204,19 @@ where type State = Z::State; } +impl Named for StdMutationalStage { + fn name(&self) -> &Cow<'static, str> { + &self.name + } +} + impl Stage for StdMutationalStage where E: UsesState, EM: UsesState, M: Mutator, Z: Evaluator, - Self::State: HasCorpus + HasRand + HasMetadata + HasExecutions, + Self::State: HasCorpus + HasRand + HasMetadata + HasExecutions + HasNamedMetadata, I: MutatedTransform + Clone, { #[inline] @@ -229,14 +236,12 @@ where ret } - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { - Ok(true) - // self.restart_helper.restart_progress_should_run(state) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + StdRestartHelper::should_restart(state, &self.name, 3) } - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { - Ok(()) - // self.restart_helper.clear_restart_progress(state) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } @@ -274,10 +279,18 @@ where /// Creates a new transforming mutational stage with the given max iterations pub fn transforming_with_max_iterations(mutator: M, max_iterations: usize) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = MUTATIONAL_STAGE_ID; + MUTATIONAL_STAGE_ID += 1; + ret + }; Self { + name: Cow::Owned( + MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(), + ), mutator, max_iterations, - restart_helper: ExecutionCountRestartHelper::default(), phantom: PhantomData, } } @@ -286,11 +299,17 @@ where /// A mutational stage that operates on multiple inputs, as returned by [`MultiMutator::multi_mutate`]. #[derive(Clone, Debug)] pub struct MultiMutationalStage { + name: Cow<'static, str>, mutator: M, #[allow(clippy::type_complexity)] phantom: PhantomData<(E, EM, I, Z)>, } +/// The unique id for multi mutational stage +static mut MULTI_MUTATIONAL_STAGE_ID: usize = 0; +/// The name for multi mutational stage +pub static MULTI_MUTATIONAL_STAGE_NAME: &str = "multimutational"; + impl UsesState for MultiMutationalStage where Z: UsesState, @@ -300,8 +319,7 @@ where impl Named for MultiMutationalStage { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed("MultiMutational"); - &NAME + &self.name } } @@ -315,15 +333,14 @@ where I: MutatedTransform + Clone, { #[inline] - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - // TODO: add proper crash/timeout handling - // For now, Make sure we don't get stuck crashing on a single testcase - RetryRestartHelper::restart_progress_should_run(state, self, 3) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // Make sure we don't get stuck crashing on a single testcase + StdRestartHelper::should_restart(state, &self.name, 3) } #[inline] - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } #[inline] @@ -370,7 +387,16 @@ where impl MultiMutationalStage { /// Creates a new transforming mutational stage pub fn transforming(mutator: M) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = MULTI_MUTATIONAL_STAGE_ID; + MULTI_MUTATIONAL_STAGE_ID += 1; + ret + }; Self { + name: Cow::Owned( + MULTI_MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(), + ), mutator, phantom: PhantomData, } diff --git a/libafl/src/stages/power.rs b/libafl/src/stages/power.rs index 41c3abea40..ab729b369e 100644 --- a/libafl/src/stages/power.rs +++ b/libafl/src/stages/power.rs @@ -1,6 +1,9 @@ //! The power schedules. This stage should be invoked after the calibration stage. -use alloc::borrow::Cow; +use alloc::{ + borrow::{Cow, ToOwned}, + string::ToString, +}; use core::{fmt::Debug, marker::PhantomData}; use libafl_bolts::Named; @@ -10,10 +13,13 @@ use crate::{ fuzzer::Evaluator, mutators::Mutator, schedulers::{testcase_score::CorpusPowerTestcaseScore, TestcaseScore}, - stages::{mutational::MutatedTransform, ExecutionCountRestartHelper, MutationalStage, Stage}, + stages::{mutational::MutatedTransform, MutationalStage, Stage, StdRestartHelper}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, UsesState}, - Error, HasMetadata, + Error, HasMetadata, HasNamedMetadata, }; + +/// The unique id for this stage +static mut POWER_MUTATIONAL_STAGE_ID: usize = 0; /// Default name for `PowerMutationalStage`; derived from AFL++ pub const POWER_MUTATIONAL_STAGE_NAME: &str = "power"; /// The mutational stage using power schedules @@ -22,8 +28,6 @@ pub struct PowerMutationalStage { name: Cow<'static, str>, /// The mutators we use mutator: M, - /// Helper for restarts - restart_helper: ExecutionCountRestartHelper, #[allow(clippy::type_complexity)] phantom: PhantomData<(E, F, EM, I, Z)>, } @@ -47,7 +51,7 @@ where EM: UsesState, F: TestcaseScore, M: Mutator, - Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions, + Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions + HasNamedMetadata, Z: Evaluator, I: MutatedTransform + Clone, { @@ -72,10 +76,6 @@ where Ok(score) } - - fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.execs_since_progress_start(state) - } } impl Stage for PowerMutationalStage @@ -84,7 +84,7 @@ where EM: UsesState, F: TestcaseScore, M: Mutator, - Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions, + Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions + HasNamedMetadata, Z: Evaluator, I: MutatedTransform + Clone, { @@ -101,14 +101,13 @@ where ret } - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { - Ok(true) - // self.restart_helper.restart_progress_should_run(state) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // Make sure we don't get stuck crashing on a single testcase + StdRestartHelper::should_restart(state, &self.name, 3) } - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { - Ok(()) - // self.restart_helper.clear_restart_progress(state) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } @@ -123,11 +122,18 @@ where { /// Creates a new [`PowerMutationalStage`] pub fn new(mutator: M) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = POWER_MUTATIONAL_STAGE_ID; + POWER_MUTATIONAL_STAGE_ID += 1; + ret + }; Self { - name: Cow::Borrowed(POWER_MUTATIONAL_STAGE_NAME), + name: Cow::Owned( + POWER_MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(), + ), mutator, phantom: PhantomData, - restart_helper: ExecutionCountRestartHelper::default(), } } } diff --git a/libafl/src/stages/stats.rs b/libafl/src/stages/stats.rs index 73ea96f5c3..fa3db5692d 100644 --- a/libafl/src/stages/stats.rs +++ b/libafl/src/stages/stats.rs @@ -126,13 +126,13 @@ where } #[inline] - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { + fn should_restart(&mut self, _state: &mut Self::State) -> Result { // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything Ok(true) } #[inline] - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything Ok(()) } diff --git a/libafl/src/stages/sync.rs b/libafl/src/stages/sync.rs index 6b6be129d9..65d899a976 100644 --- a/libafl/src/stages/sync.rs +++ b/libafl/src/stages/sync.rs @@ -1,6 +1,6 @@ //! The [`SyncFromDiskStage`] is a stage that imports inputs from disk for e.g. sync with AFL -use alloc::borrow::Cow; +use alloc::borrow::{Cow, ToOwned}; use core::marker::PhantomData; use std::{ fs, @@ -20,7 +20,7 @@ use crate::{ executors::{Executor, ExitKind, HasObservers}, fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor}, inputs::{Input, InputConverter, UsesInput}, - stages::{RetryRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, state::{HasCorpus, HasExecutions, HasRand, State, UsesState}, Error, HasMetadata, HasNamedMetadata, }; @@ -149,24 +149,24 @@ where } #[inline] - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { + fn should_restart(&mut self, state: &mut Self::State) -> 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 - RetryRestartHelper::restart_progress_should_run(state, self, 3) + StdRestartHelper::no_retry(state, &self.name) } #[inline] - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } impl SyncFromDiskStage { /// Creates a new [`SyncFromDiskStage`] #[must_use] - pub fn new(sync_dir: PathBuf, load_callback: CB) -> Self { + pub fn new(sync_dir: PathBuf, load_callback: CB, name: &str) -> Self { Self { - name: Cow::Borrowed(SYNC_FROM_DISK_STAGE_NAME), + name: Cow::Owned(SYNC_FROM_DISK_STAGE_NAME.to_owned() + ":" + name), phantom: PhantomData, sync_dir, load_callback, @@ -367,13 +367,13 @@ where } #[inline] - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { + fn should_restart(&mut self, _state: &mut Self::State) -> Result { // No restart handling needed - does not execute the target. Ok(true) } #[inline] - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { // Not needed - does not execute the target. Ok(()) } diff --git a/libafl/src/stages/tmin.rs b/libafl/src/stages/tmin.rs index 6dd7799ee1..ab250a3092 100644 --- a/libafl/src/stages/tmin.rs +++ b/libafl/src/stages/tmin.rs @@ -1,6 +1,9 @@ //! The [`TMinMutationalStage`] is a stage which will attempt to minimize corpus entries. -use alloc::borrow::Cow; +use alloc::{ + borrow::{Cow, ToOwned}, + string::ToString, +}; use core::{borrow::BorrowMut, fmt::Debug, hash::Hash, marker::PhantomData}; use ahash::RandomState; @@ -29,7 +32,8 @@ use crate::{ state::{ HasCorpus, HasCurrentTestcase, HasExecutions, HasMaxSize, HasSolutions, State, UsesState, }, - Error, ExecutesInput, ExecutionProcessor, HasFeedback, HasMetadata, HasScheduler, + Error, ExecutesInput, ExecutionProcessor, HasFeedback, HasMetadata, HasNamedMetadata, + HasScheduler, }; #[cfg(feature = "introspection")] use crate::{monitors::PerfFeature, state::HasClientPerfMonitor}; @@ -79,8 +83,14 @@ where let orig_max_size = state.max_size(); // basically copy-pasted from mutational.rs - let num = self.iterations(state)? - - usize::try_from(self.execs_since_progress_start(state)?).unwrap(); + let num = self + .iterations(state)? + .saturating_sub(usize::try_from(self.execs_since_progress_start(state)?)?); + + // If num is negative, then quit. + if num == 0 { + return Ok(()); + } start_timer!(state); let transformed = @@ -203,6 +213,8 @@ where /// The default corpus entry minimising mutational stage #[derive(Clone, Debug)] pub struct StdTMinMutationalStage { + /// The name + name: Cow<'static, str>, /// The mutator(s) this stage uses mutator: M, /// The factory @@ -231,16 +243,17 @@ where FF: FeedbackFactory, F: Feedback, Self::Input: MutatedTransform + Clone + HasLen + Hash, - Self::State: HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize, + Self::State: + HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize + HasNamedMetadata, M: Mutator, IP: MutatedTransformPost + Clone, { - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.restart_progress_should_run(state) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + self.restart_helper.should_restart(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - self.restart_helper.clear_restart_progress(state) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + self.restart_helper.clear_progress(state) } fn perform( @@ -270,6 +283,17 @@ where } } +impl Named for StdTMinMutationalStage { + fn name(&self) -> &Cow<'static, str> { + &self.name + } +} + +/// The counter for giving this stage unique id +static mut TMIN_STAGE_ID: usize = 0; +/// The name for tmin stage +pub static TMIN_STAGE_NAME: &str = "tmin"; + impl TMinMutationalStage for StdTMinMutationalStage where @@ -280,7 +304,8 @@ where FF: FeedbackFactory, F: Feedback, Self::Input: MutatedTransform + Clone + HasLen + Hash, - Self::State: HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize, + Self::State: + HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize + HasNamedMetadata, M: Mutator, IP: MutatedTransformPost + Clone, { @@ -302,14 +327,22 @@ where } fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.execs_since_progress_start(state) + self.restart_helper + .execs_since_progress_start(state, &self.name) } } impl StdTMinMutationalStage { /// Creates a new minimizing mutational stage that will minimize provided corpus entries pub fn new(mutator: M, factory: FF, runs: usize) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = TMIN_STAGE_ID; + TMIN_STAGE_ID += 1; + ret + }; Self { + name: Cow::Owned(TMIN_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str()), mutator, factory, runs, diff --git a/libafl/src/stages/tracing.rs b/libafl/src/stages/tracing.rs index 4dcd0103c4..6487ddca7b 100644 --- a/libafl/src/stages/tracing.rs +++ b/libafl/src/stages/tracing.rs @@ -1,6 +1,9 @@ //! The tracing stage can trace the target and enrich a testcase with metadata, for example for `CmpLog`. -use alloc::borrow::Cow; +use alloc::{ + borrow::{Cow, ToOwned}, + string::ToString, +}; use core::{fmt::Debug, marker::PhantomData}; use libafl_bolts::Named; @@ -9,7 +12,7 @@ use crate::{ executors::{Executor, HasObservers, ShadowExecutor}, mark_feature_time, observers::ObserversTuple, - stages::{RetryRestartHelper, Stage}, + stages::{Stage, StdRestartHelper}, start_timer, state::{HasCorpus, HasCurrentTestcase, HasExecutions, State, UsesState}, Error, HasNamedMetadata, @@ -20,8 +23,8 @@ use crate::{monitors::PerfFeature, state::HasClientPerfMonitor}; /// A stage that runs a tracer executor #[derive(Clone, Debug)] pub struct TracingStage { + name: Cow<'static, str>, tracer_executor: TE, - max_retries: usize, #[allow(clippy::type_complexity)] phantom: PhantomData<(EM, TE, Z)>, } @@ -42,7 +45,7 @@ where { #[allow(rustdoc::broken_intra_doc_links)] /// Perform tracing on the given `CorpusId`. Useful for if wrapping [`TracingStage`] with your - /// own stage and you need to manage [`super::NestedStageRestartHelper`] differently + /// own stage and you need to manage [`super::NestedStageStdRestartHelper`] differently /// see [`super::ConcolicTracingStage`]'s implementation as an example of usage. pub fn trace( &mut self, @@ -96,40 +99,43 @@ where self.trace(fuzzer, state, manager) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + StdRestartHelper::no_retry(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } impl Named for TracingStage { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed("TracingStage"); - &NAME + &self.name } } +/// The counter for giving this stage unique id +static mut TRACING_STAGE_ID: usize = 0; +/// The name for tracing stage +pub static TRACING_STAGE_NAME: &str = "tracing"; + impl TracingStage { /// Creates a new default stage pub fn new(tracer_executor: TE) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = TRACING_STAGE_ID; + TRACING_STAGE_ID += 1; + ret + }; + Self { + name: Cow::Owned(TRACING_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_ref()), tracer_executor, - max_retries: 10, phantom: PhantomData, } } - /// Specify how many times that this stage will try again to trace the input before giving up - /// and not processing the input again. 0 retries means that the trace will be tried only once. - #[must_use] - pub fn with_retries(mut self, retries: usize) -> Self { - self.max_retries = retries; - self - } - /// Gets the underlying tracer executor pub fn executor(&self) -> &TE { &self.tracer_executor @@ -144,7 +150,7 @@ impl TracingStage { /// A stage that runs the shadow executor using also the shadow observers #[derive(Clone, Debug)] pub struct ShadowTracingStage { - max_retries: usize, + name: Cow<'static, str>, #[allow(clippy::type_complexity)] phantom: PhantomData<(E, EM, SOT, Z)>, } @@ -155,14 +161,17 @@ where { type State = E::State; } +/// The counter for giving this stage unique id +static mut SHADOW_TRACING_STAGE_ID: usize = 0; +/// Name for shadow tracing stage +pub static SHADOW_TRACING_STAGE_NAME: &str = "shadow"; impl Named for ShadowTracingStage where E: UsesState, { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed("ShadowTracingStage"); - &NAME + &self.name } } @@ -210,12 +219,12 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + StdRestartHelper::no_retry(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - RetryRestartHelper::clear_restart_progress(state, self) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + StdRestartHelper::clear_progress(state, &self.name) } } @@ -229,17 +238,17 @@ where { /// Creates a new default stage pub fn new(_executor: &mut ShadowExecutor) -> Self { + // unsafe but impossible that you create two threads both instantiating this instance + let stage_id = unsafe { + let ret = SHADOW_TRACING_STAGE_ID; + SHADOW_TRACING_STAGE_ID += 1; + ret + }; Self { - max_retries: 10, + name: Cow::Owned( + SHADOW_TRACING_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(), + ), phantom: PhantomData, } } - - /// Specify how many times that this stage will try again to trace the input before giving up - /// and not processing the input again. 0 retries means that the trace will be tried only once. - #[must_use] - pub fn with_retries(mut self, retries: usize) -> Self { - self.max_retries = retries; - self - } } diff --git a/libafl/src/stages/tuneable.rs b/libafl/src/stages/tuneable.rs index b5d57d5f3a..749e74095a 100644 --- a/libafl/src/stages/tuneable.rs +++ b/libafl/src/stages/tuneable.rs @@ -248,10 +248,6 @@ where 1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS), ) } - - fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.execs_since_progress_start(state) - } } impl UsesState for TuneableMutationalStage @@ -287,12 +283,12 @@ where ret } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - self.restart_helper.restart_progress_should_run(state) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + self.restart_helper.should_restart(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { - self.restart_helper.clear_restart_progress(state) + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + self.restart_helper.clear_progress(state) } } @@ -303,9 +299,17 @@ where M: Mutator::State>, Z: Evaluator, ::State: - HasCorpus + HasRand + HasNamedMetadata + HasMetadata + HasExecutions, + HasCorpus + HasRand + HasNamedMetadata + HasExecutions + HasMetadata, I: MutatedTransform::State> + Clone, { + fn execs_since_progress_start( + &mut self, + state: &mut ::State, + ) -> Result { + self.restart_helper + .execs_since_progress_start(state, &self.name) + } + /// Creates a new default tuneable mutational stage #[must_use] pub fn new(state: &mut ::State, mutator: M) -> Self { diff --git a/libafl/src/stages/unicode.rs b/libafl/src/stages/unicode.rs index 0a661ab306..98957a015c 100644 --- a/libafl/src/stages/unicode.rs +++ b/libafl/src/stages/unicode.rs @@ -127,13 +127,13 @@ where } #[inline] - fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result { + fn should_restart(&mut self, _state: &mut Self::State) -> Result { // Stage does not run the target. No reset helper needed. Ok(true) } #[inline] - fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { // Stage does not run the target. No reset helper needed. Ok(()) } diff --git a/libafl_bolts/src/lib.rs b/libafl_bolts/src/lib.rs index 18ff284c2a..1a9efaba89 100644 --- a/libafl_bolts/src/lib.rs +++ b/libafl_bolts/src/lib.rs @@ -31,7 +31,7 @@ clippy::module_name_repetitions, clippy::ptr_cast_constness, clippy::negative_feature_names, - clippy::too_many_lines, + clippy::too_many_lines )] #![cfg_attr(not(test), warn( missing_debug_implementations, diff --git a/libafl_bolts/src/llmp.rs b/libafl_bolts/src/llmp.rs index a45f12610a..d1581ec291 100644 --- a/libafl_bolts/src/llmp.rs +++ b/libafl_bolts/src/llmp.rs @@ -1289,7 +1289,9 @@ where log::debug!( "[{} - {:#x}] Send message with id {}", - self.id.0, self as *const Self as u64, mid + self.id.0, + self as *const Self as u64, + mid ); Ok(()) @@ -2369,13 +2371,11 @@ impl Brokers { loop { self.llmp_brokers.retain_mut(|broker| { if broker.is_shutting_down() { - broker.send_buf(LLMP_TAG_EXITING, &[]).expect( "Error when shutting down broker: Could not send LLMP_TAG_EXITING msg.", ); - return false - + return false; } if current_milliseconds() > end_time { diff --git a/libafl_targets/src/cmps/stages/aflpptracing.rs b/libafl_targets/src/cmps/stages/aflpptracing.rs index a248e6d760..8ae08b55ec 100644 --- a/libafl_targets/src/cmps/stages/aflpptracing.rs +++ b/libafl_targets/src/cmps/stages/aflpptracing.rs @@ -1,11 +1,11 @@ -use alloc::borrow::Cow; +use alloc::borrow::{Cow, ToOwned}; use core::marker::PhantomData; use libafl::{ executors::{Executor, HasObservers}, inputs::{BytesInput, UsesInput}, observers::ObserversTuple, - stages::{colorization::TaintMetadata, RetryRestartHelper, Stage}, + stages::{colorization::TaintMetadata, Stage, StdRestartHelper}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState}, Error, HasMetadata, HasNamedMetadata, }; @@ -22,11 +22,14 @@ pub struct AFLppCmplogTracingStage<'a, EM, TE, Z> where TE: UsesState, { + name: Cow<'static, str>, tracer_executor: TE, - cmplog_observer_handle: Option::State>>>, + cmplog_observer_handle: Handle::State>>, #[allow(clippy::type_complexity)] phantom: PhantomData<(EM, TE, Z)>, } +/// The name for aflpp tracing stage +pub static AFLPP_CMPLOG_TRACING_STAGE_NAME: &str = "aflpptracing"; impl UsesState for AFLppCmplogTracingStage<'_, EM, TE, Z> where @@ -40,8 +43,7 @@ where TE: UsesState, { fn name(&self) -> &Cow<'static, str> { - static NAME: Cow<'static, str> = Cow::Borrowed("AFLppCmplogTracingStage"); - &NAME + &self.name } } @@ -65,19 +67,17 @@ where // First run with the un-mutated input let unmutated_input = state.current_input_cloned()?; - if let Some(observer_handle) = &self.cmplog_observer_handle { - if let Some(ob) = self - .tracer_executor - .observers_mut() - .get_mut(observer_handle) - { - // This is not the original input, - // Set it to false - ob.set_original(true); - } - // I can't think of any use of this stage if you don't use AFLppCmpLogObserver - // but do nothing ofcourse + if let Some(ob) = self + .tracer_executor + .observers_mut() + .get_mut(&self.cmplog_observer_handle) + { + // This is not the original input, + // Set it to false + ob.set_original(true); } + // I can't think of any use of this stage if you don't use AFLppCmpLogObserver + // but do nothing ofcourse self.tracer_executor .observers_mut() @@ -97,19 +97,17 @@ where None => return Err(Error::unknown("No metadata found")), }; - if let Some(observer_handle) = &self.cmplog_observer_handle { - if let Some(ob) = self - .tracer_executor - .observers_mut() - .get_mut(observer_handle) - { - // This is not the original input, - // Set it to false - ob.set_original(false); - } - // I can't think of any use of this stage if you don't use AFLppCmpLogObserver - // but do nothing ofcourse + if let Some(ob) = self + .tracer_executor + .observers_mut() + .get_mut(&self.cmplog_observer_handle) + { + // This is not the original input, + // Set it to false + ob.set_original(false); } + // I can't think of any use of this stage if you don't use AFLppCmpLogObserver + // but do nothing ofcourse self.tracer_executor .observers_mut() @@ -126,14 +124,15 @@ where Ok(()) } - fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result { - // TODO: this may need better resumption? (Or is it always used with a forkserver?) - RetryRestartHelper::restart_progress_should_run(state, self, 3) + fn should_restart(&mut self, state: &mut Self::State) -> Result { + // Tracing stage is always deterministic + // don't restart + StdRestartHelper::no_retry(state, &self.name) } - fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { + fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { // TODO: this may need better resumption? (Or is it always used with a forkserver?) - RetryRestartHelper::clear_restart_progress(state, self) + StdRestartHelper::clear_progress(state, &self.name) } } @@ -141,22 +140,19 @@ impl<'a, EM, TE, Z> AFLppCmplogTracingStage<'a, EM, TE, Z> where TE: UsesState, { - /// Creates a new default stage - pub fn new(tracer_executor: TE) -> Self { - Self { - cmplog_observer_handle: None, - tracer_executor, - phantom: PhantomData, - } - } - /// With cmplog observer - pub fn with_cmplog_observer( + pub fn new( tracer_executor: TE, observer_handle: Handle>, ) -> Self { + let observer_name = observer_handle.name().clone(); Self { - cmplog_observer_handle: Some(observer_handle), + name: Cow::Owned( + AFLPP_CMPLOG_TRACING_STAGE_NAME.to_owned() + + ":" + + observer_name.into_owned().as_str(), + ), + cmplog_observer_handle: observer_handle, tracer_executor, phantom: PhantomData, } diff --git a/utils/multi_machine_generator/src/graph.rs b/utils/multi_machine_generator/src/graph.rs index 113e0321f1..0cfc59c264 100644 --- a/utils/multi_machine_generator/src/graph.rs +++ b/utils/multi_machine_generator/src/graph.rs @@ -1,23 +1,20 @@ use std::net::SocketAddr; -use petgraph::{Direction, Graph}; -use petgraph::graph::NodeIndex; + +use petgraph::{graph::NodeIndex, Direction, Graph}; /// A node of the network -#[derive(Debug, Clone)] -pub struct MultiMachineNode { - addr: SocketAddr -} +#[derive(Debug, Clone, Default)] +pub struct MultiMachineNode {} /// The tree pub struct MultiMachineTree { - pub graph: Graph + pub graph: Graph, } impl MultiMachineNode { - pub fn new(addr: SocketAddr) -> Self { - Self { - addr - } + #[must_use] + pub fn new() -> Self { + Self {} } } @@ -26,21 +23,19 @@ impl MultiMachineTree { /// /// /// - machines: machines to add. - /// - max_children_per_parent: each parent will have at most this amount of children + /// - `max_children_per_parent`: each parent will have at most this amount of children + #[must_use] pub fn generate(machines: &[SocketAddr], max_children_per_parent: u64) -> Self { let mut graph = Graph::::new(); let mut machines = Vec::from(machines); - let root = if let Some(root) = machines.pop() { - graph.add_node(MultiMachineNode::new(root)) + let root = if let Some(_root) = machines.pop() { + graph.add_node(MultiMachineNode::new()) } else { - return Self { - graph - }; + return Self { graph }; }; - let mut graph = Self { graph - }; + let mut graph = Self { graph }; let mut populate_idx = 0u64; // round-robin population to avoid congestion let mut nodes_to_populate_now: Vec = vec![root]; // current nodes we are working on @@ -48,13 +43,18 @@ impl MultiMachineTree { let mut nodes_to_populate_later: Vec = Vec::new(); // place all the machines in the graph - while let Some(machine) = machines.pop() { - if graph.nb_children(nodes_to_populate_now[populate_idx as usize]) == max_children_per_parent { - nodes_to_populate_now = nodes_to_populate_later.drain(..).collect(); + while let Some(_machine) = machines.pop() { + if graph.nb_children(nodes_to_populate_now[populate_idx as usize]) + == max_children_per_parent + { + nodes_to_populate_now = core::mem::take(&mut nodes_to_populate_later); populate_idx = 0; // should be useless } - let new_child = graph.add_child(nodes_to_populate_now[populate_idx as usize], MultiMachineNode::new(machine)); + let new_child = graph.add_child( + nodes_to_populate_now[populate_idx as usize], + MultiMachineNode::new(), + ); nodes_to_populate_later.push(new_child); populate_idx = (populate_idx + 1) % nodes_to_populate_now.len() as u64; @@ -70,6 +70,8 @@ impl MultiMachineTree { } fn nb_children(&self, node: NodeIndex) -> u64 { - self.graph.neighbors_directed(node, Direction::Incoming).count() as u64 + self.graph + .neighbors_directed(node, Direction::Incoming) + .count() as u64 } -} \ No newline at end of file +} diff --git a/utils/multi_machine_generator/src/main.rs b/utils/multi_machine_generator/src/main.rs index 389bae6e18..dbbe3fe988 100644 --- a/utils/multi_machine_generator/src/main.rs +++ b/utils/multi_machine_generator/src/main.rs @@ -6,11 +6,10 @@ //! //! We suppose everyone is on the same network and the machines have the fuzzer ready to run on each machine. +use std::{fs, net::SocketAddr, str::FromStr}; -use std::fs; -use std::net::SocketAddr; -use std::str::FromStr; use petgraph::dot::Dot; + use crate::graph::MultiMachineTree; pub mod graph; @@ -41,5 +40,5 @@ fn main() { let dot = Dot::new(&multi_machine_graph.graph); - fs::write("multi_machine.dot", format!("{:?}", dot)).unwrap(); + fs::write("multi_machine.dot", format!("{dot:?}")).unwrap(); }