From 2977c9ca47a6088f964b1e77438b750fae7518b7 Mon Sep 17 00:00:00 2001
From: Bet4 <0xbet4@gmail.com>
Date: Fri, 21 Jan 2022 22:37:15 +0800
Subject: [PATCH] Support to pass unicorn handle to rust through FFI

---
 bindings/rust/src/lib.rs           | 25 ++++++++++++++++++++++++-
 bindings/rust/src/unicorn_const.rs | 19 +++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs
index 203b6e0d..4a7e6a55 100644
--- a/bindings/rust/src/lib.rs
+++ b/bindings/rust/src/lib.rs
@@ -128,6 +128,7 @@ impl<'a> MmioCallbackScope<'a> {
 
 pub struct UnicornInner<'a, D> {
     pub uc: uc_handle,
+    pub ffi: bool,
     pub arch: Arch,
     /// to keep ownership over the hook for this uc instance's lifetime
     pub hooks: Vec<(ffi::uc_hook, Box<dyn ffi::IsUcHook<'a> + 'a>)>,
@@ -139,7 +140,7 @@ pub struct UnicornInner<'a, D> {
 /// Drop UC
 impl<'a, D> Drop for UnicornInner<'a, D> {
     fn drop(&mut self) {
-        if !self.uc.is_null() {
+        if !self.ffi && !self.uc.is_null() {
             unsafe { ffi::uc_close(self.uc) };
         }
         self.uc = ptr::null_mut();
@@ -157,6 +158,27 @@ impl<'a> Unicorn<'a, ()> {
     pub fn new(arch: Arch, mode: Mode) -> Result<Unicorn<'a, ()>, uc_error> {
         Self::new_with_data(arch, mode, ())
     }
+
+    pub fn new_with_handle(handle: uc_handle) -> Result<Unicorn<'a, ()>, uc_error> {
+        if handle == ptr::null_mut() {
+            return Err(uc_error::HANDLE);
+        }
+        let mut arch: libc::size_t = Default::default();
+        let err = unsafe { ffi::uc_query(handle, Query::ARCH, &mut arch) };
+        if err != uc_error::OK {
+            return Err(err);
+        }
+        Ok(Unicorn {
+            inner: Rc::new(UnsafeCell::from(UnicornInner {
+                uc: handle,
+                ffi: true,
+                arch: arch.try_into()?,
+                data: (),
+                hooks: vec![],
+                mmio_callbacks: vec![],
+            })),
+        })
+    }
 }
 
 impl<'a, D> Unicorn<'a, D>
@@ -172,6 +194,7 @@ where
             Ok(Unicorn {
                 inner: Rc::new(UnsafeCell::from(UnicornInner {
                     uc: handle,
+                    ffi: false,
                     arch,
                     data,
                     hooks: vec![],
diff --git a/bindings/rust/src/unicorn_const.rs b/bindings/rust/src/unicorn_const.rs
index 01975bdf..ddc7878a 100644
--- a/bindings/rust/src/unicorn_const.rs
+++ b/bindings/rust/src/unicorn_const.rs
@@ -131,6 +131,25 @@ pub enum Arch {
     MAX = 9,
 }
 
+impl TryFrom<usize> for Arch {
+    type Error = uc_error;
+
+    fn try_from(v: usize) -> Result<Self, Self::Error> {
+        match v {
+            x if x == Self::ARM as usize => Ok(Self::ARM),
+            x if x == Self::ARM64 as usize => Ok(Self::ARM64),
+            x if x == Self::MIPS as usize => Ok(Self::MIPS),
+            x if x == Self::X86 as usize => Ok(Self::X86),
+            x if x == Self::PPC as usize => Ok(Self::PPC),
+            x if x == Self::SPARC as usize => Ok(Self::SPARC),
+            x if x == Self::M68K as usize => Ok(Self::M68K),
+            x if x == Self::RISCV as usize => Ok(Self::RISCV),
+            x if x == Self::MAX as usize => Ok(Self::MAX),
+            _ => Err(uc_error::ARCH),
+        }
+    }
+}
+
 bitflags! {
     #[repr(C)]
     pub struct Mode: i32 {