From 3a7c0da037a145ebd95063be8e36467fe1cf864c Mon Sep 17 00:00:00 2001
From: Alwin Berger <alwin.berger@tu-dortmund.de>
Date: Wed, 19 Jun 2024 14:22:04 +0200
Subject: [PATCH] fast snapshots by default

---
 fuzzers/FRET/Cargo.toml       |  2 +-
 fuzzers/FRET/src/fuzzer.rs    | 29 ++++++++++++++++++-----------
 fuzzers/FRET/src/qemustate.rs |  7 +++++--
 3 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/fuzzers/FRET/Cargo.toml b/fuzzers/FRET/Cargo.toml
index 2770b993df..8767916a1f 100644
--- a/fuzzers/FRET/Cargo.toml
+++ b/fuzzers/FRET/Cargo.toml
@@ -5,7 +5,7 @@ authors = ["Alwin Berger <alwin.berger@tu-dortmund.de>"]
 edition = "2021"
 
 [features]
-default = ["std", "snapshot_restore", "singlecore", "restarting", "do_hash_notify_state", "config_stg", "fuzz_int" ]
+default = ["std", "snapshot_restore", "snapshot_fast", "singlecore", "restarting", "do_hash_notify_state", "config_stg", "fuzz_int" ]
 std = []
 # Exec environemnt basics
 snapshot_restore = []
diff --git a/fuzzers/FRET/src/fuzzer.rs b/fuzzers/FRET/src/fuzzer.rs
index 3c9f25be49..24239a77c9 100644
--- a/fuzzers/FRET/src/fuzzer.rs
+++ b/fuzzers/FRET/src/fuzzer.rs
@@ -423,20 +423,27 @@ let mut run_client = |state: Option<_>, mut mgr, _core_id| {
         // "-semihosting",
         // "--semihosting-config",
         // "enable=on,target=native",
-        // "-snapshot",
-        // "-drive",
-        // "if=none,format=qcow2,file=dummy.qcow2",
+        #[cfg(not(feature = "snapshot_fast"))]
+        "-snapshot",
+        #[cfg(not(feature = "snapshot_fast"))]
+        "-drive",
+        #[cfg(not(feature = "snapshot_fast"))]
+        "if=none,format=qcow2,file=dummy.qcow2",
     ].into_iter().map(String::from).collect();
     let env: Vec<(String, String)> = env::vars().collect();
     let emu = Qemu::init(&args, &env).expect("Emulator creation failed");
 
-    // if let Some(main_addr) = main_addr {
-    //     unsafe {
-    //         emu.set_breakpoint(main_addr);
-    //         emu.run();
-    //         emu.remove_breakpoint(main_addr);
-    //     }
-    // }
+    if let Some(main_addr) = main_addr {
+        unsafe {
+            libafl_qemu::sys::libafl_qemu_set_native_breakpoint(main_addr as u64);
+            emu.run();
+            libafl_qemu::sys::libafl_qemu_remove_native_breakpoint(main_addr as u64);
+        }
+    }
+    #[cfg(feature = "snapshot_fast")]
+    let initial_snap = Some(emu.create_fast_snapshot(true));
+    #[cfg(not(feature = "snapshot_fast"))]
+    let initial_snap = None;
 
     unsafe { emu.set_breakpoint(breakpoint); }// BREAKPOINT 
 
@@ -590,7 +597,7 @@ let mut run_client = |state: Option<_>, mut mgr, _core_id| {
         let qhelpers = (QemuSystemStateHelper::new(api_addreses,api_ranges,isr_addreses,isr_ranges,curr_tcb_pointer,task_queue_addr,task_delay_addr,task_delay_overflow_addr,scheduler_lock,scheduler_running, critical_section,input_counter_ptr,app_range.clone()), qhelpers);
         #[cfg(feature = "observe_edges")]
         let qhelpers = (QemuEdgeCoverageHelper::new(denylist, QemuFilterList::None), qhelpers);
-        let qhelpers = (QemuStateRestoreHelper::new(), qhelpers);
+        let qhelpers = (QemuStateRestoreHelper::with_fast(initial_snap), qhelpers);
 
         let mut hooks = QemuHooks::new(emu.clone(),qhelpers);
 
diff --git a/fuzzers/FRET/src/qemustate.rs b/fuzzers/FRET/src/qemustate.rs
index 390a16992c..97a97ffa13 100644
--- a/fuzzers/FRET/src/qemustate.rs
+++ b/fuzzers/FRET/src/qemustate.rs
@@ -16,7 +16,6 @@ use libafl_qemu::{
 #[derive(Debug)]
 pub struct QemuStateRestoreHelper {
     has_snapshot: bool,
-    use_snapshot: bool,
     saved_cpu_states: Vec<CPUArchState>,
     fastsnap: Option<FastSnapshotPtr>
 }
@@ -26,11 +25,15 @@ impl QemuStateRestoreHelper {
     pub fn new() -> Self {
         Self {
             has_snapshot: false,
-            use_snapshot: true,
             saved_cpu_states: vec![],
             fastsnap: None
         }
     }
+    pub fn with_fast(fastsnap: Option<FastSnapshotPtr>) -> Self {
+        let mut r = Self::new();
+        r.fastsnap = fastsnap;
+        r
+    }
 }
 
 impl Default for QemuStateRestoreHelper {