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
- fuzzers-preflight
strategy:
fail-fast: true
matrix:
os: [ ubuntu-latest ]
fuzzer:

View File

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

View File

@ -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();

View File

@ -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));

View File

@ -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)?;

View File

@ -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)?);
}

View File

@ -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<C, E, O, OT> {
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<bool, Error> {
// 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<bool, Error> {
// Calibration stage disallow restarts
// If a testcase that causes crash/timeout in the queue, we need to remove it from the queue immediately.
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<Observer = C> + 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<Observer = C> + 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
}
}

View File

@ -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<bool, Error> {
// 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<bool, Error> {
// 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,
}
}

View File

@ -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<EM, TE, Z>,
observer_handle: Handle<ConcolicObserver<'a>>,
}
@ -49,10 +48,12 @@ where
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> {
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<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, 3)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// 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<EM, TE, Z>,
observer_handle: Handle<ConcolicObserver<'a>>,
) -> 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<Item = (SymExprRef, SymExpr)>) -> 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<Z> {
/// The helper keeps track of progress for timeouting/restarting targets
restart_helper: ExecutionCountRestartHelper,
name: Cow<'static, str>,
phantom: PhantomData<Z>,
}
@ -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<Z> Named for SimpleConcolicMutationalStage<Z> {
fn name(&self) -> &Cow<'static, str> {
&self.name
}
}
#[cfg(feature = "concolic_mutation")]
impl<E, EM, Z> Stage<E, EM, Z> for SimpleConcolicMutationalStage<Z>
where
@ -373,7 +395,7 @@ where
EM: UsesState<State = Self::State>,
Z: Evaluator<E, EM>,
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<bool, Error> {
self.restart_helper.restart_progress_should_run(state)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// 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<Z> Default for SimpleConcolicMutationalStage<Z> {
fn default() -> Self {
impl<Z> SimpleConcolicMutationalStage<Z> {
#[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,
}
}

View File

@ -115,13 +115,13 @@ where
}
#[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
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(())
}

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`.
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<u8>], 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<C, EM, O, OT, Z> {
name: Cow<'static, str>,
map_observer_handle: Handle<C>,
#[allow(clippy::type_complexity)]
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> {
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<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
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,
}

View File

@ -54,11 +54,13 @@ where
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)
}
fn clear_restart_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
fn clear_progress(&mut self, _state: &mut Self::State) -> Result<(), Error> {
Ok(())
}
}

View File

@ -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<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
impl NestedStageStdRestartHelper {
fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
where
S: HasNestedStageStatus,
{
@ -21,7 +21,7 @@ impl NestedStageRestartHelper {
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
S: HasNestedStageStatus,
{
@ -70,12 +70,12 @@ where
Ok(())
}
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<bool, Error> {
NestedStageRestartHelper::restart_progress_should_run(state, self)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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)
}
}

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.
*/
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<bool, Error>;
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error>;
/// 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<CB, E, EM, Z> {
name: Cow<'static, str>,
closure: CB,
phantom: PhantomData<(E, EM, Z)>,
}
@ -299,8 +310,7 @@ where
impl<CB, E, EM, Z> Named for ClosureStage<CB, E, EM, Z> {
fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("<unnamed fn>");
&NAME
&self.name
}
}
@ -323,14 +333,15 @@ where
}
#[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// 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<bool, Error> {
// 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<CB, E, EM, Z> ClosureStage<CB, E, EM, Z> {
/// 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<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`]
#[allow(clippy::type_complexity)]
#[derive(Debug)]
pub struct PushStageAdapter<CS, EM, OT, PS, Z> {
name: Cow<'static, str>,
push_stage: PS,
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`]
#[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<CS, EM, OT, PS, Z> UsesState for PushStageAdapter<CS, EM, OT, PS, Z>
where
@ -384,11 +405,23 @@ where
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>
where
CS: Scheduler,
Self::State:
HasExecutions + HasMetadata + HasRand + HasCorpus + HasLastReportTime + HasCurrentCorpusId,
Self::State: HasExecutions
+ HasRand
+ HasCorpus
+ HasLastReportTime
+ HasCurrentCorpusId
+ HasNamedMetadata
+ HasMetadata,
E: Executor<EM, Z> + HasObservers<Observers = OT, State = Self::State>,
EM: EventFirer<State = Self::State>
+ EventRestarter
@ -446,48 +479,51 @@ where
}
#[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...
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<usize>,
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
///
/// Returns `true` if the stage should run
pub fn restart_progress_should_run<S, ST>(
state: &mut S,
stage: &ST,
max_retries: usize,
) -> Result<bool, Error>
pub fn should_restart<S>(state: &mut S, name: &str, max_retries: usize) -> Result<bool, Error>
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<S, ST>(state: &mut S, stage: &ST) -> Result<(), Error>
pub fn clear_progress<S>(state: &mut S, name: &str) -> Result<(), Error>
where
S: HasNamedMetadata,
ST: Named,
{
state
.named_metadata_mut::<Self>(stage.name())?
.tries_remaining = None;
state.named_metadata_mut::<Self>(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<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
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::<ExecutionCountRestartHelperMetadata>()
.named_metadata::<ExecutionCountRestartHelperMetadata>(name)
.map(|x| {
self.started_at_execs = Some(x.started_at_execs);
x.started_at_execs
@ -620,12 +653,13 @@ impl ExecutionCountRestartHelper {
}
/// 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
S: HasMetadata + HasExecutions,
S: HasNamedMetadata + HasExecutions,
{
let executions = *state.executions();
let metadata = state.metadata_or_insert_with(|| ExecutionCountRestartHelperMetadata {
let metadata =
state.named_metadata_or_insert_with(name, || ExecutionCountRestartHelperMetadata {
started_at_execs: executions,
});
self.started_at_execs = Some(metadata.started_at_execs);
@ -633,13 +667,13 @@ impl ExecutionCountRestartHelper {
}
/// 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
S: HasMetadata,
{
self.started_at_execs = None;
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(())
}
}
@ -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<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
fn should_restart<S, ST>(state: &mut S, _stage: &ST) -> Result<bool, Error>
where
S: HasMetadata,
{
@ -690,7 +724,7 @@ pub mod test {
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
S: HasMetadata,
{
@ -727,12 +761,12 @@ pub mod test {
Ok(())
}
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
TestProgress::restart_progress_should_run(state, self)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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(())
}

View File

@ -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<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
#[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<E, EM, I, M, Z> {
/// 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<State = Self::State>,
M: Mutator<I, Self::State>,
Z: Evaluator<E, EM>,
Self::State: HasCorpus + HasRand + HasExecutions + HasMetadata,
Self::State: HasCorpus + HasRand + HasExecutions + HasMetadata + HasNamedMetadata,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
/// The mutator, added to this stage
@ -190,12 +190,13 @@ where
fn iterations(&self, state: &mut Self::State) -> Result<usize, Error> {
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>
where
Z: UsesState,
@ -203,13 +204,19 @@ where
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>
where
E: UsesState<State = Self::State>,
EM: UsesState<State = Self::State>,
M: Mutator<I, Self::State>,
Z: Evaluator<E, EM>,
Self::State: HasCorpus + HasRand + HasMetadata + HasExecutions,
Self::State: HasCorpus + HasRand + HasMetadata + HasExecutions + HasNamedMetadata,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
#[inline]
@ -229,14 +236,12 @@ where
ret
}
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
Ok(true)
// self.restart_helper.restart_progress_should_run(state)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<E, EM, I, M, Z> {
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<E, EM, I, M, Z> UsesState for MultiMutationalStage<E, EM, I, M, Z>
where
Z: UsesState,
@ -300,8 +319,7 @@ where
impl<E, EM, I, M, Z> Named for MultiMutationalStage<E, EM, I, M, Z> {
fn name(&self) -> &Cow<'static, str> {
static NAME: Cow<'static, str> = Cow::Borrowed("MultiMutational");
&NAME
&self.name
}
}
@ -315,15 +333,14 @@ where
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
#[inline]
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// 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<bool, Error> {
// 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<E, EM, I, M, Z> MultiMutationalStage<E, EM, I, M, Z> {
/// 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,
}

View File

@ -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<E, F, EM, I, M, Z> {
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<State = Self::State>,
F: TestcaseScore<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>,
I: MutatedTransform<E::Input, Self::State> + Clone,
{
@ -72,10 +76,6 @@ where
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>
@ -84,7 +84,7 @@ where
EM: UsesState<State = Self::State>,
F: TestcaseScore<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>,
I: MutatedTransform<Self::Input, Self::State> + Clone,
{
@ -101,14 +101,13 @@ where
ret
}
fn restart_progress_should_run(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
Ok(true)
// self.restart_helper.restart_progress_should_run(state)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
// 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(),
}
}
}

View File

@ -126,13 +126,13 @@ where
}
#[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
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(())
}

View File

@ -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<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
// 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<CB, E, EM, Z> SyncFromDiskStage<CB, E, EM, Z> {
/// 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<bool, Error> {
fn should_restart(&mut self, _state: &mut Self::State) -> Result<bool, Error> {
// 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(())
}

View File

@ -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<E, EM, F, FF, IP, M, Z> {
/// The name
name: Cow<'static, str>,
/// The mutator(s) this stage uses
mutator: M,
/// The factory
@ -231,16 +243,17 @@ where
FF: FeedbackFactory<F, E::Observers>,
F: Feedback<Self::State>,
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>,
IP: MutatedTransformPost<Self::State> + Clone,
{
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
self.restart_helper.restart_progress_should_run(state)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<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>
for StdTMinMutationalStage<E, EM, F, FF, IP, M, Z>
where
@ -280,7 +304,8 @@ where
FF: FeedbackFactory<F, E::Observers>,
F: Feedback<Self::State>,
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>,
IP: MutatedTransformPost<Self::State> + Clone,
{
@ -302,14 +327,22 @@ where
}
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> {
/// 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,

View File

@ -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<EM, TE, Z> {
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<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<EM, TE, Z> Named for TracingStage<EM, TE, Z> {
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<EM, TE, Z> TracingStage<EM, TE, Z> {
/// 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<EM, TE, Z> TracingStage<EM, TE, Z> {
/// A stage that runs the shadow executor using also the shadow observers
#[derive(Clone, Debug)]
pub struct ShadowTracingStage<E, EM, SOT, Z> {
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<E, EM, SOT, Z> Named for ShadowTracingStage<E, EM, SOT, Z>
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<bool, Error> {
RetryRestartHelper::restart_progress_should_run(state, self, self.max_retries)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<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 {
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
}
}

View File

@ -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<u64, Error> {
self.restart_helper.execs_since_progress_start(state)
}
}
impl<E, EM, I, M, Z> UsesState for TuneableMutationalStage<E, EM, I, M, Z>
@ -287,12 +283,12 @@ where
ret
}
fn restart_progress_should_run(&mut self, state: &mut Self::State) -> Result<bool, Error> {
self.restart_helper.restart_progress_should_run(state)
fn should_restart(&mut self, state: &mut Self::State) -> Result<bool, Error> {
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<I, <Self as UsesState>::State>,
Z: Evaluator<E, EM>,
<Self as UsesState>::State:
HasCorpus + HasRand + HasNamedMetadata + HasMetadata + HasExecutions,
HasCorpus + HasRand + HasNamedMetadata + HasExecutions + HasMetadata,
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
#[must_use]
pub fn new(state: &mut <Self as UsesState>::State, mutator: M) -> Self {

View File

@ -127,13 +127,13 @@ where
}
#[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.
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(())
}

View File

@ -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,

View File

@ -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 {

View File

@ -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<Handle<AFLppCmpLogObserver<'a, <Self as UsesState>::State>>>,
cmplog_observer_handle: Handle<AFLppCmpLogObserver<'a, <Self as UsesState>::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<EM, TE, Z> 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,11 +67,10 @@ 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)
.get_mut(&self.cmplog_observer_handle)
{
// This is not the original input,
// Set it to false
@ -77,7 +78,6 @@ where
}
// 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,11 +97,10 @@ 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)
.get_mut(&self.cmplog_observer_handle)
{
// This is not the original input,
// Set it to false
@ -109,7 +108,6 @@ where
}
// 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<bool, Error> {
// 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<bool, Error> {
// 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<AFLppCmpLogObserver<'a, TE::State>>,
) -> 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,
}

View File

@ -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<MultiMachineNode, ()>
pub graph: Graph<MultiMachineNode, ()>,
}
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::<MultiMachineNode, ()>::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<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();
// 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
}
}

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.
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();
}