From 3a015153672ae16d46b15041e3c4b8e3faae4915 Mon Sep 17 00:00:00 2001
From: Sai Ashwin <sai.ashwin2409@gmail.com>
Date: Sat, 7 Dec 2024 12:19:23 +0530
Subject: [PATCH] Rust Bindings: reg_read_batch and reg_write_batch  (#2060)

* Added binding for rust reg_{read,write}_batch

* Fix reg_write_batch values pointer
---
 bindings/rust/src/ffi.rs | 12 ++++++++++
 bindings/rust/src/lib.rs | 49 ++++++++++++++++++++++++++++++++++++++++
 tests/rust-tests/main.rs | 32 ++++++++++++++++++++++++++
 3 files changed, 93 insertions(+)

diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs
index 742e22412c..7f7a205bf9 100644
--- a/bindings/rust/src/ffi.rs
+++ b/bindings/rust/src/ffi.rs
@@ -23,7 +23,19 @@ extern "C" {
     pub fn uc_errno(engine: uc_handle) -> uc_error;
     pub fn uc_strerror(error_code: uc_error) -> *const c_char;
     pub fn uc_reg_write(engine: uc_handle, regid: c_int, value: *const c_void) -> uc_error;
+    pub fn uc_reg_write_batch(
+        engine: uc_handle,
+        regids: *const c_int,
+        values: *const *const c_void,
+        count: c_int,
+    ) -> uc_error;
     pub fn uc_reg_read(engine: uc_handle, regid: c_int, value: *mut c_void) -> uc_error;
+    pub fn uc_reg_read_batch(
+        engine: uc_handle,
+        regids: *const c_int,
+        values: *const *mut c_void,
+        count: c_int,
+    ) -> uc_error;
     pub fn uc_mem_write(
         engine: uc_handle,
         address: u64,
diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs
index a1e7b16925..cd84b4eaa1 100644
--- a/bindings/rust/src/lib.rs
+++ b/bindings/rust/src/lib.rs
@@ -523,6 +523,28 @@ impl<'a, D> Unicorn<'a, D> {
             .into()
     }
 
+    /// Write values into batch of registers
+    pub fn reg_write_batch<T: Into<i32>>(
+        &self,
+        regids: &[T],
+        values: &[u64],
+        count: i32,
+    ) -> Result<(), uc_error> {
+        let mut values_ptrs: Vec<*const u64> = vec![0 as *const u64; count as usize];
+        for i in 0..values.len() {
+            values_ptrs[i as usize] = &values[i] as *const u64;
+        }
+        unsafe {
+            ffi::uc_reg_write_batch(
+                self.get_handle(),
+                regids.as_ptr() as *const i32,
+                values_ptrs.as_ptr() as *const *const c_void,
+                count,
+            )
+        }
+        .into()
+    }
+
     /// Write variable sized values into registers.
     ///
     /// The user has to make sure that the buffer length matches the register size.
@@ -540,6 +562,33 @@ impl<'a, D> Unicorn<'a, D> {
             .and(Ok(value))
     }
 
+    /// Read batch of registers
+    ///
+    /// Not to be used with registers larger than 64 bit
+    pub fn reg_read_batch<T: Into<i32>>(
+        &self,
+        regids: &[T],
+        count: i32,
+    ) -> Result<Vec<u64>, uc_error> {
+        unsafe {
+            let mut addrs_vec = vec![0u64; count as usize];
+            let addrs = addrs_vec.as_mut_slice();
+            for i in 0..count {
+                addrs[i as usize] = &mut addrs[i as usize] as *mut u64 as u64;
+            }
+            let res = ffi::uc_reg_read_batch(
+                self.get_handle(),
+                regids.as_ptr() as *const i32,
+                addrs.as_ptr() as *const *mut c_void,
+                count,
+            );
+            match res {
+                uc_error::OK => Ok(addrs_vec),
+                _ => Err(res),
+            }
+        }
+    }
+
     /// Read 128, 256 or 512 bit register value into heap allocated byte array.
     ///
     /// This adds safe support for registers >64 bit (GDTR/IDTR, XMM, YMM, ZMM, ST (x86); Q, V (arm64)).
diff --git a/tests/rust-tests/main.rs b/tests/rust-tests/main.rs
index b35e395f34..1661890e8d 100644
--- a/tests/rust-tests/main.rs
+++ b/tests/rust-tests/main.rs
@@ -815,3 +815,35 @@ fn x86_tlb_callback() {
     assert_eq!(emu.remove_hook(tlb_hook), Ok(()));
     assert_eq!(emu.remove_hook(syscall_hook), Ok(()));
 }
+
+#[test]
+fn x86_reg_rw_batch() {
+    // mov rax, 0x10, mov rbx, 0x20, mov rcx, 0x30, mov rdx, 0x40
+    let code: Vec<u8> = vec![0x48, 0xC7, 0xC0, 0x10, 0x00, 0x00, 0x00, 0x48, 0xC7, 0xC3, 0x20, 0x00, 0x00, 0x00, 0x48, 0xC7, 0xC1, 0x30, 0x00, 0x00, 0x00, 0x48, 0xC7, 0xC2, 0x40, 0x00, 0x00, 0x00];
+    let expect: Vec<u64> = vec![0x10, 0x20, 0x30, 0x40];
+    let mut emu = unicorn_engine::Unicorn::new(Arch::X86, Mode::MODE_64)
+        .expect("failed to initialize unicorn instance");
+    assert_eq!(emu.mem_map(0x1000, 0x1000, Permission::ALL), Ok(()));
+    assert_eq!(emu.mem_write(0x1000, &code), Ok(()));
+
+    assert_eq!(
+        emu.emu_start(0x1000, (0x1000 + code.len()) as u64, 0, 0),
+        Ok(())
+    );
+
+    let regids = vec![RegisterX86::RAX, RegisterX86::RBX, RegisterX86::RCX, RegisterX86::RDX];
+    assert_eq!(
+        emu.reg_read_batch(&regids, 4),
+        Ok(expect)
+    );
+    let regvals = vec![0x50, 0x60, 0x70, 0x80];
+    assert_eq!(
+        emu.reg_write_batch(&regids, &regvals, 4),
+        Ok(())
+    );
+    assert_eq!(
+        emu.reg_read_batch(&regids, 4),
+        Ok(vec![0x50, 0x60, 0x70, 0x80])
+    );
+
+}