diff --git a/libafl/src/stages/logics.rs b/libafl/src/stages/logics.rs index 9de107bd2a..5f431fca38 100644 --- a/libafl/src/stages/logics.rs +++ b/libafl/src/stages/logics.rs @@ -3,7 +3,7 @@ use core::marker::PhantomData; use crate::{ - stages::{HasCurrentStage, HasNestedStageStatus, Stage, StagesTuple}, + stages::{HasCurrentStage, HasNestedStageStatus, Stage, StageId, StagesTuple}, state::UsesState, Error, }; @@ -72,7 +72,8 @@ where state: &mut E::State, manager: &mut EM, ) -> Result<(), Error> { - while state.current_stage()?.is_some() || (self.closure)(fuzzer, executor, state, manager)? + while state.current_stage_idx()?.is_some() + || (self.closure)(fuzzer, executor, state, manager)? { self.stages.perform_all(fuzzer, executor, state, manager)?; } @@ -150,7 +151,8 @@ where state: &mut E::State, manager: &mut EM, ) -> Result<(), Error> { - if state.current_stage()?.is_some() || (self.closure)(fuzzer, executor, state, manager)? { + if state.current_stage_idx()?.is_some() || (self.closure)(fuzzer, executor, state, manager)? + { self.if_stages .perform_all(fuzzer, executor, state, manager)?; } @@ -231,21 +233,21 @@ where state: &mut E::State, manager: &mut EM, ) -> Result<(), Error> { - let current = state.current_stage()?; + let current = state.current_stage_idx()?; let fresh = current.is_none(); let closure_return = fresh && (self.closure)(fuzzer, executor, state, manager)?; - if current == Some(0) || closure_return { + if current == Some(StageId(0)) || closure_return { if fresh { - state.set_stage(0)?; + state.set_current_stage_idx(StageId(0))?; } state.enter_inner_stage()?; self.if_stages .perform_all(fuzzer, executor, state, manager)?; } else { if fresh { - state.set_stage(1)?; + state.set_current_stage_idx(StageId(1))?; } state.enter_inner_stage()?; self.else_stages diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 1b6813a9d8..96638b3fb3 100644 --- a/libafl/src/stages/mod.rs +++ b/libafl/src/stages/mod.rs @@ -5,7 +5,7 @@ Other stages may enrich [`crate::corpus::Testcase`]s with metadata. */ use alloc::{borrow::Cow, boxed::Box, vec::Vec}; -use core::marker::PhantomData; +use core::{fmt, marker::PhantomData}; pub use calibrate::CalibrationStage; pub use colorization::*; @@ -153,7 +153,7 @@ where stage: &mut S, _: &mut EM, ) -> Result<(), Error> { - if stage.current_stage()?.is_some() { + if stage.current_stage_idx()?.is_some() { Err(Error::illegal_state( "Got to the end of the tuple without completing resume.", )) @@ -179,11 +179,11 @@ where state: &mut Head::State, manager: &mut EM, ) -> Result<(), Error> { - match state.current_stage()? { - Some(idx) if idx < Self::LEN => { + match state.current_stage_idx()? { + Some(idx) if idx < StageId(Self::LEN) => { // do nothing; we are resuming } - Some(idx) if idx == Self::LEN => { + Some(idx) if idx == StageId(Self::LEN) => { // perform the stage, but don't set it let stage = &mut self.0; @@ -191,12 +191,12 @@ where state.clear_stage()?; } - Some(idx) if idx > Self::LEN => { + Some(idx) if idx > StageId(Self::LEN) => { unreachable!("We should clear the stage index before we get here..."); } // this is None, but the match can't deduce that _ => { - state.set_stage(Self::LEN)?; + state.set_current_stage_idx(StageId(Self::LEN))?; let stage = &mut self.0; stage.perform_restartable(fuzzer, executor, state, manager)?; @@ -540,16 +540,27 @@ impl RetryRestartHelper { } } +/// The index of a stage +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] +#[repr(transparent)] +pub struct StageId(pub(crate) usize); + +impl fmt::Display for StageId { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", self.0) + } +} + /// Trait for types which track the current stage pub trait HasCurrentStage { /// Set the current stage; we have started processing this stage - fn set_stage(&mut self, idx: usize) -> Result<(), Error>; + fn set_current_stage_idx(&mut self, idx: StageId) -> Result<(), Error>; /// Clear the current stage; we are done processing this stage fn clear_stage(&mut self) -> Result<(), Error>; /// Fetch the current stage -- typically used after a state recovery or transfer - fn current_stage(&self) -> Result, Error>; + fn current_stage_idx(&self) -> Result, Error>; /// Notify of a reset from which we may recover fn on_restart(&mut self) -> Result<(), Error> { diff --git a/libafl/src/state/mod.rs b/libafl/src/state/mod.rs index eeed0de554..1697b08ed7 100644 --- a/libafl/src/state/mod.rs +++ b/libafl/src/state/mod.rs @@ -1,5 +1,6 @@ //! The fuzzer, and state are the core pieces of every good fuzzer +#[cfg(feature = "std")] use alloc::vec::Vec; use core::{ borrow::BorrowMut, @@ -22,6 +23,9 @@ use libafl_bolts::{ }; use serde::{de::DeserializeOwned, Deserialize, Serialize}; +mod stack; +pub use stack::StageStack; + #[cfg(feature = "introspection")] use crate::monitors::ClientPerfMonitor; #[cfg(feature = "scalability_introspection")] @@ -33,7 +37,7 @@ use crate::{ fuzzer::{Evaluator, ExecuteInputResult}, generators::Generator, inputs::{Input, UsesInput}, - stages::{HasCurrentStage, HasNestedStageStatus}, + stages::{HasCurrentStage, HasNestedStageStatus, StageId}, Error, HasMetadata, HasNamedMetadata, }; @@ -257,10 +261,7 @@ pub struct StdState { last_report_time: Option, /// The current index of the corpus; used to record for resumable fuzzing. corpus_idx: Option, - /// The stage indexes for each nesting of stages - stage_idx_stack: Vec, - /// The current stage depth - stage_depth: usize, + stage_stack: StageStack, phantom: PhantomData, } @@ -532,41 +533,30 @@ where } impl HasCurrentStage for StdState { - fn set_stage(&mut self, idx: usize) -> Result<(), Error> { - // ensure we are in the right frame - if self.stage_depth != self.stage_idx_stack.len() { - return Err(Error::illegal_state( - "stage not resumed before setting stage", - )); - } - self.stage_idx_stack.push(idx); - Ok(()) + fn set_current_stage_idx(&mut self, idx: StageId) -> Result<(), Error> { + self.stage_stack.set_current_stage_idx(idx) } fn clear_stage(&mut self) -> Result<(), Error> { - self.stage_idx_stack.truncate(self.stage_depth); - Ok(()) + self.stage_stack.clear_stage() } - fn current_stage(&self) -> Result, Error> { - Ok(self.stage_idx_stack.get(self.stage_depth).copied()) + fn current_stage_idx(&self) -> Result, Error> { + self.stage_stack.current_stage_idx() } fn on_restart(&mut self) -> Result<(), Error> { - self.stage_depth = 0; // reset the stage depth so that we may resume inward - Ok(()) + self.stage_stack.on_restart() } } impl HasNestedStageStatus for StdState { fn enter_inner_stage(&mut self) -> Result<(), Error> { - self.stage_depth += 1; - Ok(()) + self.stage_stack.enter_inner_stage() } fn exit_inner_stage(&mut self) -> Result<(), Error> { - self.stage_depth -= 1; - Ok(()) + self.stage_stack.exit_inner_stage() } } @@ -1107,8 +1097,7 @@ where dont_reenter: None, last_report_time: None, corpus_idx: None, - stage_depth: 0, - stage_idx_stack: Vec::new(), + stage_stack: StageStack::default(), phantom: PhantomData, #[cfg(feature = "std")] multicore_inputs_processed: None, @@ -1239,7 +1228,7 @@ impl HasCurrentCorpusId for NopState { } impl HasCurrentStage for NopState { - fn set_stage(&mut self, _idx: usize) -> Result<(), Error> { + fn set_current_stage_idx(&mut self, _idx: StageId) -> Result<(), Error> { Ok(()) } @@ -1247,7 +1236,7 @@ impl HasCurrentStage for NopState { Ok(()) } - fn current_stage(&self) -> Result, Error> { + fn current_stage_idx(&self) -> Result, Error> { Ok(None) } } diff --git a/libafl/src/state/stack.rs b/libafl/src/state/stack.rs new file mode 100644 index 0000000000..72c5828df5 --- /dev/null +++ b/libafl/src/state/stack.rs @@ -0,0 +1,54 @@ +use alloc::vec::Vec; + +use libafl_bolts::Error; +use serde::{Deserialize, Serialize}; + +use crate::stages::{HasCurrentStage, HasNestedStageStatus, StageId}; + +/// A stack to keep track of which stage is executing +#[derive(Serialize, Deserialize, Clone, Debug, Default)] +pub struct StageStack { + /// The stage indexes for each nesting of stages + stage_idx_stack: Vec, + /// The current stage depth + stage_depth: usize, +} + +impl HasCurrentStage for StageStack { + fn set_current_stage_idx(&mut self, idx: StageId) -> Result<(), Error> { + // ensure we are in the right frame + if self.stage_depth != self.stage_idx_stack.len() { + return Err(Error::illegal_state( + "stage not resumed before setting stage", + )); + } + self.stage_idx_stack.push(idx); + Ok(()) + } + + fn clear_stage(&mut self) -> Result<(), Error> { + self.stage_idx_stack.truncate(self.stage_depth); + Ok(()) + } + + fn current_stage_idx(&self) -> Result, Error> { + Ok(self.stage_idx_stack.get(self.stage_depth).copied()) + } + + fn on_restart(&mut self) -> Result<(), Error> { + self.stage_depth = 0; // reset the stage depth so that we may resume inward + Ok(()) + } +} + +impl HasNestedStageStatus for StageStack { + fn enter_inner_stage(&mut self) -> Result<(), Error> { + self.stage_depth += 1; + Ok(()) + } + + fn exit_inner_stage(&mut self) -> Result<(), Error> { + self.stage_depth -= 1; + Ok(()) + } +}