From 0f42efa12bbcda9b3eea88303580b8e4d0ea8c71 Mon Sep 17 00:00:00 2001
From: Stefan Zabka <git@zabka.it>
Date: Wed, 24 Apr 2024 11:58:46 +0200
Subject: [PATCH] refactor(snapshot.rs): use const generics (#2069)

* refactor(snapshot.rs): use const generics

* Add docs to read hooks

---------

Co-authored-by: Romain Malmain <romain.malmain@pm.me>
---
 libafl_qemu/src/emu.rs      | 12 +++++++++
 libafl_qemu/src/snapshot.rs | 51 +++++--------------------------------
 2 files changed, 18 insertions(+), 45 deletions(-)

diff --git a/libafl_qemu/src/emu.rs b/libafl_qemu/src/emu.rs
index 7a0a5bacc8..73623e5489 100644
--- a/libafl_qemu/src/emu.rs
+++ b/libafl_qemu/src/emu.rs
@@ -1161,6 +1161,18 @@ impl Qemu {
         }
     }
 
+    /// `data` can be used to pass data that can be accessed as the first argument in the `gen` and the `exec` functions
+    ///
+    /// `gen` gets passed the current programm counter, mutable access to a `TCGTemp` and information about the memory
+    /// access being performed.
+    ///  The `u64` return value is an id that gets passed to the `exec` functions as their second argument.
+    ///
+    /// `exec` hooks get invoked on every read performed by the guest
+    ///
+    /// `exec1`-`exec8` special case accesses of width 1-8
+    ///
+    /// If there is no specialized hook for a given read width, the `exec_n` will be
+    /// called and its last argument will specify the access width
     #[allow(clippy::missing_transmute_annotations)]
     pub fn add_read_hooks<T: Into<HookData>>(
         &self,
diff --git a/libafl_qemu/src/snapshot.rs b/libafl_qemu/src/snapshot.rs
index 994ee1831b..5259181ea3 100644
--- a/libafl_qemu/src/snapshot.rs
+++ b/libafl_qemu/src/snapshot.rs
@@ -498,10 +498,10 @@ where
             // The ASan helper, if present, will call the tracer hook for the snapshot helper as opt
             hooks.writes(
                 Hook::Empty,
-                Hook::Function(trace_write1_snapshot::<QT, S>),
-                Hook::Function(trace_write2_snapshot::<QT, S>),
-                Hook::Function(trace_write4_snapshot::<QT, S>),
-                Hook::Function(trace_write8_snapshot::<QT, S>),
+                Hook::Function(trace_write_snapshot::<QT, S, 1>),
+                Hook::Function(trace_write_snapshot::<QT, S, 2>),
+                Hook::Function(trace_write_snapshot::<QT, S, 4>),
+                Hook::Function(trace_write_snapshot::<QT, S, 8>),
                 Hook::Function(trace_write_n_snapshot::<QT, S>),
             );
         }
@@ -521,7 +521,7 @@ where
     }
 }
 
-pub fn trace_write1_snapshot<QT, S>(
+pub fn trace_write_snapshot<QT, S, const SIZE: usize>(
     hooks: &mut QemuHooks<QT, S>,
     _state: Option<&mut S>,
     _id: u64,
@@ -531,46 +531,7 @@ pub fn trace_write1_snapshot<QT, S>(
     QT: QemuHelperTuple<S>,
 {
     let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
-    h.access(addr, 1);
-}
-
-pub fn trace_write2_snapshot<QT, S>(
-    hooks: &mut QemuHooks<QT, S>,
-    _state: Option<&mut S>,
-    _id: u64,
-    addr: GuestAddr,
-) where
-    S: UsesInput,
-    QT: QemuHelperTuple<S>,
-{
-    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
-    h.access(addr, 2);
-}
-
-pub fn trace_write4_snapshot<QT, S>(
-    hooks: &mut QemuHooks<QT, S>,
-    _state: Option<&mut S>,
-    _id: u64,
-    addr: GuestAddr,
-) where
-    S: UsesInput,
-    QT: QemuHelperTuple<S>,
-{
-    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
-    h.access(addr, 4);
-}
-
-pub fn trace_write8_snapshot<QT, S>(
-    hooks: &mut QemuHooks<QT, S>,
-    _state: Option<&mut S>,
-    _id: u64,
-    addr: GuestAddr,
-) where
-    S: UsesInput,
-    QT: QemuHelperTuple<S>,
-{
-    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
-    h.access(addr, 8);
+    h.access(addr, SIZE);
 }
 
 pub fn trace_write_n_snapshot<QT, S>(