Light refactor of EmulatorModules

* qemu is now a parameter to EmulatorModule callbacks and most function hooks.
* EmulatorModules is initialized before QEMU is initialized.

libafl_qemu/src/command/mod.rs

@@ -452,8 +452,8 @@ where
         // Unleash hooks if locked
         if emu.driver_mut().unlock_hooks() {
             // Prepare hooks
-            emu.modules_mut().first_exec_all(state);
-            emu.modules_mut().pre_exec_all(state, input);
+            emu.modules_mut().first_exec_all(qemu, state);
+            emu.modules_mut().pre_exec_all(qemu, state, input);
         }
 
         // Auto page filtering if option is enabled

libafl_qemu/src/emu/builder.rs

@@ -12,8 +12,8 @@ use crate::{
     command::{CommandManager, NopCommandManager, StdCommandManager},
     config::QemuConfig,
     modules::{EmulatorModule, EmulatorModuleTuple},
-    Emulator, EmulatorHooks, NopEmulatorDriver, NopSnapshotManager, Qemu, QemuHooks, QemuInitError,
-    QemuParams, StdEmulatorDriver, StdSnapshotManager,
+    Emulator, EmulatorHooks, EmulatorModules, NopEmulatorDriver, NopSnapshotManager, Qemu,
+    QemuHooks, QemuInitError, QemuParams, StdEmulatorDriver, StdSnapshotManager,
 };
 
 #[derive(Clone, Debug)]
