From 0e4c6722f0789c024027e05a6e6b68337d9e40bc Mon Sep 17 00:00:00 2001 From: "Dongjia \"toka\" Zhang" Date: Tue, 18 Feb 2025 19:55:02 +0100 Subject: [PATCH] Replay stage (#3003) * replay stage * seems to be working * a * rever * rev * rdy for merge * fmt * lol --- fuzzers/baby/tutorial/src/lib.rs | 2 +- fuzzers/binary_only/qemu_launcher/Justfile | 2 +- .../binary_only/qemu_launcher/src/instance.rs | 2 +- .../full_system/nyx_launcher/src/instance.rs | 2 +- fuzzers/inprocess/fuzzbench/Justfile | 3 +- fuzzers/inprocess/fuzzbench_ctx/Justfile | 2 +- fuzzers/inprocess/libfuzzer_libpng/src/lib.rs | 2 +- .../libfuzzer_libpng_cmin/src/lib.rs | 2 +- .../libfuzzer_libpng_norestart/src/lib.rs | 2 +- .../libfuzzer_libpng_tcp_manager/src/lib.rs | 2 +- .../libfuzzer_windows_asan/src/lib.rs | 2 +- libafl/src/events/centralized.rs | 4 +- libafl/src/events/llmp/restarting.rs | 2 +- libafl/src/events/mod.rs | 10 +- libafl/src/events/simple.rs | 4 +- libafl/src/events/tcp.rs | 4 +- libafl/src/executors/inprocess/mod.rs | 2 +- libafl/src/fuzzer/mod.rs | 4 +- libafl/src/stages/mod.rs | 3 + libafl/src/stages/replay.rs | 171 ++++++++++++++++++ libafl_sugar/src/forkserver.rs | 4 +- libafl_sugar/src/inmemory.rs | 8 +- libafl_sugar/src/qemu.rs | 8 +- 23 files changed, 211 insertions(+), 36 deletions(-) create mode 100644 libafl/src/stages/replay.rs diff --git a/fuzzers/baby/tutorial/src/lib.rs b/fuzzers/baby/tutorial/src/lib.rs index b86a5677f5..3e44abff34 100644 --- a/fuzzers/baby/tutorial/src/lib.rs +++ b/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(()) } diff --git a/fuzzers/binary_only/qemu_launcher/Justfile b/fuzzers/binary_only/qemu_launcher/Justfile index 36288989d0..149d4690a2 100644 --- a/fuzzers/binary_only/qemu_launcher/Justfile +++ b/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] diff --git a/fuzzers/binary_only/qemu_launcher/src/instance.rs b/fuzzers/binary_only/qemu_launcher/src/instance.rs index 966d04d6ff..245afc3b91 100644 --- a/fuzzers/binary_only/qemu_launcher/src/instance.rs +++ b/fuzzers/binary_only/qemu_launcher/src/instance.rs @@ -469,7 +469,7 @@ impl 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)?; diff --git a/fuzzers/full_system/nyx_launcher/src/instance.rs b/fuzzers/full_system/nyx_launcher/src/instance.rs index b705f1790b..568297844b 100644 --- a/fuzzers/full_system/nyx_launcher/src/instance.rs +++ b/fuzzers/full_system/nyx_launcher/src/instance.rs @@ -250,7 +250,7 @@ impl 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)?; } diff --git a/fuzzers/inprocess/fuzzbench/Justfile b/fuzzers/inprocess/fuzzbench/Justfile index 4892a9d85f..6316f40605 100644 --- a/fuzzers/inprocess/fuzzbench/Justfile +++ b/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: diff --git a/fuzzers/inprocess/fuzzbench_ctx/Justfile b/fuzzers/inprocess/fuzzbench_ctx/Justfile index aad1ca4816..a0bf7510da 100644 --- a/fuzzers/inprocess/fuzzbench_ctx/Justfile +++ b/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 diff --git a/fuzzers/inprocess/libfuzzer_libpng/src/lib.rs b/fuzzers/inprocess/libfuzzer_libpng/src/lib.rs index fb7abc8774..09d5357bfd 100644 --- a/fuzzers/inprocess/libfuzzer_libpng/src/lib.rs +++ b/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(()) } diff --git a/fuzzers/inprocess/libfuzzer_libpng_cmin/src/lib.rs b/fuzzers/inprocess/libfuzzer_libpng_cmin/src/lib.rs index 5f144c37c8..850d15ed27 100644 --- a/fuzzers/inprocess/libfuzzer_libpng_cmin/src/lib.rs +++ b/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(()) } diff --git a/fuzzers/inprocess/libfuzzer_libpng_norestart/src/lib.rs b/fuzzers/inprocess/libfuzzer_libpng_norestart/src/lib.rs index 479faa3b27..684d616f76 100644 --- a/fuzzers/inprocess/libfuzzer_libpng_norestart/src/lib.rs +++ b/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(()) }; diff --git a/fuzzers/inprocess/libfuzzer_libpng_tcp_manager/src/lib.rs b/fuzzers/inprocess/libfuzzer_libpng_tcp_manager/src/lib.rs index ebd7bea130..65a2044d19 100644 --- a/fuzzers/inprocess/libfuzzer_libpng_tcp_manager/src/lib.rs +++ b/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(()) } diff --git a/fuzzers/inprocess/libfuzzer_windows_asan/src/lib.rs b/fuzzers/inprocess/libfuzzer_windows_asan/src/lib.rs index c02ffb863e..89b50d70ab 100644 --- a/fuzzers/inprocess/libfuzzer_windows_asan/src/lib.rs +++ b/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(()) } diff --git a/libafl/src/events/centralized.rs b/libafl/src/events/centralized.rs index 64283e0070..781a17d3ae 100644 --- a/libafl/src/events/centralized.rs +++ b/libafl/src/events/centralized.rs @@ -272,9 +272,9 @@ where SP: ShMemProvider, { #[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(()) } } diff --git a/libafl/src/events/llmp/restarting.rs b/libafl/src/events/llmp/restarting.rs index 00c84c6cc5..31722536cd 100644 --- a/libafl/src/events/llmp/restarting.rs +++ b/libafl/src/events/llmp/restarting.rs @@ -256,7 +256,7 @@ where SP: ShMemProvider, { /// 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 { diff --git a/libafl/src/events/mod.rs b/libafl/src/events/mod.rs index c8ebcafedc..154889e2ec 100644 --- a/libafl/src/events/mod.rs +++ b/libafl/src/events/mod.rs @@ -559,10 +559,10 @@ pub trait EventRestarter { /// 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(restarter: &mut EM, state: &mut S) -> Result<(), Error> where @@ -644,7 +644,7 @@ impl EventRestarter 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, { #[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) } } diff --git a/libafl/src/events/simple.rs b/libafl/src/events/simple.rs index bbaab9dc35..a31499b7f7 100644 --- a/libafl/src/events/simple.rs +++ b/libafl/src/events/simple.rs @@ -112,7 +112,7 @@ impl EventRestarter for SimpleEventManager 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 diff --git a/libafl/src/events/tcp.rs b/libafl/src/events/tcp.rs index b0d2885f95..48c8fb5811 100644 --- a/libafl/src/events/tcp.rs +++ b/libafl/src/events/tcp.rs @@ -626,7 +626,7 @@ impl EventRestarter for TcpEventManager 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, { /// 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 diff --git a/libafl/src/executors/inprocess/mod.rs b/libafl/src/executors/inprocess/mod.rs index 5dd5b640d1..d8450dd27e 100644 --- a/libafl/src/executors/inprocess/mod.rs +++ b/libafl/src/executors/inprocess/mod.rs @@ -430,7 +430,7 @@ pub fn run_observers_and_save_state( } // 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!"); } diff --git a/libafl/src/fuzzer/mod.rs b/libafl/src/fuzzer/mod.rs index c161cc152e..a3d4790d23 100644 --- a/libafl/src/fuzzer/mod.rs +++ b/libafl/src/fuzzer/mod.rs @@ -207,7 +207,7 @@ pub trait Fuzzer { /// 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 { /// 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, diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs index 135465c7cc..6c77aa518e 100644 --- a/libafl/src/stages/mod.rs +++ b/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; diff --git a/libafl/src/stages/replay.rs b/libafl/src/stages/replay.rs new file mode 100644 index 0000000000..dcef475387 --- /dev/null +++ b/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 { + name: Cow<'static, str>, + restart_helper: ReplayRestartingHelper, + phantom: PhantomData, +} + +impl Default for ReplayStage { + fn default() -> Self { + Self::new() + } +} + +impl Named for ReplayStage { + 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, + done_solution: HashSet, +} + +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 ReplayStage { + #[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 Stage for ReplayStage +where + S: HasCorpus + HasSolutions, + Z: Evaluator, + I: Clone, +{ + fn should_restart(&mut self, _state: &mut S) -> Result { + 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 = 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 = 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(()) + } +} diff --git a/libafl_sugar/src/forkserver.rs b/libafl_sugar/src/forkserver.rs index 58be325d73..2e285897ee 100644 --- a/libafl_sugar/src/forkserver.rs +++ b/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)?; diff --git a/libafl_sugar/src/inmemory.rs b/libafl_sugar/src/inmemory.rs index 1a635ba10b..e26765a26d 100644 --- a/libafl_sugar/src/inmemory.rs +++ b/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)?; diff --git a/libafl_sugar/src/qemu.rs b/libafl_sugar/src/qemu.rs index a2a6d849fc..50c8d3a284 100644 --- a/libafl_sugar/src/qemu.rs +++ b/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)?;