Replay stage (#3003)

* replay stage

* seems to be working

* a

* rever

* rev

* rdy for merge

* fmt

* lol

fuzzers/baby/tutorial/src/lib.rs

@@ -186,7 +186,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.mgr_on_restart(&mut state)?;
+    restarting_mgr.on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/binary_only/qemu_launcher/Justfile

@@ -55,7 +55,7 @@ test_inner: harness build
     ./tests/injection/test.sh || exit 1
 
     # complie again with simple mgr
-    cargo build --profile={{PROFILE}} --features="simplemgr,{{ARCH}}" --target-dir={{ TARGET_DIR }}
+    cargo build --profile={{PROFILE}} --features="simplemgr,{{ARCH}}" --target-dir={{ TARGET_DIR }} || exit 1
     ./tests/qasan/test.sh || exit 1
 
 [unix]

fuzzers/binary_only/qemu_launcher/src/instance.rs

@@ -469,7 +469,7 @@ impl<M: Monitor> Instance<'_, M> {
 
                 // It's important, that we store the state before restarting!
                 // Else, the parent will not respawn a new child and quit.
-                self.mgr.mgr_on_restart(state)?;
+                self.mgr.on_restart(state)?;
             }
         } else {
             fuzzer.fuzz_loop(stages, executor, state, &mut self.mgr)?;

fuzzers/full_system/nyx_launcher/src/instance.rs

@@ -250,7 +250,7 @@ impl<M: Monitor> Instance<'_, M> {
 
             // It's important, that we store the state before restarting!
             // Else, the parent will not respawn a new child and quit.
-            self.mgr.mgr_on_restart(state)?;
+            self.mgr.on_restart(state)?;
         } else {
             fuzzer.fuzz_loop(stages, executor, state, &mut self.mgr)?;
         }

fuzzers/inprocess/fuzzbench/Justfile

@@ -40,12 +40,13 @@ fuzzer:
 
 [linux]
 [macos]
-run: cxx fuzz_o
+run: cxx fuzz_o fuzzer
 	#!/bin/bash
 	rm -rf libafl_unix_shmem_server || true
 	mkdir in || true
 	echo a > in/a
 	./{{FUZZER_NAME}} -o out -i in
+	RUST_LOG=info ./{{FUZZER_NAME}} -o out -i seed
 
 [windows]
 run:

fuzzers/inprocess/fuzzbench_ctx/Justfile

@@ -40,7 +40,7 @@ fuzzer:
 
 [linux]
 [macos]
-run: cxx fuzz_o
+run: cxx fuzz_o fuzzer
 	#!/bin/bash
 	rm -rf libafl_unix_shmem_server || true
 	mkdir in || true

fuzzers/inprocess/libfuzzer_libpng/src/lib.rs

@@ -222,7 +222,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.mgr_on_restart(&mut state)?;
+    restarting_mgr.on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/inprocess/libfuzzer_libpng_cmin/src/lib.rs

@@ -236,7 +236,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.mgr_on_restart(&mut state)?;
+    restarting_mgr.on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/inprocess/libfuzzer_libpng_norestart/src/lib.rs

@@ -274,7 +274,7 @@ pub extern "C" fn libafl_main() {
             &mut restarting_mgr,
             opt.loop_iters,
         )?;
-        restarting_mgr.mgr_on_restart(&mut state)?;
+        restarting_mgr.on_restart(&mut state)?;
 
         Ok(())
     };

fuzzers/inprocess/libfuzzer_libpng_tcp_manager/src/lib.rs

@@ -219,7 +219,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.mgr_on_restart(&mut state)?;
+    restarting_mgr.on_restart(&mut state)?;
 
     Ok(())
 }

fuzzers/inprocess/libfuzzer_windows_asan/src/lib.rs

@@ -184,7 +184,7 @@ fn fuzz(corpus_dirs: &[PathBuf], objective_dir: PathBuf, broker_port: u16) -> Re
 
     // It's important, that we store the state before restarting!
     // Else, the parent will not respawn a new child and quit.
-    restarting_mgr.mgr_on_restart(&mut state)?;
+    restarting_mgr.on_restart(&mut state)?;
 
     Ok(())
 }

libafl/src/events/centralized.rs

@@ -272,9 +272,9 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     #[inline]
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         self.client.await_safe_to_unmap_blocking();
-        self.inner.mgr_on_restart(state)?;
+        self.inner.on_restart(state)?;
         Ok(())
     }
 }

libafl/src/events/llmp/restarting.rs

