Fix Rust binding memory leak

unicorn-engine · Jan 18, 2023 · 014cee5 · 014cee5
1 parent 549f34f
commit 014cee5
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 54 deletions.
diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs
@@ -1,10 +1,11 @@
 #![allow(non_camel_case_types)]
 #![allow(dead_code)]
 
-use crate::Unicorn;
+use crate::{Unicorn, UnicornInner};
 
 use super::unicorn_const::{uc_error, Arch, HookType, MemRegion, MemType, Mode, Query};
-use core::ffi::c_void;
+use alloc::rc::Weak;
+use core::{cell::UnsafeCell, ffi::c_void};
 use libc::{c_char, c_int};
 
 pub type uc_handle = *mut c_void;
@@ -89,7 +90,7 @@ extern "C" {
 
 pub struct UcHook<'a, D: 'a, F: 'a> {
     pub callback: F,
-    pub uc: Unicorn<'a, D>,
+    pub uc: Weak<UnsafeCell<UnicornInner<'a, D>>>,
 }
 
 pub trait IsUcHook<'a> {}
@@ -106,8 +107,11 @@ where
     F: FnMut(&mut crate::Unicorn<D>, u64, usize) -> u64,
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, offset, size)
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, offset, size)
 }
 
 pub extern "C" fn mmio_write_callback_proxy<D, F>(
@@ -120,8 +124,11 @@ pub extern "C" fn mmio_write_callback_proxy<D, F>(
     F: FnMut(&mut crate::Unicorn<D>, u64, usize, u64),
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, offset, size, value);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, offset, size, value);
 }
 
 pub extern "C" fn code_hook_proxy<D, F>(
@@ -133,8 +140,11 @@ pub extern "C" fn code_hook_proxy<D, F>(
     F: FnMut(&mut crate::Unicorn<D>, u64, u32),
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, address, size);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, address, size);
 }
 
 pub extern "C" fn block_hook_proxy<D, F>(
@@ -146,8 +156,11 @@ pub extern "C" fn block_hook_proxy<D, F>(
     F: FnMut(&mut crate::Unicorn<D>, u64, u32),
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, address, size);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, address, size);
 }
 
 pub extern "C" fn mem_hook_proxy<D, F>(
@@ -162,17 +175,23 @@ where
     F: FnMut(&mut crate::Unicorn<D>, MemType, u64, usize, i64) -> bool,
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, mem_type, address, size as usize, value)
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, mem_type, address, size as usize, value)
 }
 
 pub extern "C" fn intr_hook_proxy<D, F>(uc: uc_handle, value: u32, user_data: *mut UcHook<D, F>)
 where
     F: FnMut(&mut crate::Unicorn<D>, u32),
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, value);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, value);
 }
 
 pub extern "C" fn insn_in_hook_proxy<D, F>(
@@ -184,17 +203,23 @@ pub extern "C" fn insn_in_hook_proxy<D, F>(
     F: FnMut(&mut crate::Unicorn<D>, u32, usize) -> u32,
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, port, size);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, port, size);
 }
 
 pub extern "C" fn insn_invalid_hook_proxy<D, F>(uc: uc_handle, user_data: *mut UcHook<D, F>) -> bool
 where
     F: FnMut(&mut crate::Unicorn<D>) -> bool,
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc)
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc)
 }
 
 pub extern "C" fn insn_out_hook_proxy<D, F>(
@@ -207,15 +232,21 @@ pub extern "C" fn insn_out_hook_proxy<D, F>(
     F: FnMut(&mut crate::Unicorn<D>, u32, usize, u32),
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc, port, size, value);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc, port, size, value);
 }
 
 pub extern "C" fn insn_sys_hook_proxy<D, F>(uc: uc_handle, user_data: *mut UcHook<D, F>)
 where
     F: FnMut(&mut crate::Unicorn<D>),
 {
     let user_data = unsafe { &mut *user_data };
-    debug_assert_eq!(uc, user_data.uc.get_handle());
-    (user_data.callback)(&mut user_data.uc);
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    (user_data.callback)(&mut user_data_uc);
 }
diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs
@@ -369,17 +369,13 @@ impl<'a, D> Unicorn<'a, D> {
         let mut read_data = read_callback.map(|c| {
             Box::new(ffi::UcHook {
                 callback: c,
-                uc: Unicorn {
-                    inner: self.inner.clone(),
-                },
+                uc: Rc::downgrade(&self.inner),
             })
         });
         let mut write_data = write_callback.map(|c| {
             Box::new(ffi::UcHook {
                 callback: c,
-                uc: Unicorn {
-                    inner: self.inner.clone(),
-                },
+                uc: Rc::downgrade(&self.inner),
             })
         });
 
@@ -586,7 +582,8 @@ impl<'a, D> Unicorn<'a, D> {
             return Err(uc_error::ARCH);
         }
 
-        let err: uc_error = unsafe { ffi::uc_reg_read(self.get_handle(), curr_reg_id, value.as_mut_ptr() as _) };
+        let err: uc_error =
+            unsafe { ffi::uc_reg_read(self.get_handle(), curr_reg_id, value.as_mut_ptr() as _) };
 
         if err == uc_error::OK {
             boxed = value.into_boxed_slice();
@@ -622,9 +619,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -654,9 +649,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -697,9 +690,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -730,9 +721,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -763,9 +752,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -796,9 +783,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -830,9 +815,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {
@@ -870,9 +853,7 @@ impl<'a, D> Unicorn<'a, D> {
         let mut hook_ptr = core::ptr::null_mut();
         let mut user_data = Box::new(ffi::UcHook {
             callback,
-            uc: Unicorn {
-                inner: self.inner.clone(),
-            },
+            uc: Rc::downgrade(&self.inner),
         });
 
         let err = unsafe {