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 <you@example.com>
This commit is contained in:
Dongjia "toka" Zhang 2024-06-20 17:38:15 +02:00 committed by GitHub
parent 042840dba1
commit e3dd7cf0dc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
27 changed files with 548 additions and 374 deletions

View File

@ -277,6 +277,7 @@ jobs:
- ubuntu - ubuntu
- fuzzers-preflight - fuzzers-preflight
strategy: strategy:
fail-fast: true
matrix: matrix:
os: [ ubuntu-latest ] os: [ ubuntu-latest ]
fuzzer: fuzzer:

View File

@ -21,6 +21,7 @@ members = [
"utils/libafl_benches", "utils/libafl_benches",
"utils/gramatron/construct_automata", "utils/gramatron/construct_automata",
"utils/desyscall", "utils/desyscall",
"utils/multi_machine_generator",
] ]
default-members = [ default-members = [
"libafl", "libafl",

View File

@ -125,7 +125,7 @@ pub fn main() -> Result<(), Error> {
let mut stages = tuple_list!(StdTMinMutationalStage::new( let mut stages = tuple_list!(StdTMinMutationalStage::new(
minimizer, minimizer,
CrashFeedback::new(), CrashFeedback::new(),
1 << 10 1 << 10,
)); ));
let scheduler = QueueScheduler::new(); let scheduler = QueueScheduler::new();

View File

@ -368,7 +368,7 @@ fn fuzz(
.build(tuple_list!(cmplog_observer)) .build(tuple_list!(cmplog_observer))
.unwrap(); .unwrap();
let tracing = AFLppCmplogTracingStage::with_cmplog_observer(cmplog_executor, cmplog_ref); let tracing = AFLppCmplogTracingStage::new(cmplog_executor, cmplog_ref);
// Setup a randomic Input2State stage // Setup a randomic Input2State stage
let rq = MultiMutationalStage::new(AFLppRedQueen::with_cmplog_options(true, true)); let rq = MultiMutationalStage::new(AFLppRedQueen::with_cmplog_options(true, true));

View File

@ -214,12 +214,12 @@ fn fuzz(
// Create a concolic trace // Create a concolic trace
ConcolicTracingStage::new( ConcolicTracingStage::new(
TracingStage::new( TracingStage::new(
MyCommandConfigurator.into_executor(tuple_list!(concolic_observer)) MyCommandConfigurator.into_executor(tuple_list!(concolic_observer)),
), ),
concolic_ref, concolic_ref,
), ),
// Use the concolic trace for z3-based solving // Use the concolic trace for z3-based solving
SimpleConcolicMutationalStage::default(), SimpleConcolicMutationalStage::new(),
); );
fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut restarting_mgr)?; fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut restarting_mgr)?;

View File

@ -27,7 +27,7 @@ use crate::{
use crate::{monitors::PerfFeature, state::HasClientPerfMonitor}; use crate::{monitors::PerfFeature, state::HasClientPerfMonitor};
/// Send a monitor update all 15 (or more) seconds /// 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 /// Holds a scheduler
pub trait HasScheduler: UsesState pub trait HasScheduler: UsesState
@ -247,7 +247,6 @@ where
let monitor_timeout = STATS_TIMEOUT_DEFAULT; let monitor_timeout = STATS_TIMEOUT_DEFAULT;
for _ in 0..iters { for _ in 0..iters {
// log::info!("Starting another fuzz_loop");
manager.maybe_report_progress(state, monitor_timeout)?; manager.maybe_report_progress(state, monitor_timeout)?;
ret = Some(self.fuzz_one(stages, executor, state, manager)?); ret = Some(self.fuzz_one(stages, executor, state, manager)?);
} }

View File

@ -1,6 +1,9 @@
//! The calibration stage. The fuzzer measures the average exec time and the bitmap size. //! 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 core::{fmt::Debug, marker::PhantomData, time::Duration};
use hashbrown::HashSet; use hashbrown::HashSet;
@ -17,7 +20,7 @@ use crate::{
monitors::{AggregatorOps, UserStats, UserStatsValue}, monitors::{AggregatorOps, UserStats, UserStatsValue},
observers::{MapObserver, ObserversTuple}, observers::{MapObserver, ObserversTuple},
schedulers::powersched::SchedulerMetadata, schedulers::powersched::SchedulerMetadata,
stages::{ExecutionCountRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
}; };
@ -75,7 +78,6 @@ pub struct CalibrationStage<C, E, O, OT> {
stage_max: usize, stage_max: usize,
/// If we should track stability /// If we should track stability
track_stability: bool, track_stability: bool,
restart_helper: ExecutionCountRestartHelper,
phantom: PhantomData<(E, O, OT)>, phantom: PhantomData<(E, O, OT)>,
} }
@ -125,8 +127,6 @@ where
let mut iter = self.stage_max; let mut iter = self.stage_max;
// If we restarted after a timeout or crash, do less iterations. // 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()?; let input = state.current_input_cloned()?;
// Run once to get the initial calibration map // Run once to get the initial calibration map
@ -350,14 +350,18 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO: Make sure this is the correct way / there may be a better way? // Calibration stage disallow restarts
self.restart_helper.restart_progress_should_run(state) // 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? // 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 where
F: HasObserverHandle<Observer = C> + Named, F: HasObserverHandle<Observer = C> + Named,
{ {
let map_name = map_feedback.name().clone();
Self { Self {
map_observer_handle: map_feedback.observer_handle().clone(), map_observer_handle: map_feedback.observer_handle().clone(),
map_name: map_feedback.name().clone(), map_name: map_name.clone(),
stage_max: CAL_STAGE_START, stage_max: CAL_STAGE_START,
track_stability: true, track_stability: true,
restart_helper: ExecutionCountRestartHelper::default(),
phantom: PhantomData, 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 where
F: HasObserverHandle<Observer = C> + Named, F: HasObserverHandle<Observer = C> + Named,
{ {
Self { let mut ret = Self::new(map_feedback);
map_observer_handle: map_feedback.observer_handle().clone(), ret.track_stability = false;
map_name: map_feedback.name().clone(), ret
stage_max: CAL_STAGE_START,
track_stability: false,
restart_helper: ExecutionCountRestartHelper::default(),
phantom: PhantomData,
name: Cow::Borrowed(CALIBRATION_STAGE_NAME),
}
} }
} }

View File

@ -1,5 +1,9 @@
//! The colorization stage from `colorization()` in afl++ //! 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 core::{cmp::Ordering, fmt::Debug, marker::PhantomData, ops::Range};
use libafl_bolts::{ use libafl_bolts::{
@ -15,7 +19,7 @@ use crate::{
inputs::HasMutatorBytes, inputs::HasMutatorBytes,
mutators::mutations::buffer_copy, mutators::mutations::buffer_copy,
observers::{MapObserver, ObserversTuple}, observers::{MapObserver, ObserversTuple},
stages::{RetryRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
state::{HasCorpus, HasCurrentTestcase, HasRand, UsesState}, state::{HasCorpus, HasCurrentTestcase, HasRand, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
}; };
@ -104,14 +108,15 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO this stage needs a proper resume // This is a deterministic stage
RetryRestartHelper::restart_progress_should_run(state, self, 3) // 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> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
// TODO this stage needs a proper resume StdRestartHelper::clear_progress(state, &self.name)
RetryRestartHelper::clear_restart_progress(state, self)
} }
} }
@ -309,9 +314,10 @@ where
#[must_use] #[must_use]
/// Creates a new [`ColorizationStage`] /// Creates a new [`ColorizationStage`]
pub fn new(map_observer: &C) -> Self { pub fn new(map_observer: &C) -> Self {
let obs_name = map_observer.name().clone().into_owned();
Self { Self {
map_observer_handle: map_observer.handle(), 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, phantom: PhantomData,
} }
} }

View File

@ -1,8 +1,7 @@
//! This module contains the `concolic` stages, which can trace a target using symbolic execution //! This module contains the `concolic` stages, which can trace a target using symbolic execution
//! and use the results for fuzzer input and mutations. //! and use the results for fuzzer input and mutations.
//! //!
use alloc::borrow::{Cow, ToOwned};
use alloc::borrow::Cow;
#[cfg(feature = "concolic_mutation")] #[cfg(feature = "concolic_mutation")]
use alloc::{string::ToString, vec::Vec}; use alloc::{string::ToString, vec::Vec};
#[cfg(feature = "concolic_mutation")] #[cfg(feature = "concolic_mutation")]
@ -20,7 +19,7 @@ use crate::state::HasClientPerfMonitor;
use crate::{ use crate::{
executors::{Executor, HasObservers}, executors::{Executor, HasObservers},
observers::concolic::ConcolicObserver, observers::concolic::ConcolicObserver,
stages::{RetryRestartHelper, Stage, TracingStage}, stages::{Stage, StdRestartHelper, TracingStage},
state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
}; };
@ -29,7 +28,6 @@ use crate::{
inputs::HasMutatorBytes, inputs::HasMutatorBytes,
mark_feature_time, mark_feature_time,
observers::concolic::{ConcolicMetadata, SymExpr, SymExprRef}, observers::concolic::{ConcolicMetadata, SymExpr, SymExprRef},
stages::ExecutionCountRestartHelper,
start_timer, start_timer,
state::State, state::State,
Evaluator, Evaluator,
@ -38,6 +36,7 @@ use crate::{
/// Wraps a [`TracingStage`] to add concolic observing. /// Wraps a [`TracingStage`] to add concolic observing.
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct ConcolicTracingStage<'a, EM, TE, Z> { pub struct ConcolicTracingStage<'a, EM, TE, Z> {
name: Cow<'static, str>,
inner: TracingStage<EM, TE, Z>, inner: TracingStage<EM, TE, Z>,
observer_handle: Handle<ConcolicObserver<'a>>, observer_handle: Handle<ConcolicObserver<'a>>,
} }
@ -49,10 +48,12 @@ where
type State = TE::State; type State = TE::State;
} }
/// The name for concolic tracer
pub const CONCOLIC_TRACING_STAGE_NAME: &str = "concolictracing";
impl<EM, TE, Z> Named for ConcolicTracingStage<'_, EM, TE, Z> { impl<EM, TE, Z> Named for ConcolicTracingStage<'_, EM, TE, Z> {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("ConcolicTracingStage"); &self.name
&NAME
} }
} }
@ -83,12 +84,15 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, 3) // 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> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
RetryRestartHelper::clear_restart_progress(state, self) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
@ -99,9 +103,13 @@ impl<'a, EM, TE, Z> ConcolicTracingStage<'a, EM, TE, Z> {
inner: TracingStage<EM, TE, Z>, inner: TracingStage<EM, TE, Z>,
observer_handle: Handle<ConcolicObserver<'a>>, observer_handle: Handle<ConcolicObserver<'a>>,
) -> Self { ) -> Self {
let observer_name = observer_handle.name().clone();
Self { Self {
inner, inner,
observer_handle, 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<Item = (SymExprRef, SymExpr)>) -> Vec<
/// A mutational stage that uses Z3 to solve concolic constraints attached to the [`crate::corpus::Testcase`] by the [`ConcolicTracingStage`]. /// A mutational stage that uses Z3 to solve concolic constraints attached to the [`crate::corpus::Testcase`] by the [`ConcolicTracingStage`].
#[cfg(feature = "concolic_mutation")] #[cfg(feature = "concolic_mutation")]
#[derive(Clone, Debug)] #[derive(Clone, Debug, Default)]
pub struct SimpleConcolicMutationalStage<Z> { pub struct SimpleConcolicMutationalStage<Z> {
/// The helper keeps track of progress for timeouting/restarting targets name: Cow<'static, str>,
restart_helper: ExecutionCountRestartHelper,
phantom: PhantomData<Z>, phantom: PhantomData<Z>,
} }
@ -366,6 +373,21 @@ where
type State = Z::State; 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<Z> Named for SimpleConcolicMutationalStage<Z> {
fn name(&self) -> &Cow<'static, str> {
&self.name
}
}
#[cfg(feature = "concolic_mutation")] #[cfg(feature = "concolic_mutation")]
impl<E, EM, Z> Stage<E, EM, Z> for SimpleConcolicMutationalStage<Z> impl<E, EM, Z> Stage<E, EM, Z> for SimpleConcolicMutationalStage<Z>
where where
@ -373,7 +395,7 @@ where
EM: UsesState<State = Self::State>, EM: UsesState<State = Self::State>,
Z: Evaluator<E, EM>, Z: Evaluator<E, EM>,
Z::Input: HasMutatorBytes, Z::Input: HasMutatorBytes,
Self::State: State + HasExecutions + HasCorpus + HasMetadata, Self::State: State + HasExecutions + HasCorpus + HasMetadata + HasNamedMetadata,
{ {
#[inline] #[inline]
fn perform( fn perform(
@ -396,11 +418,8 @@ where
mutations mutations
}); });
let post_restart_skip_cnt =
usize::try_from(self.restart_helper.execs_since_progress_start(state)?)?;
if let Some(mutations) = mutations { 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()?; let mut input_copy = state.current_input_cloned()?;
for (index, new_byte) in mutation { for (index, new_byte) in mutation {
input_copy.bytes_mut()[index] = new_byte; input_copy.bytes_mut()[index] = new_byte;
@ -413,21 +432,34 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
self.restart_helper.restart_progress_should_run(state) // This is a deterministic stage
// Once it failed, then don't retry,
// It will just fail again
StdRestartHelper::no_retry(state, &self.name)
} }
#[inline] #[inline]
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
self.restart_helper.clear_restart_progress(state) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
#[cfg(feature = "concolic_mutation")] #[cfg(feature = "concolic_mutation")]
impl<Z> Default for SimpleConcolicMutationalStage<Z> { impl<Z> SimpleConcolicMutationalStage<Z> {
fn default() -> Self { #[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 { Self {
restart_helper: ExecutionCountRestartHelper::default(), name: Cow::Owned(
SIMPLE_CONCOLIC_MUTATIONAL_NAME.to_owned() + ":" + stage_id.to_string().as_str(),
),
phantom: PhantomData, phantom: PhantomData,
} }
} }

View File

@ -115,13 +115,13 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
// Not executing the target, so restart safety is not needed // Not executing the target, so restart safety is not needed
Ok(true) Ok(true)
} }
#[inline] #[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 // Not executing the target, so restart safety is not needed
Ok(()) Ok(())
} }