@@ -256,7 +256,7 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     /// Reset the single page (we reuse it over and over from pos 0), then send the current state to the next runner.
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         state.on_restart()?;
 
         if let Some(sr) = &mut self.staterestorer {

libafl/src/events/mod.rs

@@ -559,10 +559,10 @@ pub trait EventRestarter<S> {
     /// You *must* ensure that [`HasCurrentStageId::on_restart`] will be invoked in this method, by you
     /// or an internal [`EventRestarter`], before the state is saved for recovery.
     /// [`std_on_restart`] is the standard implementation that you can call.
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error>;
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error>;
 }
 
-/// Default implementation of [`EventRestarter::mgr_on_restart`] for implementors with the given
+/// Default implementation of [`EventRestarter::on_restart`] for implementors with the given
 /// constraints
 pub fn std_on_restart<EM, S>(restarter: &mut EM, state: &mut S) -> Result<(), Error>
 where
@@ -644,7 +644,7 @@ impl<S> EventRestarter<S> for NopEventManager
 where
     S: HasCurrentStageId,
 {
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         std_on_restart(self, state)
     }
 }
@@ -769,8 +769,8 @@ where
     EM: EventRestarter<S>,
 {
     #[inline]
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
-        self.inner.mgr_on_restart(state)
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+        self.inner.on_restart(state)
     }
 }
 

libafl/src/events/simple.rs

@@ -112,7 +112,7 @@ impl<I, MT, S> EventRestarter<S> for SimpleEventManager<I, MT, S>
 where
     S: HasCurrentStageId,
 {
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         std_on_restart(self, state)
     }
 }
@@ -326,7 +326,7 @@ where
     MT: Monitor,
 {
     /// Reset the single page (we reuse it over and over from pos 0), then send the current state to the next runner.
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         state.on_restart()?;
 
         // First, reset the page to 0 so the next iteration can read read from the beginning of this page

libafl/src/events/tcp.rs

@@ -626,7 +626,7 @@ impl<EMH, I, S> EventRestarter<S> for TcpEventManager<EMH, I, S>
 where
     S: HasCurrentStageId,
 {
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         std_on_restart(self, state)
     }
 }
