Pre init module hooks (#2704)

* differenciate pre qemu init and post qemu init hooks

* api breakage: Emulator::new_with_qemu is not public anymore.

libafl_qemu/src/emu/builder.rs

@@ -12,8 +12,8 @@ use crate::{
     command::{CommandManager, NopCommandManager, StdCommandManager},
     config::QemuConfig,
     modules::{EmulatorModule, EmulatorModuleTuple},
-    Emulator, NopEmulatorDriver, NopSnapshotManager, Qemu, QemuInitError, StdEmulatorDriver,
+    Emulator, EmulatorHooks, NopEmulatorDriver, NopSnapshotManager, Qemu, QemuHooks, QemuInitError,
-    StdSnapshotManager,
+    StdEmulatorDriver, StdSnapshotManager,
 };
 
 #[derive(Clone, Debug)]
@@ -118,6 +118,10 @@ where
     {
         let qemu_builder = self.qemu_builder.ok_or(QemuInitError::EmptyArgs)?;
 
+        let mut emulator_hooks = unsafe { EmulatorHooks::new(QemuHooks::get_unchecked()) };
+
+        self.modules.pre_qemu_init_all(&mut emulator_hooks);
+
         let qemu: Qemu = match qemu_builder {
             QemuBuilder::Qemu(qemu) => qemu,
             QemuBuilder::QemuConfig(qemu_config) => {
@@ -127,13 +131,16 @@ where
             QemuBuilder::QemuString(qemu_string) => Qemu::init(&qemu_string)?,
         };
 
-        Emulator::new_with_qemu(
+        unsafe {
+            Ok(Emulator::new_with_qemu(
                 qemu,
+                emulator_hooks,
                 self.modules,
                 self.driver,
                 self.snapshot_manager,
                 self.command_manager,
-        )
+            ))
+        }
     }
 }
 

libafl_qemu/src/emu/hooks.rs

@@ -64,7 +64,7 @@ macro_rules! hook_to_repr {
     };
 }
 
-static mut EMULATOR_TOOLS: *mut () = ptr::null_mut();
+static mut EMULATOR_MODULES: *mut () = ptr::null_mut();
 
 #[cfg(feature = "usermode")]
 pub extern "C" fn crash_hook_wrapper<ET, S>(target_sig: i32)
@@ -919,14 +919,14 @@ where
     pub unsafe fn emulator_modules_mut_unchecked<'a>() -> &'a mut EmulatorModules<ET, S> {
         #[cfg(debug_assertions)]
         {
-            (EMULATOR_TOOLS as *mut EmulatorModules<ET, S>)
+            (EMULATOR_MODULES as *mut EmulatorModules<ET, S>)
                 .as_mut()
                 .unwrap()
         }
 
         #[cfg(not(debug_assertions))]
         {
-            &mut *(EMULATOR_TOOLS as *mut EmulatorModules<ET, S>)
+            &mut *(EMULATOR_MODULES as *mut EmulatorModules<ET, S>)
         }
     }
 
@@ -940,7 +940,7 @@ where
     /// generic use (it will suppose they are the same as the ones used at initialization time).
     #[must_use]
     pub unsafe fn emulator_modules_mut<'a>() -> Option<&'a mut EmulatorModules<ET, S>> {
-        unsafe { (EMULATOR_TOOLS as *mut EmulatorModules<ET, S>).as_mut() }
+        unsafe { (EMULATOR_MODULES as *mut EmulatorModules<ET, S>).as_mut() }
     }
 }
 
