Qemu new syscall hook and more python API (#306)

* new syscall hook

* expose more qemu to pylibafl

* hook syscalls from python

* update python example

* clippy

* clippy

bindings/pylibafl/src/lib.rs

@@ -1,5 +1,5 @@
-use libafl_sugar;
 use libafl_qemu;
+use libafl_sugar;
 use pyo3::prelude::*;
 
 #[pymodule]

fuzzers/python_qemu/fuzzer.py

@@ -3,11 +3,7 @@
 from pylibafl import sugar, qemu
 import lief
 
-RSI = 4
+MAX_SIZE = 0x100
-RDI = 5
-RSP = 7
-RIP = 8
-
 BINARY_PATH = './a.out'
 
 qemu.init(['qemu-x86_64', BINARY_PATH], [])
@@ -21,27 +17,26 @@ print('LLVMFuzzerTestOneInput @ 0x%x' % test_one_input)
 qemu.set_breakpoint(test_one_input)
 qemu.run()
 
-buf = qemu.read_reg(RDI)
+sp = qemu.read_reg(qemu.amd64.Rsp)
-size = qemu.read_reg(RSI)
-sp = qemu.read_reg(RSP)
-print('buf  = 0x%x' % buf)
-print('size = 0x%x' % size)
 print('SP   = 0x%x' % sp)
 
 retaddr = int.from_bytes(qemu.read_mem(sp, 8), 'little')
 print('RET  = 0x%x' % retaddr)
 
+inp = qemu.map_private(0, MAX_SIZE, qemu.mmap.ReadWrite)
+assert(inp > 0)
+
 qemu.remove_breakpoint(test_one_input)
 qemu.set_breakpoint(retaddr)
 
 def harness(b):
-    if len(b) > size:
+    if len(b) > MAX_SIZE:
-        b = b[:size]
+        b = b[:MAX_SIZE]
-    qemu.write_mem(buf, b)
+    qemu.write_mem(inp, b)
-    qemu.write_reg(RSI, size)
+    qemu.write_reg(qemu.amd64.Rsi, len(b))
-    qemu.write_reg(RDI, buf)
+    qemu.write_reg(qemu.amd64.Rdi, inp)
-    qemu.write_reg(RSP, sp)
+    qemu.write_reg(qemu.amd64.Rsp, sp)
-    qemu.write_reg(RIP, test_one_input)
+    qemu.write_reg(qemu.amd64.Rip, test_one_input)
     qemu.run()
 
 fuzz = sugar.QemuBytesCoverageSugar(['./in'], './out', 3456, [0,1,2,3])

libafl/src/executors/inprocess.rs

@@ -285,6 +285,21 @@ pub static mut GLOBAL_STATE: InProcessExecutorHandlerData = InProcessExecutorHan
     timer_queue: ptr::null_mut(),
 };
 
