Introduce StageId (extends #2201) (#2239)

* Factor out `StageStack` from `StdState`

Refactoring for the sake of abstraction and maintainability.

Previously, the `stage_idx_stack` and `stage_depth` fields of
`StdState` had an implicit relationship to one another. This commit
reifies this relationship into a new `struct`. Its fields are hidden
from everything else in `StdState`. It `impl`s `HasCurrentStage` and
`HasNestedStageStatus`, which `impl`s now no longer have mutable access
to the whole state. The `impl`s for `StdState` just forward to these
ones.

Hopefully, this refactoring would make it easier to `impl State` for
other types in the future.

* Create `StageIdx` newtype

Like `CorpusId`, this is just a thin wrapper around `usize`.

* unused import

* Rename a few functions to be more about indices

* Rename Index to Id, fmt

* Move StageId around

---------

Co-authored-by: Langston Barrett <langston.barrett@gmail.com>

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

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()? {
+        match state.current_stage_idx()? {
-            Some(idx) if idx < Self::LEN => {
+            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<Option<usize>, Error>;
+    fn current_stage_idx(&self) -> Result<Option<StageId>, Error>;
 
     /// Notify of a reset from which we may recover
     fn on_restart(&mut self) -> Result<(), Error> {

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<I, C, R, SC> {
     last_report_time: Option<Duration>,
     /// The current index of the corpus; used to record for resumable fuzzing.
     corpus_idx: Option<CorpusId>,
-    /// The stage indexes for each nesting of stages
+    stage_stack: StageStack,
-    stage_idx_stack: Vec<usize>,
-    /// The current stage depth
-    stage_depth: usize,
     phantom: PhantomData<I>,
 }
 
@@ -532,41 +533,30 @@ where
 }
 
 impl<I, C, R, SC> HasCurrentStage for StdState<I, C, R, SC> {
-    fn set_stage(&mut self, idx: usize) -> Result<(), Error> {
+    fn set_current_stage_idx(&mut self, idx: StageId) -> Result<(), Error> {
-        // ensure we are in the right frame
+        self.stage_stack.set_current_stage_idx(idx)
-        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);
+        self.stage_stack.clear_stage()
-        Ok(())
     }
 
-    fn current_stage(&self) -> Result<Option<usize>, Error> {
+    fn current_stage_idx(&self) -> Result<Option<StageId>, Error> {
-        Ok(self.stage_idx_stack.get(self.stage_depth).copied())
+        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
+        self.stage_stack.on_restart()
-        Ok(())
     }
 }
 
 impl<I, C, R, SC> HasNestedStageStatus for StdState<I, C, R, SC> {
     fn enter_inner_stage(&mut self) -> Result<(), Error> {
-        self.stage_depth += 1;
+        self.stage_stack.enter_inner_stage()
-        Ok(())
     }
 
     fn exit_inner_stage(&mut self) -> Result<(), Error> {
-        self.stage_depth -= 1;
+        self.stage_stack.exit_inner_stage()
-        Ok(())
     }
 }
 
@@ -1107,8 +1097,7 @@ where
             dont_reenter: None,
             last_report_time: None,
             corpus_idx: None,
-            stage_depth: 0,
+            stage_stack: StageStack::default(),
-            stage_idx_stack: Vec::new(),
             phantom: PhantomData,
             #[cfg(feature = "std")]
             multicore_inputs_processed: None,
@@ -1239,7 +1228,7 @@ impl<I> HasCurrentCorpusId for NopState<I> {
 }
 
 impl<I> HasCurrentStage for NopState<I> {
-    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<I> HasCurrentStage for NopState<I> {
         Ok(())
     }
 
-    fn current_stage(&self) -> Result<Option<usize>, Error> {
+    fn current_stage_idx(&self) -> Result<Option<StageId>, Error> {
         Ok(None)
     }
 }

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<StageId>,
+    /// 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<Option<StageId>, 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(())
+    }
+}