@@ -111,17 +111,23 @@ where
     {
         let qemu_parameters = self.qemu_parameters.ok_or(QemuInitError::EmptyArgs)?;
 
-        let mut emulator_hooks = unsafe { EmulatorHooks::new(QemuHooks::get_unchecked()) };
+        let emulator_hooks = unsafe { EmulatorHooks::new(QemuHooks::get_unchecked()) };
 
-        self.modules.pre_qemu_init_all(&mut emulator_hooks);
+        let mut emulator_modules = EmulatorModules::new(emulator_hooks, self.modules);
+
+        // TODO: fix things there properly. The biggest issue being that it creates 2 mut ref to the module with the callback being called
+        unsafe {
+            emulator_modules
+                .modules_mut()
+                .pre_qemu_init_all(EmulatorModules::<ET, S>::emulator_modules_mut_unchecked());
+        }
 
         let qemu = Qemu::init(qemu_parameters)?;
 
         unsafe {
             Ok(Emulator::new_with_qemu(
                 qemu,
-                emulator_hooks,
-                self.modules,
+                emulator_modules,
                 self.driver,
                 self.snapshot_manager,
                 self.command_manager,

libafl_qemu/src/emu/drivers.rs

@@ -54,7 +54,7 @@ where
     /// Just before calling user's harness for the first time.
     /// Called only once
     fn first_harness_exec(emulator: &mut Emulator<CM, Self, ET, S, SM>, state: &mut S) {
-        emulator.modules.first_exec_all(state);
+        emulator.modules.first_exec_all(emulator.qemu, state);
     }
 
     /// Just before calling user's harness
@@ -63,7 +63,7 @@ where
         state: &mut S,
         input: &S::Input,
     ) {
-        emulator.modules.pre_exec_all(state, input);
+        emulator.modules.pre_exec_all(emulator.qemu, state, input);
     }
 
     /// Just after returning from user's harness
@@ -78,7 +78,7 @@ where
     {
         emulator
             .modules
-            .post_exec_all(state, input, observers, exit_kind);
+            .post_exec_all(emulator.qemu, state, input, observers, exit_kind);
     }
 
     /// Just before entering QEMU
@@ -169,7 +169,7 @@ where
 {
     fn first_harness_exec(emulator: &mut Emulator<CM, Self, ET, S, SM>, state: &mut S) {
         if !emulator.driver.hooks_locked {
-            emulator.modules.first_exec_all(state);
+            emulator.modules.first_exec_all(emulator.qemu, state);
         }
     }
 
@@ -179,7 +179,7 @@ where
         input: &S::Input,
     ) {
         if !emulator.driver.hooks_locked {
-            emulator.modules.pre_exec_all(state, input);
+            emulator.modules.pre_exec_all(emulator.qemu, state, input);
         }
 
         let input_location = { emulator.driver.input_location.get().cloned() };
@@ -206,7 +206,7 @@ where
         if !emulator.driver.hooks_locked {
             emulator
                 .modules
-                .post_exec_all(state, input, observers, exit_kind);
+                .post_exec_all(emulator.qemu, state, input, observers, exit_kind);
         }
     }
 

libafl_qemu/src/emu/hooks.rs

@@ -3,7 +3,7 @@
 use std::{fmt::Debug, marker::PhantomData, mem::transmute, pin::Pin, ptr};
 
 use libafl::{executors::ExitKind, inputs::UsesInput, observers::ObserversTuple};
-use libafl_qemu_sys::{CPUArchStatePtr, CPUStatePtr, FatPtr, GuestAddr, GuestUsize, TCGTemp};
+use libafl_qemu_sys::{CPUStatePtr, FatPtr, GuestAddr, GuestUsize, TCGTemp};
 
 #[cfg(feature = "usermode")]
 use crate::qemu::{
@@ -37,9 +37,15 @@ use crate::{
         ReadExecNHook, ReadGenHook, ReadHookId, TcgHookState, WriteExecHook, WriteExecNHook,
         WriteGenHook, WriteHookId,
     },
-    CpuPostRunHook, CpuPreRunHook, CpuRunHookId, HookState, MemAccessInfo, Qemu,
+    CpuPostRunHook, CpuPreRunHook, CpuRunHookId, HookState, MemAccessInfo, NewThreadHookFn, Qemu,
 };
 
+/// Get a C-compatible function pointer from the input hook.
+/// If the hook was already a c function, nothing is done.
+///
+/// h: input hook
+/// replacement: C-compatible function to call when C side should call the hook
+/// fntype: type used to cast h into replacement
 macro_rules! get_raw_hook {
     ($h:expr, $replacement:expr, $fntype:ty) => {
         match $h {
@@ -100,7 +106,6 @@ pub struct EmulatorModules<ET, S>
 where
     S: UsesInput,
 {
-    qemu: Qemu,
     modules: Pin<Box<ET>>,
     hooks: EmulatorHooks<ET, S>,
     phantom: PhantomData<S>,
@@ -679,10 +684,7 @@ where
         }
     }
 
-    pub fn backdoor_function(
-        &self,
-        hook: fn(&mut EmulatorModules<ET, S>, Option<&mut S>, cpu: CPUArchStatePtr, pc: GuestAddr),
-    ) -> BackdoorHookId {
+    pub fn backdoor_function(&self, hook: BackdoorHookFn<ET, S>) -> BackdoorHookId {
         unsafe {
             self.qemu_hooks
                 .add_backdoor_hook(transmute(hook), func_backdoor_hook_wrapper::<ET, S>)
@@ -715,15 +717,7 @@ where
         }
     }
 
-    pub fn thread_creation_function(
-        &mut self,
-        hook: fn(
-            &mut EmulatorModules<ET, S>,
-            Option<&mut S>,
-            env: CPUArchStatePtr,
-            tid: u32,
-        ) -> bool,
-    ) -> NewThreadHookId {
+    pub fn thread_creation_function(&mut self, hook: NewThreadHookFn<ET, S>) -> NewThreadHookId {
         unsafe {
             self.qemu_hooks
                 .add_new_thread_hook(transmute(hook), func_new_thread_hook_wrapper::<ET, S>)
@@ -965,7 +959,7 @@ where
     pub fn instruction_function(
         &mut self,
         addr: GuestAddr,
-        hook: fn(&mut EmulatorModules<ET, S>, Option<&mut S>, GuestAddr),
+        hook: InstructionHookFn<ET, S>,
         invalidate_block: bool,
     ) -> InstructionHookId {
         self.hooks
@@ -1076,15 +1070,7 @@ where
         self.hooks.thread_creation(hook)
     }
 
-    pub fn thread_creation_function(
-        &mut self,
-        hook: fn(
-            &mut EmulatorModules<ET, S>,
-            Option<&mut S>,
-            env: CPUArchStatePtr,
-            tid: u32,
-        ) -> bool,
-    ) -> NewThreadHookId {
+    pub fn thread_creation_function(&mut self, hook: NewThreadHookFn<ET, S>) -> NewThreadHookId {
         self.hooks.thread_creation_function(hook)
     }
 
@@ -1101,13 +1087,8 @@ where
     ET: EmulatorModuleTuple<S>,
     S: UsesInput + Unpin,
 {
-    pub(super) fn new(
-        qemu: Qemu,
-        emulator_hooks: EmulatorHooks<ET, S>,
-        modules: ET,
-    ) -> Pin<Box<Self>> {
+    pub(super) fn new(emulator_hooks: EmulatorHooks<ET, S>, modules: ET) -> Pin<Box<Self>> {
         let mut modules = Box::pin(Self {
-            qemu,
             modules: Box::pin(modules),
             hooks: emulator_hooks,
             phantom: PhantomData,
@@ -1130,35 +1111,41 @@ where
         modules
     }
 
-    pub fn post_qemu_init_all(&mut self) {
+    /// Run post-QEMU init module callbacks.
+    pub unsafe fn post_qemu_init_all(&mut self, qemu: Qemu) {
         // We give access to EmulatorModuleTuple<S> during init, the compiler complains (for good reasons)
         // TODO: We should find a way to be able to check for a module without giving full access to the tuple.
         unsafe {
             self.modules_mut()
-                .post_qemu_init_all(Self::emulator_modules_mut_unchecked());
+                .post_qemu_init_all(qemu, Self::emulator_modules_mut_unchecked());
         }
     }
 
-    pub fn first_exec_all(&mut self, state: &mut S) {
+    pub fn first_exec_all(&mut self, qemu: Qemu, state: &mut S) {
         // # Safety
         // We assume that the emulator was initialized correctly
         unsafe {
             self.modules_mut()
-                .first_exec_all(Self::emulator_modules_mut_unchecked(), state);
+                .first_exec_all(qemu, Self::emulator_modules_mut_unchecked(), state);
         }
     }
 
-    pub fn pre_exec_all(&mut self, state: &mut S, input: &S::Input) {
+    pub fn pre_exec_all(&mut self, qemu: Qemu, state: &mut S, input: &S::Input) {
         // # Safety
         // We assume that the emulator was initialized correctly
         unsafe {
-            self.modules_mut()
-                .pre_exec_all(Self::emulator_modules_mut_unchecked(), state, input);
+            self.modules_mut().pre_exec_all(
+                qemu,
+                Self::emulator_modules_mut_unchecked(),
+                state,
+                input,
+            );
         }
     }
 
     pub fn post_exec_all<OT>(
         &mut self,
+        qemu: Qemu,
         state: &mut S,
         input: &S::Input,
         observers: &mut OT,
@@ -1168,6 +1155,7 @@ where
     {
         unsafe {
             self.modules_mut().post_exec_all(
+                qemu,
                 Self::emulator_modules_mut_unchecked(),
                 state,
                 input,
@@ -1199,11 +1187,6 @@ impl<ET, S> EmulatorModules<ET, S>
 where
     S: UsesInput,
 {
-    #[must_use]
-    pub fn qemu(&self) -> Qemu {
-        self.qemu
-    }
-
     #[must_use]
     pub fn modules(&self) -> &ET {
         self.modules.as_ref().get_ref()

libafl_qemu/src/emu/mod.rs

@@ -329,17 +329,22 @@ where
         snapshot_manager: SM,
         command_manager: CM,
     ) -> Result<Self, QemuInitError> {
-        let mut emulator_hooks = unsafe { EmulatorHooks::new(QemuHooks::get_unchecked()) };
+        let emulator_hooks = unsafe { EmulatorHooks::new(QemuHooks::get_unchecked()) };
+        let mut emulator_modules = EmulatorModules::new(emulator_hooks, modules);
 
-        modules.pre_qemu_init_all(&mut emulator_hooks);
+        // TODO: fix things there properly. The biggest issue being that it creates 2 mut ref to the module with the callback being called
+        unsafe {
+            emulator_modules
+                .modules_mut()
+                .pre_qemu_init_all(EmulatorModules::<ET, S>::emulator_modules_mut_unchecked());
+        }
 
         let qemu = Qemu::init(qemu_args)?;
 
         unsafe {
             Ok(Self::new_with_qemu(
                 qemu,
-                emulator_hooks,
-                modules,
+                emulator_modules,
                 driver,
                 snapshot_manager,
                 command_manager,
@@ -352,17 +357,16 @@ where
     ///
     /// # Safety
     ///
-    /// pre-init qemu hooks should be run by then.
+    /// pre-init qemu hooks should be run before calling this.
     pub(crate) unsafe fn new_with_qemu(
         qemu: Qemu,
-        emulator_hooks: EmulatorHooks<ET, S>,
-        modules: ET,
+        emulator_modules: Pin<Box<EmulatorModules<ET, S>>>,
         driver: ED,
         snapshot_manager: SM,
         command_manager: CM,
     ) -> Self {
         let mut emulator = Emulator {
-            modules: EmulatorModules::new(qemu, emulator_hooks, modules),
+            modules: emulator_modules,
             command_manager,
             snapshot_manager,
             driver,
@@ -371,7 +375,7 @@ where
             qemu,
         };
 
-        emulator.modules.post_qemu_init_all();
+        emulator.modules.post_qemu_init_all(qemu);
 
         emulator
     }