View File

@ -1,6 +1,9 @@
//! The tracing stage can trace the target and enrich a [`crate::corpus::Testcase`] with metadata, for example for `CmpLog`. //! 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 core::{fmt::Debug, marker::PhantomData};
use libafl_bolts::{ use libafl_bolts::{
@ -16,7 +19,7 @@ use crate::{
mark_feature_time, mark_feature_time,
observers::{CanTrack, MapObserver, ObserversTuple}, observers::{CanTrack, MapObserver, ObserversTuple},
require_novelties_tracking, require_novelties_tracking,
stages::{RetryRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
start_timer, start_timer,
state::{HasCorpus, HasExecutions, UsesState}, state::{HasCorpus, HasExecutions, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
@ -40,9 +43,13 @@ fn find_next_char(list: &[Option<u8>], mut idx: usize, ch: u8) -> usize {
idx idx
} }
/// The name for generalization stage
pub static GENERALIZATION_STAGE_NAME: &str = "generalization";
/// A stage that runs a tracer executor /// A stage that runs a tracer executor
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct GeneralizationStage<C, EM, O, OT, Z> { pub struct GeneralizationStage<C, EM, O, OT, Z> {
name: Cow<'static, str>,
map_observer_handle: Handle<C>, map_observer_handle: Handle<C>,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(EM, O, OT, Z)>, phantom: PhantomData<(EM, O, OT, Z)>,
@ -50,8 +57,7 @@ pub struct GeneralizationStage<C, EM, O, OT, Z> {
impl<C, EM, O, OT, Z> Named for GeneralizationStage<C, EM, O, OT, Z> { impl<C, EM, O, OT, Z> Named for GeneralizationStage<C, EM, O, OT, Z> {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("GeneralizationStage"); &self.name
&NAME
} }
} }
@ -320,15 +326,15 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO: We need to be able to resume better if something crashes or times out // 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] #[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 // 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] #[must_use]
pub fn new(map_observer: &C) -> Self { pub fn new(map_observer: &C) -> Self {
require_novelties_tracking!("GeneralizationStage", C); require_novelties_tracking!("GeneralizationStage", C);
let name = map_observer.name().clone();
Self { Self {
name: Cow::Owned(
GENERALIZATION_STAGE_NAME.to_owned() + ":" + name.into_owned().as_str(),
),
map_observer_handle: map_observer.handle(), map_observer_handle: map_observer.handle(),
phantom: PhantomData, phantom: PhantomData,
} }

View File

@ -54,11 +54,13 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
// It's a random generation stage
// so you can restart for whatever times you want
Ok(true) 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(()) Ok(())
} }
} }

View File

@ -10,10 +10,10 @@ use crate::{
/// Progress for nested stages. This merely enters/exits the inner stage's scope. /// Progress for nested stages. This merely enters/exits the inner stage's scope.
#[derive(Debug)] #[derive(Debug)]
pub struct NestedStageRestartHelper; pub struct NestedStageStdRestartHelper;
impl NestedStageRestartHelper { impl NestedStageStdRestartHelper {
fn restart_progress_should_run<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error> fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
where where
S: HasNestedStageStatus, S: HasNestedStageStatus,
{ {
@ -21,7 +21,7 @@ impl NestedStageRestartHelper {
Ok(true) Ok(true)
} }
fn clear_restart_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error> fn clear_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
where where
S: HasNestedStageStatus, S: HasNestedStageStatus,
{ {
@ -70,12 +70,12 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self) NestedStageStdRestartHelper::should_restart(state, self)
} }
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
NestedStageRestartHelper::clear_restart_progress(state, self) NestedStageStdRestartHelper::clear_progress(state, self)
} }
} }
@ -134,12 +134,12 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self) NestedStageStdRestartHelper::should_restart(state, self)
} }
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
NestedStageRestartHelper::clear_restart_progress(state, self) NestedStageStdRestartHelper::clear_progress(state, self)
} }
} }
@ -219,12 +219,12 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self) NestedStageStdRestartHelper::should_restart(state, self)
} }
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
NestedStageRestartHelper::clear_restart_progress(state, self) NestedStageStdRestartHelper::clear_progress(state, self)
} }
} }
@ -280,12 +280,12 @@ where
} }
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self) NestedStageStdRestartHelper::should_restart(state, self)
} }
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
NestedStageRestartHelper::clear_restart_progress(state, self) NestedStageStdRestartHelper::clear_progress(state, self)
} }
} }