@@ -1101,11 +1101,15 @@ where
     ET: EmulatorModuleTuple<S>,
     S: UsesInput + Unpin,
 {
-    pub(super) fn new(qemu: Qemu, modules: ET) -> Pin<Box<Self>> {
+    pub(super) fn new(
+        qemu: Qemu,
+        emulator_hooks: EmulatorHooks<ET, S>,
+        modules: ET,
+    ) -> Pin<Box<Self>> {
         let mut modules = Box::pin(Self {
             qemu,
             modules: Box::pin(modules),
-            hooks: EmulatorHooks::default(),
+            hooks: emulator_hooks,
             phantom: PhantomData,
         });
 
@@ -1116,22 +1120,23 @@ where
 
         // Set global EmulatorModules pointer
         unsafe {
-            if EMULATOR_TOOLS.is_null() {
+            if EMULATOR_MODULES.is_null() {
-                EMULATOR_TOOLS = ptr::from_mut::<Self>(modules.as_mut().get_mut()) as *mut ();
+                EMULATOR_MODULES = ptr::from_mut::<Self>(modules.as_mut().get_mut()) as *mut ();
             } else {
                 panic!("Emulator Modules have already been set and is still active. It is not supported to have multiple instances of `EmulatorModules` at the same time yet.")
             }
         }
 
-        unsafe {
-            // 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.
         modules
-                .modules
-                .init_modules_all(Self::emulator_modules_mut_unchecked());
     }
 
-        modules
+    pub fn post_qemu_init_all(&mut self) {
+        // 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());
+        }
     }
 
     pub fn first_exec_all(&mut self, state: &mut S) {
@@ -1336,7 +1341,7 @@ where
     fn drop(&mut self) {
         // Make the global pointer null at drop time
         unsafe {
-            EMULATOR_TOOLS = ptr::null_mut();
+            EMULATOR_MODULES = ptr::null_mut();
         }
     }
 }

libafl_qemu/src/emu/mod.rs

@@ -19,8 +19,8 @@ use crate::{
     command::{CommandError, CommandManager, NopCommandManager, StdCommandManager},
     modules::EmulatorModuleTuple,
     sync_exit::SyncExit,
-    Qemu, QemuExitError, QemuExitReason, QemuInitError, QemuMemoryChunk, QemuShutdownCause, Regs,
+    Qemu, QemuExitError, QemuExitReason, QemuHooks, QemuInitError, QemuMemoryChunk,
-    CPU,
+    QemuShutdownCause, Regs, CPU,
 };
 
 mod hooks;
@@ -325,31 +325,55 @@ where
     pub fn new(
         qemu_args: &[String],
         modules: ET,
-        drivers: ED,
-        snapshot_manager: SM,
-        command_manager: CM,
-    ) -> Result<Self, QemuInitError> {
-        let qemu = Qemu::init(qemu_args)?;
-
-        Self::new_with_qemu(qemu, modules, drivers, snapshot_manager, command_manager)
-    }
-
-    pub fn new_with_qemu(
-        qemu: Qemu,
-        modules: ET,
         driver: ED,
         snapshot_manager: SM,
         command_manager: CM,
     ) -> Result<Self, QemuInitError> {
-        Ok(Emulator {
+        let mut emulator_hooks = unsafe { EmulatorHooks::new(QemuHooks::get_unchecked()) };
-            modules: EmulatorModules::new(qemu, modules),
+
+        modules.pre_qemu_init_all(&mut emulator_hooks);
+
+        let qemu = Qemu::init(qemu_args)?;
+
+        unsafe {
+            Ok(Self::new_with_qemu(
+                qemu,
+                emulator_hooks,
+                modules,
+                driver,
+                snapshot_manager,
+                command_manager,
+            ))
+        }
+    }
+
+    /// New emulator with already initialized QEMU.
+    /// We suppose modules init hooks have already been run.
+    ///
+    /// # Safety
+    ///
+    /// pre-init qemu hooks should be run by then.
+    pub(crate) unsafe fn new_with_qemu(
+        qemu: Qemu,
+        emulator_hooks: EmulatorHooks<ET, S>,
+        modules: ET,
+        driver: ED,
+        snapshot_manager: SM,
+        command_manager: CM,
+    ) -> Self {
+        let mut emulator = Emulator {
+            modules: EmulatorModules::new(qemu, emulator_hooks, modules),
             command_manager,
             snapshot_manager,
             driver,
             breakpoints_by_addr: RefCell::new(HashMap::new()),
             breakpoints_by_id: RefCell::new(HashMap::new()),
             qemu,
-        })
+        };
+
+        emulator.modules.post_qemu_init_all();
+
+        emulator
     }
 }
 

libafl_qemu/src/modules/calls.rs

@@ -399,7 +399,7 @@ where
     #[cfg(feature = "systemmode")]
     type ModulePageFilter = NopPageFilter;
 
-    fn init_module<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    fn post_qemu_init<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {

libafl_qemu/src/modules/drcov.rs

@@ -264,7 +264,7 @@ where
     #[cfg(feature = "systemmode")]
     type ModulePageFilter = NopPageFilter;
 
-    fn init_module<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    fn post_qemu_init<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {

libafl_qemu/src/modules/mod.rs

@@ -40,7 +40,7 @@ pub mod drcov;
 #[cfg(not(cpu_target = "hexagon"))]
 pub use drcov::{DrCovMetadata, DrCovModule, DrCovModuleBuilder};
 
-use crate::{emu::EmulatorModules, Qemu};
+use crate::{emu::EmulatorModules, EmulatorHooks, Qemu};
 
 /// A module for `libafl_qemu`.
 // TODO remove 'static when specialization will be stable
@@ -55,10 +55,19 @@ where
 
     const HOOKS_DO_SIDE_EFFECTS: bool = true;
 
-    /// Initialize the module, mostly used to install some hooks early.
+    /// Hook run **before** QEMU is initialized.
     /// This is always run when Emulator gets initialized, in any case.
-    /// Install here hooks that should be alive for the whole execution of the VM.
+    /// Install here hooks that should be alive for the whole execution of the VM, even before QEMU gets initialized.
-    fn init_module<ET>(&self, _emulator_modules: &mut EmulatorModules<ET, S>)
+    fn pre_qemu_init<ET>(&self, _emulator_hooks: &mut EmulatorHooks<ET, S>)
+    where
+        ET: EmulatorModuleTuple<S>,
+    {
+    }
+
+    /// Hook run **after** QEMU is initialized.
+    /// This is always run when Emulator gets initialized, in any case.
+    /// Install here hooks that should be alive for the whole execution of the VM, after QEMU gets initialized.
+    fn post_qemu_init<ET>(&self, _emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {
@@ -137,7 +146,11 @@ where
 {
     const HOOKS_DO_SIDE_EFFECTS: bool;
 
-    fn init_modules_all<ET>(&self, _emulator_modules: &mut EmulatorModules<ET, S>)
+    fn pre_qemu_init_all<ET>(&self, _emulator_hooks: &mut EmulatorHooks<ET, S>)
+    where
+        ET: EmulatorModuleTuple<S>;
+
+    fn post_qemu_init_all<ET>(&self, _emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>;
 
@@ -186,7 +199,13 @@ where
 {
     const HOOKS_DO_SIDE_EFFECTS: bool = false;
 
-    fn init_modules_all<ET>(&self, _emulator_modules: &mut EmulatorModules<ET, S>)
+    fn pre_qemu_init_all<ET>(&self, _emulator_hooks: &mut EmulatorHooks<ET, S>)
+    where
+        ET: EmulatorModuleTuple<S>,
+    {
+    }
+
+    fn post_qemu_init_all<ET>(&self, _emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {
@@ -239,12 +258,20 @@ where
 {
     const HOOKS_DO_SIDE_EFFECTS: bool = Head::HOOKS_DO_SIDE_EFFECTS || Tail::HOOKS_DO_SIDE_EFFECTS;
 
-    fn init_modules_all<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    fn pre_qemu_init_all<ET>(&self, emulator_hooks: &mut EmulatorHooks<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {
-        self.0.init_module(emulator_modules);
+        self.0.pre_qemu_init(emulator_hooks);
-        self.1.init_modules_all(emulator_modules);
+        self.1.pre_qemu_init_all(emulator_hooks);
+    }
+
+    fn post_qemu_init_all<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    where
+        ET: EmulatorModuleTuple<S>,
+    {
+        self.0.post_qemu_init(emulator_modules);
+        self.1.post_qemu_init_all(emulator_modules);
     }
 
     fn first_exec_all<ET>(&mut self, emulator_modules: &mut EmulatorModules<ET, S>, state: &mut S)

libafl_qemu/src/modules/usermode/asan.rs

@@ -929,7 +929,7 @@ where
     type ModuleAddressFilter = StdAddressFilter;
     const HOOKS_DO_SIDE_EFFECTS: bool = false;
 
-    fn init_module<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    fn post_qemu_init<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {

libafl_qemu/src/modules/usermode/injections.rs

@@ -262,7 +262,7 @@ where
 {
     type ModuleAddressFilter = NopAddressFilter;
 
-    fn init_module<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    fn post_qemu_init<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {

libafl_qemu/src/modules/usermode/snapshot.rs

@@ -675,7 +675,7 @@ where
 {
     type ModuleAddressFilter = NopAddressFilter;
 
-    fn init_module<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
+    fn post_qemu_init<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
     where
         ET: EmulatorModuleTuple<S>,
     {