+#[must_use]
+pub fn inprocess_run_state<'a, S>() -> Option<&'a mut S> {
+    unsafe { (GLOBAL_STATE.state_ptr as *mut S).as_mut() }
+}
+
+#[must_use]
+pub fn inprocess_run_event_manager<'a, EM>() -> Option<&'a mut EM> {
+    unsafe { (GLOBAL_STATE.event_mgr_ptr as *mut EM).as_mut() }
+}
+
+#[must_use]
+pub fn inprocess_run_event_fuzzer<'a, F>() -> Option<&'a mut F> {
+    unsafe { (GLOBAL_STATE.fuzzer_ptr as *mut F).as_mut() }
+}
+
 #[cfg(unix)]
 mod unix_signal_handler {
     use alloc::vec::Vec;
@@ -382,53 +397,53 @@ mod unix_signal_handler {
             #[cfg(feature = "std")]
             println!("TIMEOUT or SIGUSR2 happened, but currently not fuzzing.");
             return;
-        } else {
-            #[cfg(feature = "std")]
-            println!("Timeout in fuzz run.");
-            #[cfg(feature = "std")]
-            let _res = stdout().flush();
-
-            let input = (data.current_input_ptr as *const I).as_ref().unwrap();
-            data.current_input_ptr = ptr::null();
-
-            let interesting = fuzzer
-                .objective_mut()
-                .is_interesting(state, event_mgr, input, observers, &ExitKind::Timeout)
-                .expect("In timeout handler objective failure.");
-
-            if interesting {
-                let mut new_testcase = Testcase::new(input.clone());
-                new_testcase.add_metadata(ExitKind::Timeout);
-                fuzzer
-                    .objective_mut()
-                    .append_metadata(state, &mut new_testcase)
-                    .expect("Failed adding metadata");
-                state
-                    .solutions_mut()
-                    .add(new_testcase)
-                    .expect("In timeout handler solutions failure.");
-                event_mgr
-                    .fire(
-                        state,
-                        Event::Objective {
-                            objective_size: state.solutions().count(),
-                        },
-                    )
-                    .expect("Could not send timeouting input");
-            }
-
-            event_mgr.on_restart(state).unwrap();
-
-            #[cfg(feature = "std")]
-            println!("Waiting for broker...");
-            event_mgr.await_restart_safe();
-            #[cfg(feature = "std")]
-            println!("Bye!");
-
-            event_mgr.await_restart_safe();
-
-            libc::_exit(55);
         }
+
+        #[cfg(feature = "std")]
+        println!("Timeout in fuzz run.");
+        #[cfg(feature = "std")]
+        let _res = stdout().flush();
+
+        let input = (data.current_input_ptr as *const I).as_ref().unwrap();
+        data.current_input_ptr = ptr::null();
+
+        let interesting = fuzzer
+            .objective_mut()
+            .is_interesting(state, event_mgr, input, observers, &ExitKind::Timeout)
+            .expect("In timeout handler objective failure.");
+
+        if interesting {
+            let mut new_testcase = Testcase::new(input.clone());
+            new_testcase.add_metadata(ExitKind::Timeout);
+            fuzzer
+                .objective_mut()
+                .append_metadata(state, &mut new_testcase)
+                .expect("Failed adding metadata");
+            state
+                .solutions_mut()
+                .add(new_testcase)
+                .expect("In timeout handler solutions failure.");
+            event_mgr
+                .fire(
+                    state,
+                    Event::Objective {
+                        objective_size: state.solutions().count(),
+                    },
+                )
+                .expect("Could not send timeouting input");
+        }
+
+        event_mgr.on_restart(state).unwrap();
+
+        #[cfg(feature = "std")]
+        println!("Waiting for broker...");
+        event_mgr.await_restart_safe();
+        #[cfg(feature = "std")]
+        println!("Bye!");
+
+        event_mgr.await_restart_safe();
+
+        libc::_exit(55);
     }
 
     /// Crash-Handler for in-process fuzzing.

libafl_qemu/Cargo.toml

@@ -23,6 +23,8 @@ num = "0.4"
 num_enum = "0.5.1"
 goblin = "0.4.2"
 libc = "0.2.97"
+strum = "0.21"
+strum_macros = "0.21"
 #pyo3 = { version = "0.14.3", features = ["extension-module"], optional = true }
 pyo3 = { version = "0.14.3", optional = true }
 

libafl_qemu/src/amd64.rs

@@ -1,6 +1,10 @@
 use num_enum::{IntoPrimitive, TryFromPrimitive};
+use strum_macros::EnumIter;
 