View File

@ -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. 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}; use core::{fmt, marker::PhantomData};
pub use calibrate::CalibrationStage; pub use calibrate::CalibrationStage;
@ -86,18 +91,19 @@ where
/// This method will be called before every call to [`Stage::perform`]. /// This method will be called before every call to [`Stage::perform`].
/// Initialize the restart tracking for this stage, _if it is not yet initialized_. /// Initialize the restart tracking for this stage, _if it is not yet initialized_.
/// On restart, this will be called again. /// 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`. /// 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<bool, Error>; fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error>;
/// Clear the current status tracking of the associated stage /// 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. /// Run the stage.
/// ///
/// Before a call to perform, [`Stage::restart_progress_should_run`] will be (must be!) 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_restart_progress`] gets 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. /// A call to [`Stage::perform_restartable`] will do these things implicitly.
/// DON'T call this function directly except from `preform_restartable` !!
fn perform( fn perform(
&mut self, &mut self,
fuzzer: &mut Z, fuzzer: &mut Z,
@ -106,7 +112,7 @@ where
manager: &mut EM, manager: &mut EM,
) -> Result<(), Error>; ) -> 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( fn perform_restartable(
&mut self, &mut self,
fuzzer: &mut Z, fuzzer: &mut Z,
@ -114,10 +120,10 @@ where
state: &mut Self::State, state: &mut Self::State,
manager: &mut EM, manager: &mut EM,
) -> Result<(), Error> { ) -> Result<(), Error> {
if self.restart_progress_should_run(state)? { if self.should_restart(state)? {
self.perform(fuzzer, executor, state, manager)?; 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 /// A [`Stage`] that will call a closure
#[derive(Debug)] #[derive(Debug)]
pub struct ClosureStage<CB, E, EM, Z> { pub struct ClosureStage<CB, E, EM, Z> {
name: Cow<'static, str>,
closure: CB, closure: CB,
phantom: PhantomData<(E, EM, Z)>, phantom: PhantomData<(E, EM, Z)>,
} }
@ -299,8 +310,7 @@ where
impl<CB, E, EM, Z> Named for ClosureStage<CB, E, EM, Z> { impl<CB, E, EM, Z> Named for ClosureStage<CB, E, EM, Z> {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("<unnamed fn>"); &self.name
&NAME
} }
} }
@ -323,14 +333,15 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// Make sure we don't get stuck crashing on a single closure // There's no restart safety in the content of the closure.
RetryRestartHelper::restart_progress_should_run(state, self, 3) // don't restart
StdRestartHelper::no_retry(state, &self.name)
} }
#[inline] #[inline]
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
RetryRestartHelper::clear_restart_progress(state, self) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
@ -339,28 +350,25 @@ impl<CB, E, EM, Z> ClosureStage<CB, E, EM, Z> {
/// Create a new [`ClosureStage`] /// Create a new [`ClosureStage`]
#[must_use] #[must_use]
pub fn new(closure: CB) -> Self { 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 { Self {
name: Cow::Owned(CLOSURE_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_ref()),
closure, closure,
phantom: PhantomData, phantom: PhantomData,
} }
} }
} }
impl<CB, E, EM, Z> From<CB> for ClosureStage<CB, E, EM, Z>
where
CB: FnMut(&mut Z, &mut E, &mut <Self as UsesState>::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`] /// Allows us to use a [`push::PushStage`] as a normal [`Stage`]
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
#[derive(Debug)] #[derive(Debug)]
pub struct PushStageAdapter<CS, EM, OT, PS, Z> { pub struct PushStageAdapter<CS, EM, OT, PS, Z> {
name: Cow<'static, str>,
push_stage: PS, push_stage: PS,
phantom: PhantomData<(CS, EM, OT, Z)>, phantom: PhantomData<(CS, EM, OT, Z)>,
} }
@ -370,12 +378,25 @@ impl<CS, EM, OT, PS, Z> PushStageAdapter<CS, EM, OT, PS, Z> {
/// to be used as a normal [`Stage`] /// to be used as a normal [`Stage`]
#[must_use] #[must_use]
pub fn new(push_stage: PS) -> Self { 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 { Self {
name: Cow::Owned(
PUSH_STAGE_ADAPTER_NAME.to_owned() + ":" + stage_id.to_string().as_str(),
),
push_stage, push_stage,
phantom: PhantomData, 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<CS, EM, OT, PS, Z> UsesState for PushStageAdapter<CS, EM, OT, PS, Z> impl<CS, EM, OT, PS, Z> UsesState for PushStageAdapter<CS, EM, OT, PS, Z>
where where
@ -384,11 +405,23 @@ where
type State = CS::State; type State = CS::State;
} }
impl<CS, EM, OT, PS, Z> Named for PushStageAdapter<CS, EM, OT, PS, Z> {
#[must_use]
fn name(&self) -> &Cow<'static, str> {
&self.name
}
}
impl<CS, E, EM, OT, PS, Z> Stage<E, EM, Z> for PushStageAdapter<CS, EM, OT, PS, Z> impl<CS, E, EM, OT, PS, Z> Stage<E, EM, Z> for PushStageAdapter<CS, EM, OT, PS, Z>
where where
CS: Scheduler, CS: Scheduler,
Self::State: Self::State: HasExecutions
HasExecutions + HasMetadata + HasRand + HasCorpus + HasLastReportTime + HasCurrentCorpusId, + HasRand
+ HasCorpus
+ HasLastReportTime
+ HasCurrentCorpusId
+ HasNamedMetadata
+ HasMetadata,
E: Executor<EM, Z> + HasObservers<Observers = OT, State = Self::State>, E: Executor<EM, Z> + HasObservers<Observers = OT, State = Self::State>,
EM: EventFirer<State = Self::State> EM: EventFirer<State = Self::State>
+ EventRestarter + EventRestarter
@ -446,48 +479,51 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO: Proper restart handling - call post_exec at the right time, etc... // TODO: Proper restart handling - call post_exec at the right time, etc...
Ok(true) StdRestartHelper::no_retry(state, &self.name)
} }
#[inline] #[inline]
fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
Ok(()) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
/// Progress which permits a fixed amount of resumes per round of fuzzing. If this amount is ever /// 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. /// exceeded, the input will no longer be executed by this stage.
#[derive(Clone, Deserialize, Serialize, Debug)] #[derive(Clone, Deserialize, Serialize, Debug)]
pub struct RetryRestartHelper { pub struct StdRestartHelper {
tries_remaining: Option<usize>, tries_remaining: Option<usize>,
skipped: HashSet<CorpusId>, skipped: HashSet<CorpusId>,
} }
impl_serdeany!(RetryRestartHelper); impl_serdeany!(StdRestartHelper);
impl StdRestartHelper {
/// Don't allow restart
pub fn no_retry<S>(state: &mut S, name: &str) -> Result<bool, Error>
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 /// Initializes (or counts down in) the progress helper, giving it the amount of max retries
/// ///
/// Returns `true` if the stage should run /// Returns `true` if the stage should run
pub fn restart_progress_should_run<S, ST>( pub fn should_restart<S>(state: &mut S, name: &str, max_retries: usize) -> Result<bool, Error>
state: &mut S,
stage: &ST,
max_retries: usize,
) -> Result<bool, Error>
where where
S: HasNamedMetadata + HasCurrentCorpusId, S: HasNamedMetadata + HasCurrentCorpusId,
ST: Named,
{ {
let corpus_id = state.current_corpus_id()?.ok_or_else(|| { let corpus_id = state.current_corpus_id()?.ok_or_else(|| {
Error::illegal_state( 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 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), tries_remaining: Some(initial_tries_remaining),
skipped: HashSet::new(), skipped: HashSet::new(),
}); });
@ -515,14 +551,11 @@ impl RetryRestartHelper {
} }
/// Clears the progress /// Clears the progress
pub fn clear_restart_progress<S, ST>(state: &mut S, stage: &ST) -> Result<(), Error> pub fn clear_progress<S>(state: &mut S, name: &str) -> Result<(), Error>
where where
S: HasNamedMetadata, S: HasNamedMetadata,
ST: Named,
{ {
state state.named_metadata_mut::<Self>(name)?.tries_remaining = None;
.named_metadata_mut::<Self>(stage.name())?
.tries_remaining = None;
Ok(()) Ok(())
} }
} }
@ -597,15 +630,15 @@ impl ExecutionCountRestartHelper {
} }
/// The execs done since start of this [`Stage`]/helper /// The execs done since start of this [`Stage`]/helper
pub fn execs_since_progress_start<S>(&mut self, state: &mut S) -> Result<u64, Error> pub fn execs_since_progress_start<S>(&mut self, state: &mut S, name: &str) -> Result<u64, Error>
where where
S: HasMetadata + HasExecutions, S: HasNamedMetadata + HasExecutions,
{ {
let started_at_execs = if let Some(started_at_execs) = self.started_at_execs { let started_at_execs = if let Some(started_at_execs) = self.started_at_execs {
started_at_execs started_at_execs
} else { } else {
state state
.metadata::<ExecutionCountRestartHelperMetadata>() .named_metadata::<ExecutionCountRestartHelperMetadata>(name)
.map(|x| { .map(|x| {
self.started_at_execs = Some(x.started_at_execs); self.started_at_execs = Some(x.started_at_execs);
x.started_at_execs x.started_at_execs
@ -620,26 +653,27 @@ impl ExecutionCountRestartHelper {
} }
/// Initialize progress for the stage this wrapper wraps. /// Initialize progress for the stage this wrapper wraps.
pub fn restart_progress_should_run<S>(&mut self, state: &mut S) -> Result<bool, Error> pub fn should_restart<S>(&mut self, state: &mut S, name: &str) -> Result<bool, Error>
where where
S: HasMetadata + HasExecutions, S: HasNamedMetadata + HasExecutions,
{ {
let executions = *state.executions(); let executions = *state.executions();
let metadata = state.metadata_or_insert_with(|| ExecutionCountRestartHelperMetadata { let metadata =
started_at_execs: executions, state.named_metadata_or_insert_with(name, || ExecutionCountRestartHelperMetadata {
}); started_at_execs: executions,
});
self.started_at_execs = Some(metadata.started_at_execs); self.started_at_execs = Some(metadata.started_at_execs);
Ok(true) Ok(true)
} }
/// Clear progress for the stage this wrapper wraps. /// Clear progress for the stage this wrapper wraps.
pub fn clear_restart_progress<S>(&mut self, state: &mut S) -> Result<(), Error> pub fn clear_progress<S>(&mut self, state: &mut S) -> Result<(), Error>
where where
S: HasMetadata, S: HasMetadata,
{ {
self.started_at_execs = None; self.started_at_execs = None;
let _metadata = state.remove_metadata::<ExecutionCountRestartHelperMetadata>(); let _metadata = state.remove_metadata::<ExecutionCountRestartHelperMetadata>();
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(()) Ok(())
} }
} }
@ -655,7 +689,7 @@ pub mod test {
use crate::{ use crate::{
corpus::{Corpus, HasCurrentCorpusId, Testcase}, corpus::{Corpus, HasCurrentCorpusId, Testcase},
inputs::NopInput, inputs::NopInput,
stages::{RetryRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
state::{test::test_std_state, HasCorpus, State, UsesState}, state::{test::test_std_state, HasCorpus, State, UsesState},
HasMetadata, HasMetadata,
}; };
@ -674,7 +708,7 @@ pub mod test {
impl TestProgress { impl TestProgress {
#[allow(clippy::unnecessary_wraps)] #[allow(clippy::unnecessary_wraps)]
fn restart_progress_should_run<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error> fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
where where
S: HasMetadata, S: HasMetadata,
{ {
@ -690,7 +724,7 @@ pub mod test {
Ok(true) Ok(true)
} }
fn clear_restart_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error> fn clear_progress<S, ST>(state: &mut S, _stage: &ST) -> Result<(), Error>
where where
S: HasMetadata, S: HasMetadata,
{ {
@ -727,12 +761,12 @@ pub mod test {
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
TestProgress::restart_progress_should_run(state, self) TestProgress::should_restart(state, self)
} }
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
TestProgress::clear_restart_progress(state, self) TestProgress::clear_progress(state, self)
} }
} }
@ -742,7 +776,7 @@ pub mod test {
// No concurrency per testcase // No concurrency per testcase
#[cfg(any(not(feature = "serdeany_autoreg"), miri))] #[cfg(any(not(feature = "serdeany_autoreg"), miri))]
unsafe { unsafe {
RetryRestartHelper::register(); StdRestartHelper::register();
} }
struct StageWithOneTry; struct StageWithOneTry;
@ -763,44 +797,58 @@ pub mod test {
for _ in 0..10 { for _ in 0..10 {
// used normally, no retries means we never skip // used normally, no retries means we never skip
assert!(RetryRestartHelper::restart_progress_should_run( assert!(StdRestartHelper::should_restart(
&mut state, &stage, 1 &mut state,
stage.name(),
1
)?); )?);
RetryRestartHelper::clear_restart_progress(&mut state, &stage)?; StdRestartHelper::clear_progress(&mut state, stage.name())?;
} }
for _ in 0..10 { for _ in 0..10 {
// used normally, only one retry means we never skip // used normally, only one retry means we never skip
assert!(RetryRestartHelper::restart_progress_should_run( assert!(StdRestartHelper::should_restart(
&mut state, &stage, 2 &mut state,
stage.name(),
2
)?); )?);
assert!(RetryRestartHelper::restart_progress_should_run( assert!(StdRestartHelper::should_restart(
&mut state, &stage, 2 &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( assert!(StdRestartHelper::should_restart(
&mut state, &stage, 2 &mut state,
stage.name(),
2
)?); )?);
// task failed, let's resume // task failed, let's resume
// we still have one more try! // we still have one more try!
assert!(RetryRestartHelper::restart_progress_should_run( assert!(StdRestartHelper::should_restart(
&mut state, &stage, 2 &mut state,
stage.name(),
2
)?); )?);
// task failed, let's resume // task failed, let's resume
// out of retries, so now we skip // out of retries, so now we skip
assert!(!RetryRestartHelper::restart_progress_should_run( assert!(!StdRestartHelper::should_restart(
&mut state, &stage, 2 &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 // we previously exhausted this testcase's retries, so we skip
assert!(!RetryRestartHelper::restart_progress_should_run( assert!(!StdRestartHelper::should_restart(
&mut state, &stage, 2 &mut state,
stage.name(),
2
)?); )?);
RetryRestartHelper::clear_restart_progress(&mut state, &stage)?; StdRestartHelper::clear_progress(&mut state, stage.name())?;
Ok(()) Ok(())
} }

View File

@ -1,7 +1,10 @@
//| The [`MutationalStage`] is the default stage used during fuzzing. //| 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. //! 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 core::marker::PhantomData;
use libafl_bolts::{rands::Rand, Named}; use libafl_bolts::{rands::Rand, Named};
@ -12,7 +15,7 @@ use crate::{
inputs::Input, inputs::Input,
mark_feature_time, mark_feature_time,
mutators::{MultiMutator, MutationResult, Mutator}, mutators::{MultiMutator, MutationResult, Mutator},
stages::{ExecutionCountRestartHelper, RetryRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
start_timer, start_timer,
state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, UsesState}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
@ -94,9 +97,6 @@ where
/// Gets the number of iterations this mutator should run for. /// Gets the number of iterations this mutator should run for.
fn iterations(&self, state: &mut Self::State) -> Result<usize, Error>; fn iterations(&self, state: &mut Self::State) -> Result<usize, Error>;
/// 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<u64, Error>;
/// Runs this (mutational) stage for the given testcase /// Runs this (mutational) stage for the given testcase
#[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely... #[allow(clippy::cast_possible_wrap)] // more than i32 stages on 32 bit system - highly unlikely...
fn perform_mutational( fn perform_mutational(
@ -155,12 +155,12 @@ pub static DEFAULT_MUTATIONAL_MAX_ITERATIONS: usize = 128;
/// The default mutational stage /// The default mutational stage
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct StdMutationalStage<E, EM, I, M, Z> { pub struct StdMutationalStage<E, EM, I, M, Z> {
/// The name
name: Cow<'static, str>,
/// The mutator(s) to use /// The mutator(s) to use
mutator: M, mutator: M,
/// The maximum amount of iterations we should do each round /// The maximum amount of iterations we should do each round
max_iterations: usize, max_iterations: usize,
/// The progress helper for this mutational stage
restart_helper: ExecutionCountRestartHelper,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(E, EM, I, Z)>, phantom: PhantomData<(E, EM, I, Z)>,
} }
@ -171,7 +171,7 @@ where
EM: UsesState<State = Self::State>, EM: UsesState<State = Self::State>,
M: Mutator<I, Self::State>, M: Mutator<I, Self::State>,
Z: Evaluator<E, EM>, Z: Evaluator<E, EM>,
Self::State: HasCorpus + HasRand + HasExecutions + HasMetadata, Self::State: HasCorpus + HasRand + HasExecutions + HasMetadata + HasNamedMetadata,
I: MutatedTransform<Self::Input, Self::State> + Clone, I: MutatedTransform<Self::Input, Self::State> + Clone,
{ {
/// The mutator, added to this stage /// The mutator, added to this stage
@ -190,12 +190,13 @@ where
fn iterations(&self, state: &mut Self::State) -> Result<usize, Error> { fn iterations(&self, state: &mut Self::State) -> Result<usize, Error> {
Ok(1 + state.rand_mut().below(self.max_iterations)) Ok(1 + state.rand_mut().below(self.max_iterations))
} }
fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result<u64, Error> {
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<E, EM, I, M, Z> UsesState for StdMutationalStage<E, EM, I, M, Z> impl<E, EM, I, M, Z> UsesState for StdMutationalStage<E, EM, I, M, Z>
where where
Z: UsesState, Z: UsesState,
@ -203,13 +204,19 @@ where
type State = Z::State; type State = Z::State;
} }
impl<E, EM, I, M, Z> Named for StdMutationalStage<E, EM, I, M, Z> {
fn name(&self) -> &Cow<'static, str> {
&self.name
}
}
impl<E, EM, I, M, Z> Stage<E, EM, Z> for StdMutationalStage<E, EM, I, M, Z> impl<E, EM, I, M, Z> Stage<E, EM, Z> for StdMutationalStage<E, EM, I, M, Z>
where where
E: UsesState<State = Self::State>, E: UsesState<State = Self::State>,
EM: UsesState<State = Self::State>, EM: UsesState<State = Self::State>,
M: Mutator<I, Self::State>, M: Mutator<I, Self::State>,
Z: Evaluator<E, EM>, Z: Evaluator<E, EM>,
Self::State: HasCorpus + HasRand + HasMetadata + HasExecutions, Self::State: HasCorpus + HasRand + HasMetadata + HasExecutions + HasNamedMetadata,
I: MutatedTransform<Self::Input, Self::State> + Clone, I: MutatedTransform<Self::Input, Self::State> + Clone,
{ {
#[inline] #[inline]
@ -229,14 +236,12 @@ where
ret ret
} }
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
Ok(true) StdRestartHelper::should_restart(state, &self.name, 3)
// self.restart_helper.restart_progress_should_run(state)
} }
fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
Ok(()) StdRestartHelper::clear_progress(state, &self.name)
// self.restart_helper.clear_restart_progress(state)
} }
} }
@ -274,10 +279,18 @@ where
/// Creates a new transforming mutational stage with the given max iterations /// Creates a new transforming mutational stage with the given max iterations
pub fn transforming_with_max_iterations(mutator: M, max_iterations: usize) -> Self { 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 { Self {
name: Cow::Owned(
MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(),
),
mutator, mutator,
max_iterations, max_iterations,
restart_helper: ExecutionCountRestartHelper::default(),
phantom: PhantomData, phantom: PhantomData,
} }
} }
@ -286,11 +299,17 @@ where
/// A mutational stage that operates on multiple inputs, as returned by [`MultiMutator::multi_mutate`]. /// A mutational stage that operates on multiple inputs, as returned by [`MultiMutator::multi_mutate`].
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct MultiMutationalStage<E, EM, I, M, Z> { pub struct MultiMutationalStage<E, EM, I, M, Z> {
name: Cow<'static, str>,
mutator: M, mutator: M,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(E, EM, I, Z)>, 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<E, EM, I, M, Z> UsesState for MultiMutationalStage<E, EM, I, M, Z> impl<E, EM, I, M, Z> UsesState for MultiMutationalStage<E, EM, I, M, Z>
where where
Z: UsesState, Z: UsesState,
@ -300,8 +319,7 @@ where
impl<E, EM, I, M, Z> Named for MultiMutationalStage<E, EM, I, M, Z> { impl<E, EM, I, M, Z> Named for MultiMutationalStage<E, EM, I, M, Z> {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("MultiMutational"); &self.name
&NAME
} }
} }
@ -315,15 +333,14 @@ where
I: MutatedTransform<Self::Input, Self::State> + Clone, I: MutatedTransform<Self::Input, Self::State> + Clone,
{ {
#[inline] #[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO: add proper crash/timeout handling // Make sure we don't get stuck crashing on a single testcase
// For now, Make sure we don't get stuck crashing on a single testcase StdRestartHelper::should_restart(state, &self.name, 3)
RetryRestartHelper::restart_progress_should_run(state, self, 3)
} }
#[inline] #[inline]
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
RetryRestartHelper::clear_restart_progress(state, self) StdRestartHelper::clear_progress(state, &self.name)
} }
#[inline] #[inline]
@ -370,7 +387,16 @@ where
impl<E, EM, I, M, Z> MultiMutationalStage<E, EM, I, M, Z> { impl<E, EM, I, M, Z> MultiMutationalStage<E, EM, I, M, Z> {
/// Creates a new transforming mutational stage /// Creates a new transforming mutational stage
pub fn transforming(mutator: M) -> Self { 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 { Self {
name: Cow::Owned(
MULTI_MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(),
),
mutator, mutator,
phantom: PhantomData, phantom: PhantomData,
} }

View File

@ -1,6 +1,9 @@
//! The power schedules. This stage should be invoked after the calibration stage. //! 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 core::{fmt::Debug, marker::PhantomData};
use libafl_bolts::Named; use libafl_bolts::Named;
@ -10,10 +13,13 @@ use crate::{
fuzzer::Evaluator, fuzzer::Evaluator,
mutators::Mutator, mutators::Mutator,
schedulers::{testcase_score::CorpusPowerTestcaseScore, TestcaseScore}, schedulers::{testcase_score::CorpusPowerTestcaseScore, TestcaseScore},
stages::{mutational::MutatedTransform, ExecutionCountRestartHelper, MutationalStage, Stage}, stages::{mutational::MutatedTransform, MutationalStage, Stage, StdRestartHelper},
state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, UsesState}, 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++ /// Default name for `PowerMutationalStage`; derived from AFL++
pub const POWER_MUTATIONAL_STAGE_NAME: &str = "power"; pub const POWER_MUTATIONAL_STAGE_NAME: &str = "power";
/// The mutational stage using power schedules /// The mutational stage using power schedules
@ -22,8 +28,6 @@ pub struct PowerMutationalStage<E, F, EM, I, M, Z> {
name: Cow<'static, str>, name: Cow<'static, str>,
/// The mutators we use /// The mutators we use
mutator: M, mutator: M,
/// Helper for restarts
restart_helper: ExecutionCountRestartHelper,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(E, F, EM, I, Z)>, phantom: PhantomData<(E, F, EM, I, Z)>,
} }
@ -47,7 +51,7 @@ where
EM: UsesState<State = Self::State>, EM: UsesState<State = Self::State>,
F: TestcaseScore<Self::State>, F: TestcaseScore<Self::State>,
M: Mutator<I, Self::State>, M: Mutator<I, Self::State>,
Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions, Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions + HasNamedMetadata,
Z: Evaluator<E, EM, State = Self::State>, Z: Evaluator<E, EM, State = Self::State>,
I: MutatedTransform<E::Input, Self::State> + Clone, I: MutatedTransform<E::Input, Self::State> + Clone,
{ {
@ -72,10 +76,6 @@ where
Ok(score) Ok(score)
} }
fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result<u64, Error> {
self.restart_helper.execs_since_progress_start(state)
}
} }
impl<E, F, EM, I, M, Z> Stage<E, EM, Z> for PowerMutationalStage<E, F, EM, I, M, Z> impl<E, F, EM, I, M, Z> Stage<E, EM, Z> for PowerMutationalStage<E, F, EM, I, M, Z>
@ -84,7 +84,7 @@ where
EM: UsesState<State = Self::State>, EM: UsesState<State = Self::State>,
F: TestcaseScore<Self::State>, F: TestcaseScore<Self::State>,
M: Mutator<I, Self::State>, M: Mutator<I, Self::State>,
Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions, Self::State: HasCorpus + HasMetadata + HasRand + HasExecutions + HasNamedMetadata,
Z: Evaluator<E, EM, State = Self::State>, Z: Evaluator<E, EM, State = Self::State>,
I: MutatedTransform<Self::Input, Self::State> + Clone, I: MutatedTransform<Self::Input, Self::State> + Clone,
{ {
@ -101,14 +101,13 @@ where
ret ret
} }
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
Ok(true) // Make sure we don't get stuck crashing on a single testcase
// self.restart_helper.restart_progress_should_run(state) StdRestartHelper::should_restart(state, &self.name, 3)
} }
fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
Ok(()) StdRestartHelper::clear_progress(state, &self.name)
// self.restart_helper.clear_restart_progress(state)
} }
} }
@ -123,11 +122,18 @@ where
{ {
/// Creates a new [`PowerMutationalStage`] /// Creates a new [`PowerMutationalStage`]
pub fn new(mutator: M) -> Self { 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 { Self {
name: Cow::Borrowed(POWER_MUTATIONAL_STAGE_NAME), name: Cow::Owned(
POWER_MUTATIONAL_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(),
),
mutator, mutator,
phantom: PhantomData, phantom: PhantomData,
restart_helper: ExecutionCountRestartHelper::default(),
} }
} }
} }

View File

@ -126,13 +126,13 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
// Not running the target so we wont't crash/timeout and, hence, don't need to restore anything // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything
Ok(true) Ok(true)
} }
#[inline] #[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 // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything
Ok(()) Ok(())
} }

View File

@ -1,6 +1,6 @@
//! The [`SyncFromDiskStage`] is a stage that imports inputs from disk for e.g. sync with AFL //! 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 core::marker::PhantomData;
use std::{ use std::{
fs, fs,
@ -20,7 +20,7 @@ use crate::{
executors::{Executor, ExitKind, HasObservers}, executors::{Executor, ExitKind, HasObservers},
fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor}, fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor},
inputs::{Input, InputConverter, UsesInput}, inputs::{Input, InputConverter, UsesInput},
stages::{RetryRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
state::{HasCorpus, HasExecutions, HasRand, State, UsesState}, state::{HasCorpus, HasExecutions, HasRand, State, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
}; };
@ -149,24 +149,24 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO: Needs proper crash handling for when an imported testcase crashes // TODO: Needs proper crash handling for when an imported testcase crashes
// For now, Make sure we don't get stuck crashing on this testcase // 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] #[inline]
fn clear_restart_progress(&mut self, state: &mut Self::State) -> Result<(), Error> { fn clear_progress(&mut self, state: &mut Self::State) -> Result<(), Error> {
RetryRestartHelper::clear_restart_progress(state, self) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
impl<CB, E, EM, Z> SyncFromDiskStage<CB, E, EM, Z> { impl<CB, E, EM, Z> SyncFromDiskStage<CB, E, EM, Z> {
/// Creates a new [`SyncFromDiskStage`] /// Creates a new [`SyncFromDiskStage`]
#[must_use] #[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 { Self {
name: Cow::Borrowed(SYNC_FROM_DISK_STAGE_NAME), name: Cow::Owned(SYNC_FROM_DISK_STAGE_NAME.to_owned() + ":" + name),
phantom: PhantomData, phantom: PhantomData,
sync_dir, sync_dir,
load_callback, load_callback,
@ -367,13 +367,13 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
// No restart handling needed - does not execute the target. // No restart handling needed - does not execute the target.
Ok(true) Ok(true)
} }
#[inline] #[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. // Not needed - does not execute the target.
Ok(()) Ok(())
} }

View File

@ -1,6 +1,9 @@
//! The [`TMinMutationalStage`] is a stage which will attempt to minimize corpus entries. //! 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 core::{borrow::BorrowMut, fmt::Debug, hash::Hash, marker::PhantomData};
use ahash::RandomState; use ahash::RandomState;
@ -29,7 +32,8 @@ use crate::{
state::{ state::{
HasCorpus, HasCurrentTestcase, HasExecutions, HasMaxSize, HasSolutions, State, UsesState, HasCorpus, HasCurrentTestcase, HasExecutions, HasMaxSize, HasSolutions, State, UsesState,
}, },
Error, ExecutesInput, ExecutionProcessor, HasFeedback, HasMetadata, HasScheduler, Error, ExecutesInput, ExecutionProcessor, HasFeedback, HasMetadata, HasNamedMetadata,
HasScheduler,
}; };
#[cfg(feature = "introspection")] #[cfg(feature = "introspection")]
use crate::{monitors::PerfFeature, state::HasClientPerfMonitor}; use crate::{monitors::PerfFeature, state::HasClientPerfMonitor};
@ -79,8 +83,14 @@ where
let orig_max_size = state.max_size(); let orig_max_size = state.max_size();
// basically copy-pasted from mutational.rs // basically copy-pasted from mutational.rs
let num = self.iterations(state)? let num = self
- usize::try_from(self.execs_since_progress_start(state)?).unwrap(); .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); start_timer!(state);
let transformed = let transformed =
@ -203,6 +213,8 @@ where
/// The default corpus entry minimising mutational stage /// The default corpus entry minimising mutational stage
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct StdTMinMutationalStage<E, EM, F, FF, IP, M, Z> { pub struct StdTMinMutationalStage<E, EM, F, FF, IP, M, Z> {
/// The name
name: Cow<'static, str>,
/// The mutator(s) this stage uses /// The mutator(s) this stage uses
mutator: M, mutator: M,
/// The factory /// The factory
@ -231,16 +243,17 @@ where
FF: FeedbackFactory<F, E::Observers>, FF: FeedbackFactory<F, E::Observers>,
F: Feedback<Self::State>, F: Feedback<Self::State>,
Self::Input: MutatedTransform<Self::Input, Self::State, Post = IP> + Clone + HasLen + Hash, Self::Input: MutatedTransform<Self::Input, Self::State, Post = IP> + Clone + HasLen + Hash,
Self::State: HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize, Self::State:
HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize + HasNamedMetadata,
M: Mutator<Self::Input, Self::State>, M: Mutator<Self::Input, Self::State>,
IP: MutatedTransformPost<Self::State> + Clone, IP: MutatedTransformPost<Self::State> + Clone,
{ {
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
self.restart_helper.restart_progress_should_run(state) self.restart_helper.should_restart(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> {
self.restart_helper.clear_restart_progress(state) self.restart_helper.clear_progress(state)
} }
fn perform( fn perform(
@ -270,6 +283,17 @@ where
} }
} }
impl<E, EM, F, FF, IP, M, Z> Named for StdTMinMutationalStage<E, EM, F, FF, IP, M, Z> {
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<E, EM, F, FF, IP, M, Z> TMinMutationalStage<E, EM, F, IP, M, Z> impl<E, EM, F, FF, IP, M, Z> TMinMutationalStage<E, EM, F, IP, M, Z>
for StdTMinMutationalStage<E, EM, F, FF, IP, M, Z> for StdTMinMutationalStage<E, EM, F, FF, IP, M, Z>
where where
@ -280,7 +304,8 @@ where
FF: FeedbackFactory<F, E::Observers>, FF: FeedbackFactory<F, E::Observers>,
F: Feedback<Self::State>, F: Feedback<Self::State>,
Self::Input: MutatedTransform<Self::Input, Self::State, Post = IP> + Clone + HasLen + Hash, Self::Input: MutatedTransform<Self::Input, Self::State, Post = IP> + Clone + HasLen + Hash,
Self::State: HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize, Self::State:
HasMetadata + HasExecutions + HasSolutions + HasCorpus + HasMaxSize + HasNamedMetadata,
M: Mutator<Self::Input, Self::State>, M: Mutator<Self::Input, Self::State>,
IP: MutatedTransformPost<Self::State> + Clone, IP: MutatedTransformPost<Self::State> + Clone,
{ {
@ -302,14 +327,22 @@ where
} }
fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result<u64, Error> { fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result<u64, Error> {
self.restart_helper.execs_since_progress_start(state) self.restart_helper
.execs_since_progress_start(state, &self.name)
} }
} }
impl<E, EM, F, FF, IP, M, Z> StdTMinMutationalStage<E, EM, F, FF, IP, M, Z> { impl<E, EM, F, FF, IP, M, Z> StdTMinMutationalStage<E, EM, F, FF, IP, M, Z> {
/// Creates a new minimizing mutational stage that will minimize provided corpus entries /// Creates a new minimizing mutational stage that will minimize provided corpus entries
pub fn new(mutator: M, factory: FF, runs: usize) -> Self { 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 { Self {
name: Cow::Owned(TMIN_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str()),
mutator, mutator,
factory, factory,
runs, runs,

View File

@ -1,6 +1,9 @@
//! The tracing stage can trace the target and enrich a testcase with metadata, for example for `CmpLog`. //! 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 core::{fmt::Debug, marker::PhantomData};
use libafl_bolts::Named; use libafl_bolts::Named;
@ -9,7 +12,7 @@ use crate::{
executors::{Executor, HasObservers, ShadowExecutor}, executors::{Executor, HasObservers, ShadowExecutor},
mark_feature_time, mark_feature_time,
observers::ObserversTuple, observers::ObserversTuple,
stages::{RetryRestartHelper, Stage}, stages::{Stage, StdRestartHelper},
start_timer, start_timer,
state::{HasCorpus, HasCurrentTestcase, HasExecutions, State, UsesState}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, State, UsesState},
Error, HasNamedMetadata, Error, HasNamedMetadata,
@ -20,8 +23,8 @@ use crate::{monitors::PerfFeature, state::HasClientPerfMonitor};
/// A stage that runs a tracer executor /// A stage that runs a tracer executor
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct TracingStage<EM, TE, Z> { pub struct TracingStage<EM, TE, Z> {
name: Cow<'static, str>,
tracer_executor: TE, tracer_executor: TE,
max_retries: usize,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(EM, TE, Z)>, phantom: PhantomData<(EM, TE, Z)>,
} }
@ -42,7 +45,7 @@ where
{ {
#[allow(rustdoc::broken_intra_doc_links)] #[allow(rustdoc::broken_intra_doc_links)]
/// Perform tracing on the given `CorpusId`. Useful for if wrapping [`TracingStage`] with your /// 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. /// see [`super::ConcolicTracingStage`]'s implementation as an example of usage.
pub fn trace( pub fn trace(
&mut self, &mut self,
@ -96,40 +99,43 @@ where
self.trace(fuzzer, state, manager) self.trace(fuzzer, state, manager)
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries) 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> {
RetryRestartHelper::clear_restart_progress(state, self) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
impl<EM, TE, Z> Named for TracingStage<EM, TE, Z> { impl<EM, TE, Z> Named for TracingStage<EM, TE, Z> {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("TracingStage"); &self.name
&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<EM, TE, Z> TracingStage<EM, TE, Z> { impl<EM, TE, Z> TracingStage<EM, TE, Z> {
/// Creates a new default stage /// Creates a new default stage
pub fn new(tracer_executor: TE) -> Self { 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 { Self {
name: Cow::Owned(TRACING_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_ref()),
tracer_executor, tracer_executor,
max_retries: 10,
phantom: PhantomData, 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 /// Gets the underlying tracer executor
pub fn executor(&self) -> &TE { pub fn executor(&self) -> &TE {
&self.tracer_executor &self.tracer_executor
@ -144,7 +150,7 @@ impl<EM, TE, Z> TracingStage<EM, TE, Z> {
/// A stage that runs the shadow executor using also the shadow observers /// A stage that runs the shadow executor using also the shadow observers
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct ShadowTracingStage<E, EM, SOT, Z> { pub struct ShadowTracingStage<E, EM, SOT, Z> {
max_retries: usize, name: Cow<'static, str>,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(E, EM, SOT, Z)>, phantom: PhantomData<(E, EM, SOT, Z)>,
} }
@ -155,14 +161,17 @@ where
{ {
type State = E::State; 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<E, EM, SOT, Z> Named for ShadowTracingStage<E, EM, SOT, Z> impl<E, EM, SOT, Z> Named for ShadowTracingStage<E, EM, SOT, Z>
where where
E: UsesState, E: UsesState,
{ {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("ShadowTracingStage"); &self.name
&NAME
} }
} }
@ -210,12 +219,12 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries) 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> {
RetryRestartHelper::clear_restart_progress(state, self) StdRestartHelper::clear_progress(state, &self.name)
} }
} }
@ -229,17 +238,17 @@ where
{ {
/// Creates a new default stage /// Creates a new default stage
pub fn new(_executor: &mut ShadowExecutor<E, SOT>) -> Self { pub fn new(_executor: &mut ShadowExecutor<E, SOT>) -> 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 { Self {
max_retries: 10, name: Cow::Owned(
SHADOW_TRACING_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_str(),
),
phantom: PhantomData, 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
}
} }

View File

@ -248,10 +248,6 @@ where
1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS), 1 + state.rand_mut().below(DEFAULT_MUTATIONAL_MAX_ITERATIONS),
) )
} }
fn execs_since_progress_start(&mut self, state: &mut Self::State) -> Result<u64, Error> {
self.restart_helper.execs_since_progress_start(state)
}
} }
impl<E, EM, I, M, Z> UsesState for TuneableMutationalStage<E, EM, I, M, Z> impl<E, EM, I, M, Z> UsesState for TuneableMutationalStage<E, EM, I, M, Z>
@ -287,12 +283,12 @@ where
ret ret
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
self.restart_helper.restart_progress_should_run(state) self.restart_helper.should_restart(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> {
self.restart_helper.clear_restart_progress(state) self.restart_helper.clear_progress(state)
} }
} }
@ -303,9 +299,17 @@ where
M: Mutator<I, <Self as UsesState>::State>, M: Mutator<I, <Self as UsesState>::State>,
Z: Evaluator<E, EM>, Z: Evaluator<E, EM>,
<Self as UsesState>::State: <Self as UsesState>::State:
HasCorpus + HasRand + HasNamedMetadata + HasMetadata + HasExecutions, HasCorpus + HasRand + HasNamedMetadata + HasExecutions + HasMetadata,
I: MutatedTransform<Z::Input, <Self as UsesState>::State> + Clone, I: MutatedTransform<Z::Input, <Self as UsesState>::State> + Clone,
{ {
fn execs_since_progress_start(
&mut self,
state: &mut <Self as UsesState>::State,
) -> Result<u64, Error> {
self.restart_helper
.execs_since_progress_start(state, &self.name)
}
/// Creates a new default tuneable mutational stage /// Creates a new default tuneable mutational stage
#[must_use] #[must_use]
pub fn new(state: &mut <Self as UsesState>::State, mutator: M) -> Self { pub fn new(state: &mut <Self as UsesState>::State, mutator: M) -> Self {

View File

@ -127,13 +127,13 @@ where
} }
#[inline] #[inline]
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
// Stage does not run the target. No reset helper needed. // Stage does not run the target. No reset helper needed.
Ok(true) Ok(true)
} }
#[inline] #[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. // Stage does not run the target. No reset helper needed.
Ok(()) Ok(())
} }

View File

@ -31,7 +31,7 @@
clippy::module_name_repetitions, clippy::module_name_repetitions,
clippy::ptr_cast_constness, clippy::ptr_cast_constness,
clippy::negative_feature_names, clippy::negative_feature_names,
clippy::too_many_lines, clippy::too_many_lines
)] )]
#![cfg_attr(not(test), warn( #![cfg_attr(not(test), warn(
missing_debug_implementations, missing_debug_implementations,

View File

@ -1289,7 +1289,9 @@ where
log::debug!( log::debug!(
"[{} - {:#x}] Send message with id {}", "[{} - {:#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(()) Ok(())
@ -2369,13 +2371,11 @@ impl Brokers {
loop { loop {
self.llmp_brokers.retain_mut(|broker| { self.llmp_brokers.retain_mut(|broker| {
if broker.is_shutting_down() { if broker.is_shutting_down() {
broker.send_buf(LLMP_TAG_EXITING, &[]).expect( broker.send_buf(LLMP_TAG_EXITING, &[]).expect(
"Error when shutting down broker: Could not send LLMP_TAG_EXITING msg.", "Error when shutting down broker: Could not send LLMP_TAG_EXITING msg.",
); );
return false return false;
} }
if current_milliseconds() > end_time { if current_milliseconds() > end_time {

View File

@ -1,11 +1,11 @@
use alloc::borrow::Cow; use alloc::borrow::{Cow, ToOwned};
use core::marker::PhantomData; use core::marker::PhantomData;
use libafl::{ use libafl::{
executors::{Executor, HasObservers}, executors::{Executor, HasObservers},
inputs::{BytesInput, UsesInput}, inputs::{BytesInput, UsesInput},
observers::ObserversTuple, observers::ObserversTuple,
stages::{colorization::TaintMetadata, RetryRestartHelper, Stage}, stages::{colorization::TaintMetadata, Stage, StdRestartHelper},
state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState}, state::{HasCorpus, HasCurrentTestcase, HasExecutions, UsesState},
Error, HasMetadata, HasNamedMetadata, Error, HasMetadata, HasNamedMetadata,
}; };
@ -22,11 +22,14 @@ pub struct AFLppCmplogTracingStage<'a, EM, TE, Z>
where where
TE: UsesState, TE: UsesState,
{ {
name: Cow<'static, str>,
tracer_executor: TE, tracer_executor: TE,
cmplog_observer_handle: Option<Handle<AFLppCmpLogObserver<'a, <Self as UsesState>::State>>>, cmplog_observer_handle: Handle<AFLppCmpLogObserver<'a, <Self as UsesState>::State>>,
#[allow(clippy::type_complexity)] #[allow(clippy::type_complexity)]
phantom: PhantomData<(EM, TE, Z)>, phantom: PhantomData<(EM, TE, Z)>,
} }
/// The name for aflpp tracing stage
pub static AFLPP_CMPLOG_TRACING_STAGE_NAME: &str = "aflpptracing";
impl<EM, TE, Z> UsesState for AFLppCmplogTracingStage<'_, EM, TE, Z> impl<EM, TE, Z> UsesState for AFLppCmplogTracingStage<'_, EM, TE, Z>
where where
@ -40,8 +43,7 @@ where
TE: UsesState, TE: UsesState,
{ {
fn name(&self) -> &Cow<'static, str> { fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("AFLppCmplogTracingStage"); &self.name
&NAME
} }
} }
@ -65,19 +67,17 @@ where
// First run with the un-mutated input // First run with the un-mutated input
let unmutated_input = state.current_input_cloned()?; let unmutated_input = state.current_input_cloned()?;
if let Some(observer_handle) = &self.cmplog_observer_handle { if let Some(ob) = self
if let Some(ob) = self .tracer_executor
.tracer_executor .observers_mut()
.observers_mut() .get_mut(&self.cmplog_observer_handle)
.get_mut(observer_handle) {
{ // This is not the original input,
// This is not the original input, // Set it to false
// Set it to false ob.set_original(true);
ob.set_original(true);
}
// I can't think of any use of this stage if you don't use AFLppCmpLogObserver
// but do nothing ofcourse
} }
// I can't think of any use of this stage if you don't use AFLppCmpLogObserver
// but do nothing ofcourse
self.tracer_executor self.tracer_executor
.observers_mut() .observers_mut()
@ -97,19 +97,17 @@ where
None => return Err(Error::unknown("No metadata found")), None => return Err(Error::unknown("No metadata found")),
}; };
if let Some(observer_handle) = &self.cmplog_observer_handle { if let Some(ob) = self
if let Some(ob) = self .tracer_executor
.tracer_executor .observers_mut()
.observers_mut() .get_mut(&self.cmplog_observer_handle)
.get_mut(observer_handle) {
{ // This is not the original input,
// This is not the original input, // Set it to false
// Set it to false ob.set_original(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
} }
// I can't think of any use of this stage if you don't use AFLppCmpLogObserver
// but do nothing ofcourse
self.tracer_executor self.tracer_executor
.observers_mut() .observers_mut()
@ -126,14 +124,15 @@ where
Ok(()) Ok(())
} }
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> { fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// TODO: this may need better resumption? (Or is it always used with a forkserver?) // Tracing stage is always deterministic
RetryRestartHelper::restart_progress_should_run(state, self, 3) // 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?) // 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 where
TE: UsesState, 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 /// With cmplog observer
pub fn with_cmplog_observer( pub fn new(
tracer_executor: TE, tracer_executor: TE,
observer_handle: Handle<AFLppCmpLogObserver<'a, TE::State>>, observer_handle: Handle<AFLppCmpLogObserver<'a, TE::State>>,
) -> Self { ) -> Self {
let observer_name = observer_handle.name().clone();
Self { 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, tracer_executor,
phantom: PhantomData, phantom: PhantomData,
} }

View File

@ -1,23 +1,20 @@
use std::net::SocketAddr; use std::net::SocketAddr;
use petgraph::{Direction, Graph};
use petgraph::graph::NodeIndex; use petgraph::{graph::NodeIndex, Direction, Graph};
/// A node of the network /// A node of the network
#[derive(Debug, Clone)] #[derive(Debug, Clone, Default)]
pub struct MultiMachineNode { pub struct MultiMachineNode {}
addr: SocketAddr
}
/// The tree /// The tree
pub struct MultiMachineTree { pub struct MultiMachineTree {
pub graph: Graph<MultiMachineNode, ()> pub graph: Graph<MultiMachineNode, ()>,
} }
impl MultiMachineNode { impl MultiMachineNode {
pub fn new(addr: SocketAddr) -> Self { #[must_use]
Self { pub fn new() -> Self {
addr Self {}
}
} }
} }
@ -26,21 +23,19 @@ impl MultiMachineTree {
/// ///
/// ///
/// - machines: machines to add. /// - 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 { pub fn generate(machines: &[SocketAddr], max_children_per_parent: u64) -> Self {
let mut graph = Graph::<MultiMachineNode, ()>::new(); let mut graph = Graph::<MultiMachineNode, ()>::new();
let mut machines = Vec::from(machines); let mut machines = Vec::from(machines);
let root = if let Some(root) = machines.pop() { let root = if let Some(_root) = machines.pop() {
graph.add_node(MultiMachineNode::new(root)) graph.add_node(MultiMachineNode::new())
} else { } else {
return Self { return Self { graph };
graph
};
}; };
let mut graph = Self { graph let mut graph = Self { graph };
};
let mut populate_idx = 0u64; // round-robin population to avoid congestion let mut populate_idx = 0u64; // round-robin population to avoid congestion
let mut nodes_to_populate_now: Vec<NodeIndex> = vec![root]; // current nodes we are working on let mut nodes_to_populate_now: Vec<NodeIndex> = vec![root]; // current nodes we are working on
@ -48,13 +43,18 @@ impl MultiMachineTree {
let mut nodes_to_populate_later: Vec<NodeIndex> = Vec::new(); let mut nodes_to_populate_later: Vec<NodeIndex> = Vec::new();
// place all the machines in the graph // place all the machines in the graph
while let Some(machine) = machines.pop() { while let Some(_machine) = machines.pop() {
if graph.nb_children(nodes_to_populate_now[populate_idx as usize]) == max_children_per_parent { if graph.nb_children(nodes_to_populate_now[populate_idx as usize])
nodes_to_populate_now = nodes_to_populate_later.drain(..).collect(); == max_children_per_parent
{
nodes_to_populate_now = core::mem::take(&mut nodes_to_populate_later);
populate_idx = 0; // should be useless 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); nodes_to_populate_later.push(new_child);
populate_idx = (populate_idx + 1) % nodes_to_populate_now.len() as u64; 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 { 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
} }
} }

View File

@ -6,11 +6,10 @@
//! //!
//! We suppose everyone is on the same network and the machines have the fuzzer ready to run on each machine. //! 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 petgraph::dot::Dot;
use crate::graph::MultiMachineTree; use crate::graph::MultiMachineTree;
pub mod graph; pub mod graph;
@ -41,5 +40,5 @@ fn main() {
let dot = Dot::new(&multi_machine_graph.graph); let dot = Dot::new(&multi_machine_graph.graph);
fs::write("multi_machine.dot", format!("{:?}", dot)).unwrap(); fs::write("multi_machine.dot", format!("{dot:?}")).unwrap();
} }