Add SnapshotModule to qemu_launcher (#2887)

* Add SnapshotModule to qemu_launcher

---------

Co-authored-by: Your Name <you@example.com>

fuzzers/binary_only/qemu_launcher/src/fuzzer.rs

@@ -6,20 +6,18 @@ use std::{
 
 use clap::Parser;
 #[cfg(feature = "simplemgr")]
-use libafl::events::SimpleEventManager;
+use libafl::events::{ClientDescription, SimpleEventManager};
 #[cfg(not(feature = "simplemgr"))]
 use libafl::events::{EventConfig, Launcher, MonitorTypedEventManager};
 use libafl::{
-    events::{ClientDescription, LlmpEventManagerBuilder},
     monitors::{tui::TuiMonitor, Monitor, MultiMonitor},
     Error,
 };
-use libafl_bolts::{core_affinity::CoreId, current_time, llmp::LlmpBroker, tuples::tuple_list};
+#[cfg(feature = "simplemgr")]
+use libafl_bolts::core_affinity::CoreId;
+use libafl_bolts::current_time;
 #[cfg(not(feature = "simplemgr"))]
-use libafl_bolts::{
-    shmem::{ShMemProvider, StdShMemProvider},
-    staterestore::StateRestorer,
-};
+use libafl_bolts::shmem::{ShMemProvider, StdShMemProvider};
 #[cfg(unix)]
 use {
     nix::unistd::dup,
@@ -85,7 +83,7 @@ impl Fuzzer {
     {
         // The shared memory allocator
         #[cfg(not(feature = "simplemgr"))]
-        let mut shmem_provider = StdShMemProvider::new()?;
+        let shmem_provider = StdShMemProvider::new()?;
         /* If we are running in verbose, don't provide a replacement stdout, otherwise, use /dev/null */
         #[cfg(not(feature = "simplemgr"))]
         let stdout = if self.options.verbose {

fuzzers/binary_only/qemu_launcher/src/instance.rs

@@ -6,19 +6,21 @@ use libafl::events::SimpleEventManager;
 #[cfg(not(feature = "simplemgr"))]
 use libafl::events::{LlmpRestartingEventManager, MonitorTypedEventManager};
 use libafl::{
-    corpus::{Corpus, InMemoryOnDiskCorpus, OnDiskCorpus},
+    corpus::{Corpus, HasCurrentCorpusId, InMemoryOnDiskCorpus, OnDiskCorpus},
     events::{ClientDescription, EventRestarter},
-    executors::{Executor, ShadowExecutor},
+    executors::{Executor, ExitKind, ShadowExecutor},
     feedback_or, feedback_or_fast,
     feedbacks::{CrashFeedback, MaxMapFeedback, TimeFeedback, TimeoutFeedback},
     fuzzer::{Evaluator, Fuzzer, StdFuzzer},
-    inputs::BytesInput,
+    inputs::{BytesInput, Input},
     monitors::Monitor,
     mutators::{
         havoc_mutations, token_mutations::I2SRandReplace, tokens_mutations, StdMOptMutator,
         StdScheduledMutator, Tokens,
     },
-    observers::{CanTrack, HitcountsMapObserver, TimeObserver, VariableMapObserver},
+    observers::{
+        CanTrack, HitcountsMapObserver, ObserversTuple, TimeObserver, VariableMapObserver,
+    },
     schedulers::{
         powersched::PowerSchedule, IndexesLenTimeMinimizerScheduler, PowerQueueScheduler,
     },
@@ -26,7 +28,7 @@ use libafl::{
         calibrate::CalibrationStage, power::StdPowerMutationalStage, AflStatsStage, IfStage,
         ShadowTracingStage, StagesTuple, StdMutationalStage,
     },
-    state::{HasCorpus, StdState},
+    state::{HasCorpus, HasExecutions, HasSolutions, StdState},
     Error, HasMetadata,
 };
 #[cfg(not(feature = "simplemgr"))]
@@ -42,7 +44,7 @@ use libafl_qemu::{
         cmplog::CmpLogObserver,
         edges::EdgeCoverageFullVariant,
         utils::filters::{HasAddressFilter, NopPageFilter, StdAddressFilter},
-        EdgeCoverageModule, EmulatorModuleTuple, StdEdgeCoverageModule,
+        EdgeCoverageModule, EmulatorModuleTuple, SnapshotModule, StdEdgeCoverageModule,
     },
     Emulator, GuestAddr, Qemu, QemuExecutor,
 };
@@ -62,6 +64,27 @@ pub type ClientMgr<M> = MonitorTypedEventManager<
     M,
 >;
 
+/*
+ * The snapshot and iterations options interact as follows:
+ *
+ * +----------+------------+-------------------------------------------+
+ * | snapshot | iterations | Functionality                             |
+ * +----------+------------+-------------------------------------------+
+ * |    N     |     N      | We set the snapshot module into manual    |
+ * |          |            | mode and never reset it.                  |
+ * +----------+------------+-------------------------------------------+
+ * |    N     |     Y      | We set the snapshot module into manual    |
+ * |          |            | mode and never reset it.                  |
+ * +----------+------------+-------------------------------------------+
+ * |    Y     |     N      | We set the snapshot module into automatic |
+ * |          |            | mode so it resets after every iteration.  |
+ * +----------+------------+-------------------------------------------+
+ * |    Y     |     Y      | We set the snapshot module into manual    |
+ * |          |            | mode and manually reset it after the      |
+ * |          |            | required number of iterations are done.   |
+ * +----------+------------+-------------------------------------------+
+ */
+
 #[derive(TypedBuilder)]
 pub struct Instance<'a, M: Monitor> {
     options: &'a FuzzerOptions,
@@ -131,7 +154,22 @@ impl<M: Monitor> Instance<'_, M> {
             .map_observer(edges_observer.as_mut())
             .build()?;
 
-        let modules = modules.prepend(edge_coverage_module);
+        let mut snapshot_module = SnapshotModule::new();
+
+        /*
+         * Since the generics for the modules are already excessive when taking
+         * into accout asan, asan guest mode, cmplog, and injection, we will
+         * always include the SnapshotModule in all configurations, but simply
+         * not use it when it is not required. See the table at the top of this
+         * file for details.
+         */
+        if !self.options.snapshots || self.options.iterations.is_some() {
+            snapshot_module.use_manual_reset();
+        }
+
+        let modules = modules
+            .prepend(edge_coverage_module)
+            .prepend(snapshot_module);
         let mut emulator = Emulator::empty()
             .qemu_parameters(args)
             .modules(modules)
@@ -266,9 +304,9 @@ impl<M: Monitor> Instance<'_, M> {
             // Create an observation channel using cmplog map
             let cmplog_observer = CmpLogObserver::new("cmplog", true);
 
-            let mut executor = ShadowExecutor::new(executor, tuple_list!(cmplog_observer));
+            let mut shadow_executor = ShadowExecutor::new(executor, tuple_list!(cmplog_observer));
 
-            let tracing = ShadowTracingStage::new(&mut executor);
+            let tracing = ShadowTracingStage::new(&mut shadow_executor);
 
             // Setup a randomic Input2State stage
             let i2s = StdMutationalStage::new(StdScheduledMutator::new(tuple_list!(
@@ -289,7 +327,14 @@ impl<M: Monitor> Instance<'_, M> {
             // The order of the stages matter!
             let mut stages = tuple_list!(calibration, tracing, i2s, power, stats_stage);
 
-            self.fuzz(&mut state, &mut fuzzer, &mut executor, &mut stages)
+            self.fuzz(
+                &mut state,
+                &mut fuzzer,
+                &mut shadow_executor,
+                Self::reset_shadow_executor_snapshot_module,
+                qemu,
+                &mut stages,
+            )
         } else {
             // Create a QEMU in-process executor
             let mut executor = QemuExecutor::new(
@@ -306,25 +351,88 @@ impl<M: Monitor> Instance<'_, M> {
             let mutator = StdScheduledMutator::new(havoc_mutations().merge(tokens_mutations()));
             let mut stages = tuple_list!(StdMutationalStage::new(mutator));
 
-            self.fuzz(&mut state, &mut fuzzer, &mut executor, &mut stages)
+            self.fuzz(
+                &mut state,
+                &mut fuzzer,
+                &mut executor,
+                Self::reset_executor_snapshot_module,
+                qemu,
+                &mut stages,
+            )
         }
     }
 
-    fn fuzz<Z, E, ST>(
+    fn reset_executor_snapshot_module<'a, C, CM, ED, EM, ET, H, I, OT, S, SM, Z>(
+        executor: &mut QemuExecutor<'a, C, CM, ED, EM, (SnapshotModule, ET), H, I, OT, S, SM, Z>,
+        qemu: Qemu,
+    ) where
+        I: Input + Unpin,
+        ET: EmulatorModuleTuple<I, S>,
+        S: HasCorpus<I> + HasCurrentCorpusId + HasSolutions<I> + HasExecutions + Unpin,
+        H: for<'e, 's, 'i> FnMut(
+            &'e mut Emulator<C, CM, ED, (SnapshotModule, ET), I, S, SM>,
+            &'s mut S,
+            &'i I,
+        ) -> ExitKind,
+        OT: ObserversTuple<I, S>,
+    {
+        executor
+            .inner_mut()
+            .exposed_executor_state_mut()
+            .modules_mut()
+            .modules_mut()
+            .0
+            .reset(qemu);
+    }
+
+    fn reset_shadow_executor_snapshot_module<'a, C, CM, ED, EM, ET, H, I, OT, S, SM, SOT, Z>(
+        executor: &mut ShadowExecutor<
+            QemuExecutor<'a, C, CM, ED, EM, (SnapshotModule, ET), H, I, OT, S, SM, Z>,
+            I,
+            S,
+            SOT,
+        >,
+        qemu: Qemu,
+    ) where
+        I: Input + Unpin,
+        ET: EmulatorModuleTuple<I, S>,
+        S: HasCorpus<I> + HasCurrentCorpusId + HasSolutions<I> + HasExecutions + Unpin,
+        H: for<'e, 's, 'i> FnMut(
+            &'e mut Emulator<C, CM, ED, (SnapshotModule, ET), I, S, SM>,
+            &'s mut S,
+            &'i I,
+        ) -> ExitKind,
+        OT: ObserversTuple<I, S>,
+        SOT: ObserversTuple<I, S>,
+    {
+        executor
+            .executor_mut()
+            .inner_mut()
+            .exposed_executor_state_mut()
+            .modules_mut()
+            .modules_mut()
+            .0
+            .reset(qemu);
+    }
+
+    fn fuzz<Z, E, ST, RSM>(
         &mut self,
         state: &mut ClientState,
         fuzzer: &mut Z,
         executor: &mut E,
+        reset_snapshot_module: RSM,
+        qemu: Qemu,
         stages: &mut ST,
     ) -> Result<(), Error>
     where
         Z: Fuzzer<E, ClientMgr<M>, BytesInput, ClientState, ST>
             + Evaluator<E, ClientMgr<M>, BytesInput, ClientState>,
         ST: StagesTuple<E, ClientMgr<M>, ClientState, Z>,
+        RSM: Fn(&mut E, Qemu),
     {
+        if state.must_load_initial_inputs() {
             let corpus_dirs = [self.options.input_dir()];
 
-        if state.must_load_initial_inputs() {
             state
                 .load_initial_inputs(fuzzer, executor, &mut self.mgr, &corpus_dirs)
                 .unwrap_or_else(|_| {
@@ -334,12 +442,23 @@ impl<M: Monitor> Instance<'_, M> {
             println!("We imported {} inputs from disk.", state.corpus().count());
         }
 
+        /*
+         * See the table a the top of this file for details on how the snapshot
+         * and iterations options interact.
+         */
         if let Some(iters) = self.options.iterations {
+            if self.options.snapshots {
+                loop {
+                    reset_snapshot_module(executor, qemu);
+                    fuzzer.fuzz_loop_for(stages, executor, state, &mut self.mgr, iters)?;
+                }
+            } else {
                 fuzzer.fuzz_loop_for(stages, executor, state, &mut self.mgr, iters)?;
 
                 // It's important, that we store the state before restarting!
                 // Else, the parent will not respawn a new child and quit.
                 self.mgr.on_restart(state)?;
+            }
         } else {
             fuzzer.fuzz_loop(stages, executor, state, &mut self.mgr)?;
         }

fuzzers/binary_only/qemu_launcher/src/options.rs

@@ -62,6 +62,9 @@ pub struct FuzzerOptions {
     #[clap(long, help = "Enable AFL++ style output", conflicts_with = "verbose")]
     pub tui: bool,
 
+    #[clap(long, help = "Enable use of snapshots to restore state")]
+    pub snapshots: bool,
+
     #[arg(long = "iterations", help = "Maximum number of iterations")]
     pub iterations: Option<u64>,
 

libafl/src/executors/shadow.rs

@@ -62,6 +62,18 @@ where
     pub fn shadow_observers_mut(&mut self) -> RefIndexable<&mut SOT, SOT> {
         RefIndexable::from(&mut self.shadow_observers)
     }
+
+    /// Inner executor
+    #[inline]
+    pub fn executor(&self) -> &E {
+        &self.executor
+    }
+
+    /// Inner executor
+    #[inline]
+    pub fn executor_mut(&mut self) -> &mut E {
+        &mut self.executor
+    }
 }
 
 impl<E, EM, I, S, SOT, Z> Executor<EM, I, S, Z> for ShadowExecutor<E, I, S, SOT>

libafl_qemu/src/modules/usermode/snapshot.rs

@@ -98,6 +98,7 @@ pub struct SnapshotModule {
     pub empty: bool,
     pub accurate_unmap: bool,
     pub interval_filter: Vec<IntervalSnapshotFilter>,
+    auto_reset: bool,
 }
 
 impl core::fmt::Debug for SnapshotModule {
@@ -130,6 +131,7 @@ impl SnapshotModule {
             empty: true,
             accurate_unmap: false,
             interval_filter: Vec::<IntervalSnapshotFilter>::new(),
+            auto_reset: true,
         }
     }
 
@@ -148,6 +150,7 @@ impl SnapshotModule {
             empty: true,
             accurate_unmap: false,
             interval_filter,
+            auto_reset: true,
         }
     }
 
@@ -166,6 +169,7 @@ impl SnapshotModule {
             empty: true,
             accurate_unmap: false,
             interval_filter: Vec::<IntervalSnapshotFilter>::new(),
+            auto_reset: true,
         }
     }
 
@@ -173,6 +177,10 @@ impl SnapshotModule {
         self.accurate_unmap = true;
     }
 
+    pub fn use_manual_reset(&mut self) {
+        self.auto_reset = false;
+    }
+
     pub fn to_skip(&self, addr: GuestAddr) -> bool {
         for filter in &self.interval_filter {
             match filter {
@@ -731,7 +739,7 @@ where
     {
         if self.empty {
             self.snapshot(qemu);
-        } else {
+        } else if self.auto_reset {
             self.reset(qemu);
         }
     }