remove the use of Box in data_notification

Author: rbran

rust/src/data_notification.rs

@@ -1,4 +1,5 @@
 use core::ffi;
+use core::pin::Pin;
 use core::ptr;
 use std::ffi::CStr;
 
@@ -40,7 +41,7 @@ macro_rules! trait_handler {
     /// Used to describe which call should be triggered by BinaryNinja.
     /// By default all calls are disabled.
     ///
-    /// Used by [CustomDataNotification::register]
+    /// Used by [CustomDataNotification::register] and [register_data_notification].
     #[derive(Default)]
     pub struct DataNotificationTriggers {
         $($fun_name: bool,)*
@@ -62,12 +63,16 @@ macro_rules! trait_handler {
         }
         )*
 
+        /// Register the data notification functions.
+        ///
+        /// If the type don't implement Unpin, use the function
+        /// [register_data_notification] with the pinned value.
         fn register<'a>(
-            self,
+            &'a mut self,
             view: &BinaryView,
             triggers: DataNotificationTriggers,
-        ) -> DataNotificationHandle<'a, Self> where Self: 'a + Sized {
-            register_data_notification(view, self, triggers)
+        ) -> DataNotificationHandle<'a, Self> where Self: 'a + Sized + Unpin {
+            register_data_notification(view, Pin::new(self), triggers)
         }
     }
     $(
@@ -79,15 +84,14 @@ macro_rules! trait_handler {
         handle.$fun_name($($value_calculated),*)
     }
     )*
-    fn register_data_notification<'a, H: CustomDataNotification + 'a>(
+    pub fn register_data_notification<'a, H: 'a + CustomDataNotification>(
         view: &BinaryView,
-        notify: H,
+        notify: Pin<&'a H>,
         triggers: DataNotificationTriggers,
     ) -> DataNotificationHandle<'a, H> {
-        // SAFETY: this leak is undone on drop
-        let leak_notify = Box::leak(Box::new(notify));
+        let notify = unsafe { Pin::into_inner_unchecked(notify) as *const H };
         let mut handle = BNBinaryDataNotification {
-            context: leak_notify as *mut _ as *mut ffi::c_void,
+            context: notify as *const ffi::c_void as *mut ffi::c_void,
             $($ffi_param_name: triggers.$fun_name.then_some($fun_name::<H>)),*
         };
         unsafe { BNRegisterDataNotification(view.handle, &mut handle) };
@@ -140,7 +144,7 @@ macro_rules! trait_handler {
     )*
 
         pub fn register(
-            self,
+            &'a self,
             view: &BinaryView,
         ) -> DataNotificationHandle<'a, Self> {
             let mut triggers = DataNotificationTriggers::default();
@@ -149,7 +153,7 @@ macro_rules! trait_handler {
                 triggers = triggers.$fun_name();
             }
             )*
-            register_data_notification(view, self, triggers)
+            register_data_notification(view, Pin::new(self), triggers)
         }
     }
 
@@ -402,36 +406,14 @@ trait_handler! {
     ),
 }
 
-pub struct DataNotificationHandle<'a, T: CustomDataNotification>
-where
-    T: 'a,
-{
+pub struct DataNotificationHandle<'a, T> {
     bv: Ref<BinaryView>,
     handle: BNBinaryDataNotification,
     _life: std::marker::PhantomData<&'a T>,
 }
 
-impl<T: CustomDataNotification> DataNotificationHandle<'_, T> {
-    unsafe fn unregister_bv(&mut self) {
-        BNUnregisterDataNotification(self.bv.handle, &mut self.handle)
-    }
-
-    unsafe fn extract_context(&mut self) -> Box<T> {
-        Box::from_raw(self.handle.context as *mut T)
-    }
-
-    pub fn unregister(self) -> T {
-        // NOTE don't drop the ctxt, return it
-        let mut slf = core::mem::ManuallyDrop::new(self);
-        unsafe { slf.unregister_bv() };
-        unsafe { *slf.extract_context() }
-    }
-}
-
-impl<T: CustomDataNotification> Drop for DataNotificationHandle<'_, T> {
+impl<'a, T> Drop for DataNotificationHandle<'a, T> {
     fn drop(&mut self) {
-        unsafe { self.unregister_bv() };
-        // drop context, avoid memory leak
-        let _ctxt = unsafe { self.extract_context() };
+        unsafe { BNUnregisterDataNotification(self.bv.handle, &mut self.handle) }
     }
 }

rust/tests/data_notification.rs

@@ -13,11 +13,12 @@ fn test_data_notification_dyn_closure() {
     let bv = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
 
     let mut func_updated_count = 0usize;
-    let custom = DataNotificationClosure::default()
-        .function_updated(|_bv: &BinaryView, _func: &Function| {
+    let custom = DataNotificationClosure::default().function_updated(
+        |_bv: &BinaryView, _func: &Function| {
             func_updated_count += 1;
-        })
-        .register(&bv);
+        },
+    );
+    custom.register(&bv);
 
     let funcs = bv.functions();
     for func in &funcs {
@@ -30,7 +31,7 @@ fn test_data_notification_dyn_closure() {
             None,
         );
     }
-    custom.unregister();
+    drop(custom);
 
     assert_eq!(funcs.len(), func_updated_count);
 }
@@ -53,15 +54,16 @@ fn test_data_notification_impl() {
     }
 
     let triggers = DataNotificationTriggers::default().tag_added();
-    let tags_lock = Tag::default().register(&bv, triggers);
+    let mut tags = Tag::default();
+    let tags_lock = tags.register(&bv, triggers);
 
     let funcs = bv.functions();
     for (i, func) in funcs.iter().enumerate() {
         let crash = bv.create_tag_type("Test", "🚧");
         func.add_tag(&crash, "Dummy tag", Some(i.try_into().unwrap()), true, None);
     }
 
-    let tags = tags_lock.unregister();
+    drop(tags_lock);
 
     assert_eq!(funcs.len(), tags.tags);
 }