From 91ce28deaccfa5ad7e884d252e4c3174dc47e91f Mon Sep 17 00:00:00 2001
From: Andrea Fioraldi <andreafioraldi@gmail.com>
Date: Fri, 1 Oct 2021 17:10:35 +0200
Subject: [PATCH] Python generic qemu hook (#314)

* python generic hook and value

* python generic hook
---
 libafl_qemu/build.rs   |  2 +-
 libafl_qemu/src/emu.rs |  6 +++---
 libafl_qemu/src/lib.rs | 25 +++++++++++++++++++++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/libafl_qemu/build.rs b/libafl_qemu/build.rs
index ec310b89bc..0e5c4deb6b 100644
--- a/libafl_qemu/build.rs
+++ b/libafl_qemu/build.rs
@@ -3,7 +3,7 @@ use which::which;
 
 const QEMU_URL: &str = "https://github.com/AFLplusplus/qemu-libafl-bridge";
 const QEMU_DIRNAME: &str = "qemu-libafl-bridge";
-const QEMU_REVISION: &str = "f9898d7db457e57f84178c3ae58b4972ad66587d";
+const QEMU_REVISION: &str = "78936b7cc7c8fcdce3858eb3a343a002ffb89f0c";
 
 fn build_dep_check(tools: &[&str]) {
     for tool in tools {
diff --git a/libafl_qemu/src/emu.rs b/libafl_qemu/src/emu.rs
index 4c458b9e49..17733db5fb 100644
--- a/libafl_qemu/src/emu.rs
+++ b/libafl_qemu/src/emu.rs
@@ -178,7 +178,7 @@ extern "C" {
     fn libafl_qemu_num_regs() -> i32;
     fn libafl_qemu_set_breakpoint(addr: u64) -> i32;
     fn libafl_qemu_remove_breakpoint(addr: u64) -> i32;
-    fn libafl_qemu_insert_hook(addr: u64, callback: /*extern "C"*/ extern fn ()) -> i32;
+    fn libafl_qemu_set_hook(addr: u64, callback: extern "C" fn(u64), val: u64) -> i32;
     fn libafl_qemu_remove_hook(addr: u64) -> i32;
     fn libafl_qemu_run() -> i32;
     fn libafl_load_addr() -> u64;
@@ -397,9 +397,9 @@ pub fn remove_breakpoint(addr: u64) {
     }
 }
 
-pub fn set_hook(addr: u64, callback: extern fn()) {
+pub fn set_hook(addr: u64, callback: extern "C" fn(u64), val: u64) {
     unsafe {
-        libafl_qemu_insert_hook(addr, callback);
+        libafl_qemu_set_hook(addr, callback, val);
     }
 }
 
diff --git a/libafl_qemu/src/lib.rs b/libafl_qemu/src/lib.rs
index 62bca35f6d..1f479d647a 100644
--- a/libafl_qemu/src/lib.rs
+++ b/libafl_qemu/src/lib.rs
@@ -45,6 +45,9 @@ 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"))]
+static mut PY_GENERIC_HOOKS: Vec<(u64, PyObject)> = vec![];
+
 #[cfg(all(target_os = "linux", feature = "python"))]
 #[pymodule]
 #[pyo3(name = "libafl_qemu")]
@@ -168,6 +171,28 @@ pub fn python_module(py: Python, m: &PyModule) -> PyResult<()> {
         emu::set_syscall_hook(py_syscall_hook_wrapper);
     }
 
+    extern "C" fn py_generic_hook_wrapper(idx: u64) {
+        let obj = unsafe { &PY_GENERIC_HOOKS[idx as usize].1 };
+        Python::with_gil(|py| {
+            obj.call0(py).expect("Error in the hook");
+        });
+    }
+    #[pyfn(m)]
+    fn set_hook(addr: u64, hook: PyObject) {
+        unsafe {
+            let idx = PY_GENERIC_HOOKS.len();
+            PY_GENERIC_HOOKS.push((addr, hook));
+            emu::set_hook(addr, py_generic_hook_wrapper, idx as u64);
+        }
+    }
+    #[pyfn(m)]
+    fn remove_hook(addr: u64) {
+        unsafe {
+            PY_GENERIC_HOOKS.retain(|(a, _)| *a != addr);
+        }
+        emu::remove_hook(addr);
+    }
+
     let x86m = PyModule::new(py, "x86")?;
     for r in x86::X86Regs::iter() {
         let v: i32 = r.into();