QEMU command system refactoring (#2189)

* implemented generic command builder

* Added builder to `Emulator`.

Author: rmalmain

fuzzers/qemu_systemmode/src/fuzzer_breakpoint.rs

@@ -29,7 +29,7 @@ use libafl_bolts::{
 };
 use libafl_qemu::{
     breakpoint::Breakpoint,
-    command::{Command, EndCommand, StartCommand},
+    command::{EndCommand, StartCommand, StdCommandManager},
     edges::{edges_map_mut_ptr, QemuEdgeCoverageHelper, EDGES_MAP_SIZE_IN_USE, MAX_EDGES_FOUND},
     elf::EasyElf,
     emu::Emulator,
@@ -95,28 +95,27 @@ pub fn fuzz() {
         // Choose Exit Handler
         let emu_exit_handler = StdEmulatorExitHandler::new(emu_snapshot_manager);
 
+        // Choose Command Manager
+        let cmd_manager = StdCommandManager::new();
+
         // Create emulator
-        let emu = Emulator::new(&args, &env, emu_exit_handler).unwrap();
+        let emu = Emulator::new(&args, &env, emu_exit_handler, cmd_manager).unwrap();
 
         // Set breakpoints of interest with corresponding commands.
         emu.add_breakpoint(
             Breakpoint::with_command(
                 main_addr,
-                Command::StartCommand(StartCommand::new(EmulatorMemoryChunk::phys(
+                StartCommand::new(EmulatorMemoryChunk::phys(
                     input_addr,
                     unsafe { MAX_INPUT_SIZE } as GuestReg,
                     None,
-                ))),
+                )),
                 true,
             ),
             true,
         );
         emu.add_breakpoint(
-            Breakpoint::with_command(
-                breakpoint,
-                Command::EndCommand(EndCommand::new(Some(ExitKind::Ok))),
-                false,
-            ),
+            Breakpoint::with_command(breakpoint, EndCommand::new(Some(ExitKind::Ok)), false),
             true,
         );
 

fuzzers/qemu_systemmode/src/fuzzer_sync_exit.rs

@@ -27,6 +27,7 @@ use libafl_bolts::{
     tuples::tuple_list,
 };
 use libafl_qemu::{
+    command::StdCommandManager,
     edges::{edges_map_mut_ptr, QemuEdgeCoverageHelper, EDGES_MAP_SIZE_IN_USE, MAX_EDGES_FOUND},
     emu::Emulator,
     executor::{stateful::StatefulQemuExecutor, QemuExecutorState},
@@ -52,11 +53,15 @@ pub fn fuzz() {
         // Initialize QEMU
         let args: Vec<String> = env::args().collect();
         let env: Vec<(String, String)> = env::vars().collect();
+
         // let emu_snapshot_manager = QemuSnapshotBuilder::new(true);
         let emu_snapshot_manager = FastSnapshotManager::new(false); // Create a snapshot manager (normal or fast for now).
         let emu_exit_handler: StdEmulatorExitHandler<FastSnapshotManager> =
             StdEmulatorExitHandler::new(emu_snapshot_manager); // Create an exit handler: it is the entity taking the decision of what should be done when QEMU returns.
-        let emu = Emulator::new(&args, &env, emu_exit_handler).unwrap(); // Create the emulator
+
+        let cmd_manager = StdCommandManager::new();
+
+        let emu = Emulator::new(&args, &env, emu_exit_handler, cmd_manager).unwrap(); // Create the emulator
 
         let devices = emu.list_devices();
         println!("Devices = {:?}", devices);

libafl_qemu/Cargo.toml

@@ -91,6 +91,7 @@ serde_yaml = { version = "0.9", optional = true } # For parsing the injections y
 toml = { version = "0.8.13", optional = true } # For parsing the injections toml file
 pyo3 = { version = "0.18", optional = true , features = ["multiple-pymethods"]}
 bytes-utils = "0.1"
+typed-builder = "0.18"
 # Document all features of this crate (for `cargo doc`)
 document-features = { version = "0.2", optional = true }
 

libafl_qemu/src/breakpoint.rs

@@ -2,52 +2,138 @@ use std::{
     borrow::Borrow,
     fmt::{Display, Formatter},
     hash::{Hash, Hasher},
+    rc::Rc,
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        OnceLock,
+    },
 };
 
+use libafl::state::{HasExecutions, State};
 use libafl_qemu_sys::GuestAddr;
 
-use crate::{command::Command, Qemu};
+use crate::{
+    command::{CommandManager, IsCommand},
+    EmulatorExitHandler, Qemu, QemuHelperTuple,
+};
+
+#[repr(transparent)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
+pub struct BreakpointId(u64);
 
 // TODO: distinguish breakpoints with IDs instead of addresses to avoid collisions.
-#[derive(Debug, Clone)]
-pub struct Breakpoint {
+#[derive(Debug)]
+pub struct Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
+    id: BreakpointId,
     addr: GuestAddr,
-    cmd: Option<Command>,
+    cmd: Option<Rc<dyn IsCommand<CM, E, QT, S>>>,
     disable_on_trigger: bool,
     enabled: bool,
 }
 
-impl Hash for Breakpoint {
+impl BreakpointId {
+    pub fn new() -> Self {
+        static mut BREAKPOINT_ID_COUNTER: OnceLock<AtomicU64> = OnceLock::new();
+
+        let counter = unsafe { BREAKPOINT_ID_COUNTER.get_or_init(|| AtomicU64::new(0)) };
+
+        BreakpointId(counter.fetch_add(1, Ordering::SeqCst))
+    }
+}
+
+impl Default for BreakpointId {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<CM, E, QT, S> Hash for Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
     fn hash<H: Hasher>(&self, state: &mut H) {
-        self.addr.hash(state);
+        self.id.hash(state);
     }
 }
 
-impl PartialEq for Breakpoint {
+impl<CM, E, QT, S> PartialEq for Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
     fn eq(&self, other: &Self) -> bool {
-        self.addr == other.addr
+        self.id == other.id
     }
 }
 
-impl Eq for Breakpoint {}
+impl<CM, E, QT, S> Eq for Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
+}
 
-impl Display for Breakpoint {
+impl<CM, E, QT, S> Display for Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "Breakpoint @vaddr 0x{:x}", self.addr)
     }
 }
 
-impl Borrow<GuestAddr> for Breakpoint {
+impl<CM, E, QT, S> Borrow<BreakpointId> for Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
+    fn borrow(&self) -> &BreakpointId {
+        &self.id
+    }
+}
+
+impl<CM, E, QT, S> Borrow<GuestAddr> for Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
     fn borrow(&self) -> &GuestAddr {
         &self.addr
     }
 }
 
-impl Breakpoint {
+impl<CM, E, QT, S> Breakpoint<CM, E, QT, S>
+where
+    CM: CommandManager<E, QT, S>,
+    E: EmulatorExitHandler<QT, S>,
+    QT: QemuHelperTuple<S>,
+    S: State + HasExecutions,
+{
     // Emu will return with the breakpoint as exit reason.
     #[must_use]
     pub fn without_command(addr: GuestAddr, disable_on_trigger: bool) -> Self {
         Self {
+            id: BreakpointId::new(),
             addr,
             cmd: None,
             disable_on_trigger,
@@ -57,15 +143,25 @@ impl Breakpoint {
 
     // Emu will execute the command when it meets the breakpoint.
     #[must_use]
-    pub fn with_command(addr: GuestAddr, cmd: Command, disable_on_trigger: bool) -> Self {
+    pub fn with_command<C: IsCommand<CM, E, QT, S> + 'static>(
+        addr: GuestAddr,
+        cmd: C,
+        disable_on_trigger: bool,
+    ) -> Self {
         Self {
+            id: BreakpointId::new(),
             addr,
-            cmd: Some(cmd),
+            cmd: Some(Rc::new(cmd)),
             disable_on_trigger,
             enabled: false,
         }
     }
 
+    #[must_use]
+    pub fn id(&self) -> BreakpointId {
+        self.id
+    }
+
     #[must_use]
     pub fn addr(&self) -> GuestAddr {
         self.addr
@@ -85,11 +181,11 @@ impl Breakpoint {
         }
     }
 
-    pub fn trigger(&mut self, qemu: &Qemu) -> Option<&Command> {
+    pub fn trigger(&mut self, qemu: &Qemu) -> Option<Rc<dyn IsCommand<CM, E, QT, S>>> {
         if self.disable_on_trigger {
             self.disable(qemu);
         }
 
-        self.cmd.as_ref()
+        self.cmd.clone()
     }
 }