-#[derive(IntoPrimitive, TryFromPrimitive, Clone, Copy)]
+#[cfg(feature = "python")]
+use pyo3::prelude::*;
+
+#[derive(IntoPrimitive, TryFromPrimitive, Debug, Clone, Copy, EnumIter)]
 #[repr(i32)]
 pub enum Amd64Regs {
     Rax = 0,
@@ -23,6 +27,14 @@ pub enum Amd64Regs {
     Rflags = 17,
 }
 
+#[cfg(feature = "python")]
+impl IntoPy<PyObject> for Amd64Regs {
+    fn into_py(self, py: Python) -> PyObject {
+        let n: i32 = self.into();
+        n.into_py(py)
+    }
+}
+
 #[allow(non_upper_case_globals)]
 pub const TARGET_NR_read: i32 = 0;
 #[allow(non_upper_case_globals)]

libafl_qemu/src/emu.rs

@@ -10,10 +10,14 @@ use core::{
 use num::Num;
 use num_enum::{IntoPrimitive, TryFromPrimitive};
 use std::{slice::from_raw_parts, str::from_utf8_unchecked};
+use strum_macros::EnumIter;
+
+#[cfg(feature = "python")]
+use pyo3::{prelude::*, PyIterProtocol};
 
 pub const SKIP_EXEC_HOOK: u64 = u64::MAX;
 
-#[derive(IntoPrimitive, TryFromPrimitive, Clone, Copy)]
+#[derive(IntoPrimitive, TryFromPrimitive, Debug, Clone, Copy, EnumIter)]
 #[repr(i32)]
 pub enum MmapPerms {
     Read = libc::PROT_READ,
@@ -25,13 +29,62 @@ pub enum MmapPerms {
     ReadWriteExecute = libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
 }
 
-#[repr(C)]
+#[cfg(feature = "python")]
-pub struct SyscallHookResult {
+impl IntoPy<PyObject> for MmapPerms {
-    retval: u64,
+    fn into_py(self, py: Python) -> PyObject {
-    skip_syscall: bool,
+        let n: i32 = self.into();
+        n.into_py(py)
+    }
 }
 
 #[repr(C)]
+#[cfg_attr(feature = "python", pyclass)]
+#[derive(FromPyObject)]
+pub struct SyscallHookResult {
+    pub retval: u64,
+    pub skip_syscall: bool,
+}
+
+#[cfg(feature = "python")]
+#[pymethods]
+impl SyscallHookResult {
+    #[new]
+    #[must_use]
+    pub fn new(value: Option<u64>) -> Self {
+        if let Some(v) = value {
+            Self {
+                retval: v,
+                skip_syscall: true,
+            }
+        } else {
+            Self {
+                retval: 0,
+                skip_syscall: false,
+            }
+        }
+    }
+}
+
+#[cfg(not(feature = "python"))]
+impl SyscallHookResult {
+    #[must_use]
+    pub fn new(value: Option<u64>) -> Self {
+        if let Some(v) = value {
+            Self {
+                retval: v,
+                skip_syscall: true,
+            }
+        } else {
+            Self {
+                retval: 0,
+                skip_syscall: false,
+            }
+        }
+    }
+}
+
+#[repr(C)]
+#[cfg_attr(feature = "python", pyclass(unsendable))]
 pub struct MapInfo {
     start: u64,
     end: u64,
@@ -41,6 +94,7 @@ pub struct MapInfo {
     is_priv: i32,
 }
 
+#[cfg_attr(feature = "python", pymethods)]
 impl MapInfo {
     #[must_use]
     pub fn start(&self) -> u64 {
@@ -161,15 +215,15 @@ pub fn init(args: &[String], env: &[(String, String)]) -> i32 {
     }
 }
 
+#[cfg_attr(feature = "python", pyclass(unsendable))]
 pub struct GuestMaps {
     orig_c_iter: *const c_void,
     c_iter: *const c_void,
 }
 
-impl GuestMaps {
+impl Default for GuestMaps {
     #[must_use]
-    #[allow(clippy::new_without_default)]
+    fn default() -> Self {
-    pub fn new() -> Self {
         unsafe {
             let maps = read_self_maps();
             Self {
@@ -180,6 +234,24 @@ impl GuestMaps {
     }
 }
 
+#[cfg(feature = "python")]
+#[pymethods]
+impl GuestMaps {
+    #[new]
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+#[cfg(not(feature = "python"))]
+impl GuestMaps {
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
 impl Iterator for GuestMaps {
     type Item = MapInfo;
 
@@ -200,6 +272,17 @@ impl Iterator for GuestMaps {
     }
 }
 
+#[cfg(feature = "python")]
+#[pyproto]
+impl PyIterProtocol for GuestMaps {
+    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
+    }
+    fn __next__(mut slf: PyRefMut<Self>) -> Option<PyObject> {
+        Python::with_gil(|py| slf.next().map(|x| x.into_py(py)))
+    }
+}
+
 impl Drop for GuestMaps {
     fn drop(&mut self) {
         unsafe {

libafl_qemu/src/executor.rs

@@ -20,6 +20,7 @@ static mut GEN_BLOCK_HOOK_PTR: *const c_void = ptr::null();
 static mut GEN_READ_HOOK_PTR: *const c_void = ptr::null();
 static mut GEN_WRITE_HOOK_PTR: *const c_void = ptr::null();
 static mut GEN_CMP_HOOK_PTR: *const c_void = ptr::null();
+static mut SYSCALL_HOOK_PTR: *const c_void = ptr::null();
 
 extern "C" fn gen_edge_hook_wrapper<S>(src: u64, dst: u64) -> u64 {
     unsafe {
@@ -61,6 +62,25 @@ extern "C" fn gen_cmp_hook_wrapper<S>(pc: u64, size: u32) -> u64 {
     }
 }
 
+extern "C" fn gen_syscall_hook_wrapper<S>(
+    sys_num: i32,
+    a0: u64,
+    a1: u64,
+    a2: u64,
+    a3: u64,
+    a4: u64,
+    a5: u64,
+    a6: u64,
+    a7: u64,
+) -> SyscallHookResult {
+    unsafe {
+        let state = (GLOBAL_STATE.state_ptr as *mut S).as_mut().unwrap();
+        let func: fn(&mut S, i32, u64, u64, u64, u64, u64, u64, u64, u64) -> SyscallHookResult =
+            transmute(SYSCALL_HOOK_PTR);
+        (func)(state, sys_num, a0, a1, a2, a3, a4, a5, a6, a7)
+    }
+}
+
 pub struct QemuExecutor<'a, H, I, OT, S>
 where
     H: FnMut(&I) -> ExitKind,
@@ -240,6 +260,17 @@ where
     ) {
         emu::set_syscall_hook(hook);
     }
+
+    #[allow(clippy::unused_self)]
+    pub fn hook_syscalls_state(
+        &self,
+        hook: fn(&mut S, sys_num: i32, u64, u64, u64, u64, u64, u64, u64, u64) -> SyscallHookResult,
+    ) {
+        unsafe {
+            SYSCALL_HOOK_PTR = hook as *const _;
+        }
+        emu::set_syscall_hook(gen_syscall_hook_wrapper::<S>);
+    }
 }
 
 impl<'a, EM, H, I, OT, S, Z> Executor<EM, I, S, Z> for QemuExecutor<'a, H, I, OT, S>

libafl_qemu/src/lib.rs

@@ -34,15 +34,20 @@ pub fn filter_qemu_args() -> Vec<String> {
 }
 
 #[cfg(all(target_os = "linux", feature = "python"))]
-use pyo3::prelude::*;
+use pyo3::{prelude::*, types::PyInt};
+
+#[cfg(all(target_os = "linux", feature = "python"))]
+static mut PY_SYSCALL_HOOK: Option<PyObject> = None;
 
 #[cfg(all(target_os = "linux", feature = "python"))]
 #[pymodule]
 #[pyo3(name = "libafl_qemu")]
 #[allow(clippy::items_after_statements)]
-pub fn python_module(_py: Python, m: &PyModule) -> PyResult<()> {
+pub fn python_module(py: Python, m: &PyModule) -> PyResult<()> {
     use core::mem::transmute;
     use pyo3::exceptions::PyValueError;
+    use std::convert::TryFrom;
+    use strum::IntoEnumIterator;
 
     #[pyfn(m)]
     #[allow(clippy::needless_pass_by_value)]
@@ -87,7 +92,7 @@ pub fn python_module(_py: Python, m: &PyModule) -> PyResult<()> {
     }
     #[pyfn(m)]
     fn g2h(addr: u64) -> u64 {
-        unsafe { emu::g2h::<*const u8>(addr) as u64 }
+        emu::g2h::<*const u8>(addr) as u64
     }
     #[pyfn(m)]
     fn h2g(addr: u64) -> u64 {
@@ -101,5 +106,81 @@ pub fn python_module(_py: Python, m: &PyModule) -> PyResult<()> {
     fn load_addr() -> u64 {
         emu::load_addr()
     }
+    #[pyfn(m)]
+    fn map_private(addr: u64, size: usize, perms: i32) -> PyResult<u64> {
+        if let Ok(p) = MmapPerms::try_from(perms) {
+            emu::map_private(addr, size, p).map_err(PyValueError::new_err)
+        } else {
+            Err(PyValueError::new_err("Invalid perms"))
+        }
+    }
+
+    extern "C" fn py_syscall_hook_wrapper(
+        sys_num: i32,
+        a0: u64,
+        a1: u64,
+        a2: u64,
+        a3: u64,
+        a4: u64,
+        a5: u64,
+        a6: u64,
+        a7: u64,
+    ) -> SyscallHookResult {
+        if let Some(obj) = unsafe { &PY_SYSCALL_HOOK } {
+            let args = (sys_num, a0, a1, a2, a3, a4, a5, a6, a7);
+            Python::with_gil(|py| {
+                let ret = obj.call1(py, args).expect("Error in the syscall hook");
+                let any = ret.as_ref(py);
+                if any.is_none() {
+                    SyscallHookResult::new(None)
+                } else {
+                    let a: Result<&PyInt, _> = any.try_into();
+                    if let Ok(i) = a {
+                        SyscallHookResult::new(Some(
+                            i.extract().expect("Invalid syscall hook return value"),
+                        ))
+                    } else {
+                        SyscallHookResult::extract(any)
+                            .expect("The syscall hook must return a SyscallHookResult")
+                    }
+                }
+            })
+        } else {
+            SyscallHookResult::new(None)
+        }
+    }
+    #[pyfn(m)]
+    fn set_syscall_hook(hook: PyObject) {
+        unsafe {
+            PY_SYSCALL_HOOK = Some(hook);
+        }
+        emu::set_syscall_hook(py_syscall_hook_wrapper);
+    }
+
+    let child_module = PyModule::new(py, "x86")?;
+    for r in x86::X86Regs::iter() {
+        let v: i32 = r.into();
+        child_module.add(&format!("{:?}", r), v)?;
+    }
+    m.add_submodule(child_module)?;
+
+    let child_module = PyModule::new(py, "amd64")?;
+    for r in amd64::Amd64Regs::iter() {
+        let v: i32 = r.into();
+        child_module.add(&format!("{:?}", r), v)?;
+    }
+    m.add_submodule(child_module)?;
+
+    let child_module = PyModule::new(py, "mmap")?;
+    for r in emu::MmapPerms::iter() {
+        let v: i32 = r.into();
+        child_module.add(&format!("{:?}", r), v)?;
+    }
+    m.add_submodule(child_module)?;
+
+    m.add_class::<emu::MapInfo>()?;
+    m.add_class::<emu::GuestMaps>()?;
+    m.add_class::<emu::SyscallHookResult>()?;
+
     Ok(())
 }

libafl_qemu/src/x86.rs

@@ -1,6 +1,10 @@
 use num_enum::{IntoPrimitive, TryFromPrimitive};
+use strum_macros::EnumIter;
 
-#[derive(IntoPrimitive, TryFromPrimitive, Clone, Copy)]
+#[cfg(feature = "python")]
+use pyo3::prelude::*;
+
+#[derive(IntoPrimitive, TryFromPrimitive, Debug, Clone, Copy, EnumIter)]
 #[repr(i32)]
 pub enum X86Regs {
     Eax = 0,
@@ -14,3 +18,11 @@ pub enum X86Regs {
     Eip = 8,
     Eflags = 9,
 }
+
+#[cfg(feature = "python")]
+impl IntoPy<PyObject> for X86Regs {
+    fn into_py(self, py: Python) -> PyObject {
+        let n: i32 = self.into();
+        n.into_py(py)
+    }
+}