@@ -860,7 +860,7 @@ where
     SP: ShMemProvider<ShMem = SHM>,
 {
     /// Reset the single page (we reuse it over and over from pos 0), then send the current state to the next runner.
-    fn mgr_on_restart(&mut self, state: &mut S) -> Result<(), Error> {
+    fn on_restart(&mut self, state: &mut S) -> Result<(), Error> {
         state.on_restart()?;
 
         // First, reset the page to 0 so the next iteration can read read from the beginning of this page

libafl/src/executors/inprocess/mod.rs

@@ -430,7 +430,7 @@ pub fn run_observers_and_save_state<E, EM, I, OF, S, Z>(
     }
 
     // Serialize the state and wait safely for the broker to read pending messages
-    event_mgr.mgr_on_restart(state).unwrap();
+    event_mgr.on_restart(state).unwrap();
 
     log::info!("Bye!");
 }

libafl/src/fuzzer/mod.rs

@@ -207,7 +207,7 @@ pub trait Fuzzer<E, EM, I, S, ST> {
     /// because each stage could run the harness for multiple times)
     ///
     /// If you use this fn in a restarting scenario to only run for `n` iterations,
-    /// before exiting, make sure you call `event_mgr.mgr_on_restart(&mut state)?;`.
+    /// before exiting, make sure you call `event_mgr.on_restart(&mut state)?;`.
     /// This way, the state will be available in the next, respawned, iteration.
     fn fuzz_one(
         &mut self,
@@ -233,7 +233,7 @@ pub trait Fuzzer<E, EM, I, S, ST> {
     /// because each stage could run the harness for multiple times)
     ///
     /// If you use this fn in a restarting scenario to only run for `n` iterations,
-    /// before exiting, make sure you call `event_mgr.mgr_on_restart(&mut state)?;`.
+    /// before exiting, make sure you call `event_mgr.on_restart(&mut state)?;`.
     /// This way, the state will be available in the next, respawned, iteration.
     fn fuzz_loop_for(
         &mut self,

libafl/src/stages/mod.rs

@@ -58,6 +58,9 @@ pub mod mutational;
 pub mod push;
 pub mod tmin;
 
+pub mod replay;
+pub use replay::*;
+
 #[cfg(feature = "std")]
 pub mod afl_stats;
 pub mod calibrate;

libafl/src/stages/replay.rs

@@ -0,0 +1,171 @@
+//! The replay stage can scan all inputs and executes them once per input
+
+use alloc::{
+    borrow::{Cow, ToOwned},
+    string::ToString,
+    vec::Vec,
+};
+use core::marker::PhantomData;
+
+use hashbrown::HashSet;
+use libafl_bolts::{impl_serdeany, Named};
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    corpus::{Corpus, CorpusId},
+    stages::Stage,
+    state::{HasCorpus, HasSolutions},
+    Error, Evaluator,
+};
+
+/// Replay all inputs
+#[derive(Debug)]
+pub struct ReplayStage<I> {
+    name: Cow<'static, str>,
+    restart_helper: ReplayRestartingHelper,
+    phantom: PhantomData<I>,
+}
+
+impl<I> Default for ReplayStage<I> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<I> Named for ReplayStage<I> {
+    fn name(&self) -> &Cow<'static, str> {
+        &self.name
+    }
+}
+
+/// Restart helper for replay stage
+#[derive(Serialize, Deserialize, Debug, Clone, Default)]
+pub struct ReplayRestartingHelper {
+    done_corpus: HashSet<CorpusId>,
+    done_solution: HashSet<CorpusId>,
+}
+
+impl ReplayRestartingHelper {
+    /// constructor
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            done_corpus: HashSet::default(),
+            done_solution: HashSet::default(),
+        }
+    }
+
+    /// clear history
+    pub fn clear(&mut self) {
+        self.done_corpus.clear();
+        self.done_solution.clear();
+    }
+
+    /// check we've scaned this corpus entry
+    pub fn corpus_probe(&mut self, id: &CorpusId) -> bool {
+        self.done_corpus.contains(id)
+    }
+
+    /// check we've scaned this solution entry
+    pub fn solution_probe(&mut self, id: &CorpusId) -> bool {
+        self.done_solution.contains(id)
+    }
+
+    /// mark this corpus entry as finished
+    pub fn corpus_finish(&mut self, id: CorpusId) {
+        self.done_corpus.insert(id);
+    }
+
+    /// mark this solution entry as finished
+    pub fn solution_finish(&mut self, id: CorpusId) {
+        self.done_solution.insert(id);
+    }
+}
+
+impl_serdeany!(ReplayRestartingHelper);
+
+/// The counter for giving this stage unique id
+static mut REPLAY_STAGE_ID: usize = 0;
+/// The name for tracing stage
+pub static REPLAY_STAGE_NAME: &str = "tracing";
+
+impl<I> ReplayStage<I> {
+    #[must_use]
+    /// Create a new replay stage
+    pub fn new() -> Self {
+        // unsafe but impossible that you create two threads both instantiating this instance
+        let stage_id = unsafe {
+            let ret = REPLAY_STAGE_ID;
+            REPLAY_STAGE_ID += 1;
+            ret
+        };
+
+        Self {
+            name: Cow::Owned(REPLAY_STAGE_NAME.to_owned() + ":" + stage_id.to_string().as_ref()),
+            restart_helper: ReplayRestartingHelper::new(),
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<E, EM, I, S, Z> Stage<E, EM, S, Z> for ReplayStage<I>
+where
+    S: HasCorpus<I> + HasSolutions<I>,
+    Z: Evaluator<E, EM, I, S>,
+    I: Clone,
+{
+    fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
+        Ok(true)
+    }
+
+    fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
+        self.restart_helper.clear();
+        Ok(())
+    }
+
+    #[inline]
+    fn perform(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+    ) -> Result<(), Error> {
+        let corpus_ids: Vec<CorpusId> = state.corpus().ids().collect();
+
+        for id in corpus_ids {
+            if self.restart_helper.corpus_probe(&id) {
+                continue;
+            }
+            log::info!("Replaying corpus: {id}");
+            let input = {
+                let mut tc = state.corpus().get(id)?.borrow_mut();
+                let input = tc.load_input(state.corpus())?;
+                input.clone()
+            };
+
+            fuzzer.evaluate_input(state, executor, manager, &input)?;
+
+            self.restart_helper.corpus_finish(id);
+        }
+
+        let solution_ids: Vec<CorpusId> = state.solutions().ids().collect();
+        for id in solution_ids {
+            if self.restart_helper.solution_probe(&id) {
+                continue;
+            }
+            log::info!("Replaying solution: {id}");
+            let input = {
+                let mut tc = state.corpus().get(id)?.borrow_mut();
+                let input = tc.load_input(state.corpus())?;
+                input.clone()
+            };
+
+            fuzzer.evaluate_input(state, executor, manager, &input)?;
+
+            self.restart_helper.solution_finish(id);
+        }
+
+        Ok(())
+    }
+}

libafl_sugar/src/forkserver.rs

@@ -257,7 +257,7 @@ impl ForkserverBytesCoverageSugar<'_> {
                         &mut mgr,
                         iters,
                     )?;
-                    mgr.mgr_on_restart(&mut state)?;
+                    mgr.on_restart(&mut state)?;
                     std::process::exit(0);
                 } else {
                     fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -278,7 +278,7 @@ impl ForkserverBytesCoverageSugar<'_> {
                         &mut mgr,
                         iters,
                     )?;
-                    mgr.mgr_on_restart(&mut state)?;
+                    mgr.on_restart(&mut state)?;
                     std::process::exit(0);
                 } else {
                     fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;

libafl_sugar/src/inmemory.rs

@@ -281,7 +281,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -296,7 +296,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -318,7 +318,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -333,7 +333,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;

libafl_sugar/src/qemu.rs

@@ -319,7 +319,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -340,7 +340,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -431,7 +431,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;
@@ -452,7 +452,7 @@ where
                             &mut mgr,
                             iters,
                         )?;
-                        mgr.mgr_on_restart(&mut state)?;
+                        mgr.on_restart(&mut state)?;
                         std::process::exit(0);
                     } else {
                         fuzzer.fuzz_loop(&mut stages, &mut executor, &mut state, &mut mgr)?;