[Rust] Move `ObjectDestructor` to own module and add some extra documentation

5 files changed, 102 insertions, 61 deletions

diff --git a/plugins/warp/src/cache.rs b/plugins/warp/src/cache.rs
index 3df2dbf8..d9a5a8d1 100644
--- a/plugins/warp/src/cache.rs
+++ b/plugins/warp/src/cache.rs
@@ -9,18 +9,14 @@ pub use type_reference::*;
 
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::function::Function as BNFunction;
+use binaryninja::object_destructor::{register_object_destructor, ObjectDestructor};
 use binaryninja::rc::Guard;
 use binaryninja::rc::Ref as BNRef;
-use binaryninja::ObjectDestructor;
 use std::hash::{DefaultHasher, Hash, Hasher};
 
 pub fn register_cache_destructor() {
-    pub static mut CACHE_DESTRUCTOR: CacheDestructor = CacheDestructor;
-    #[allow(static_mut_refs)]
-    // SAFETY: This can be done as the backing data is an opaque ZST.
-    unsafe {
-        CACHE_DESTRUCTOR.register()
-    };
+    let destructor = register_object_destructor(CacheDestructor);
+    std::mem::forget(destructor);
 }
 
 /// A unique view ID, used for caching.
diff --git a/plugins/warp/src/plugin/debug.rs b/plugins/warp/src/plugin/debug.rs
index d220d329..4198ea30 100644
--- a/plugins/warp/src/plugin/debug.rs
+++ b/plugins/warp/src/plugin/debug.rs
@@ -3,7 +3,7 @@ use crate::{build_function, cache};
 use binaryninja::binary_view::BinaryView;
 use binaryninja::command::{Command, FunctionCommand};
 use binaryninja::function::Function;
-use binaryninja::ObjectDestructor;
+use binaryninja::object_destructor::ObjectDestructor;
 
 pub struct DebugFunction;
 
diff --git a/plugins/workflow_objc/src/metadata/global_state.rs b/plugins/workflow_objc/src/metadata/global_state.rs
index db29af49..6c128e22 100644
--- a/plugins/workflow_objc/src/metadata/global_state.rs
+++ b/plugins/workflow_objc/src/metadata/global_state.rs
@@ -1,11 +1,12 @@
 use binaryninja::file_metadata::SessionId;
+use binaryninja::object_destructor::register_object_destructor;
 use binaryninja::{
     binary_view::{BinaryView, BinaryViewBase, BinaryViewExt},
     file_metadata::FileMetadata,
     metadata::Metadata,
+    object_destructor::ObjectDestructor,
     rc::Ref,
     settings::{QueryOptions, Settings},
-    ObjectDestructor,
 };
 use dashmap::DashMap;
 use once_cell::sync::Lazy;
@@ -66,8 +67,8 @@ pub struct GlobalState;
 
 impl GlobalState {
     pub fn register_cleanup() {
-        let observer = Box::leak(Box::new(ObjectLifetimeObserver));
-        observer.register();
+        let destructor = register_object_destructor(ObjectLifetimeObserver);
+        std::mem::forget(destructor);
     }
 
     fn id(bv: &BinaryView) -> SessionId {
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index 254dd2f8..5f23a9f6 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -65,6 +65,7 @@ pub mod low_level_il;
 pub mod main_thread;
 pub mod medium_level_il;
 pub mod metadata;
+pub mod object_destructor;
 pub mod platform;
 pub mod progress;
 pub mod project;
@@ -90,8 +91,6 @@ pub mod websocket;
 pub mod worker_thread;
 pub mod workflow;
 
-use crate::file_metadata::FileMetadata;
-use crate::function::Function;
 use crate::progress::{NoProgressCallback, ProgressCallback};
 use crate::string::raw_to_string;
 use binary_view::BinaryView;
@@ -428,54 +427,6 @@ pub fn memory_info() -> HashMap<String, u64> {
     usage
 }
 
-/// The trait required for receiving core object destruction callbacks.
-pub trait ObjectDestructor: 'static + Sync + Sized {
-    fn destruct_view(&self, _view: &BinaryView) {}
-    fn destruct_file_metadata(&self, _metadata: &FileMetadata) {}
-    fn destruct_function(&self, _func: &Function) {}
-
-    unsafe extern "C" fn cb_destruct_binary_view(ctxt: *mut c_void, view: *mut BNBinaryView) {
-        ffi_wrap!("ObjectDestructor::destruct_view", {
-            let view_type = &*(ctxt as *mut Self);
-            let view = BinaryView { handle: view };
-            view_type.destruct_view(&view);
-        })
-    }
-
-    unsafe extern "C" fn cb_destruct_file_metadata(ctxt: *mut c_void, file: *mut BNFileMetadata) {
-        ffi_wrap!("ObjectDestructor::destruct_file_metadata", {
-            let view_type = &*(ctxt as *mut Self);
-            let file = FileMetadata::from_raw(file);
-            view_type.destruct_file_metadata(&file);
-        })
-    }
-
-    unsafe extern "C" fn cb_destruct_function(ctxt: *mut c_void, func: *mut BNFunction) {
-        ffi_wrap!("ObjectDestructor::destruct_function", {
-            let view_type = &*(ctxt as *mut Self);
-            let func = Function { handle: func };
-            view_type.destruct_function(&func);
-        })
-    }
-
-    unsafe fn as_callbacks(&'static mut self) -> BNObjectDestructionCallbacks {
-        BNObjectDestructionCallbacks {
-            context: std::mem::transmute(&self),
-            destructBinaryView: Some(Self::cb_destruct_binary_view),
-            destructFileMetadata: Some(Self::cb_destruct_file_metadata),
-            destructFunction: Some(Self::cb_destruct_function),
-        }
-    }
-
-    fn register(&'static mut self) {
-        unsafe { BNRegisterObjectDestructionCallbacks(&mut self.as_callbacks()) };
-    }
-
-    fn unregister(&'static mut self) {
-        unsafe { BNUnregisterObjectDestructionCallbacks(&mut self.as_callbacks()) };
-    }
-}
-
 pub fn version() -> String {
     unsafe { BnString::into_string(BNGetVersionString()) }
 }
diff --git a/rust/src/object_destructor.rs b/rust/src/object_destructor.rs
new file mode 100644
index 00000000..aff5e4f0
--- /dev/null
+++ b/rust/src/object_destructor.rs
@@ -0,0 +1,93 @@
+//! Register callbacks for when core objects like [`BinaryView`]s or [`Function`]s are destroyed.
+
+use crate::binary_view::BinaryView;
+use crate::file_metadata::FileMetadata;
+use crate::function::Function;
+use binaryninjacore_sys::*;
+use std::ffi::c_void;
+
+/// Registers a destructor which will be called when certain core objects are about to be destroyed.
+///
+/// Returns a handle to the registered destructor. The destructor will be unregistered when the handle is dropped.
+///
+/// To keep the destructor alive forever, move the [`ObjectDestructorHandle`] into [`std::mem::ManuallyDrop`].
+#[must_use = "The destructor will be unregistered when the handle is dropped"]
+pub fn register_object_destructor<'a, D: ObjectDestructor>(
+    destructor: D,
+) -> ObjectDestructorHandle<'a, D> {
+    let destructor = Box::leak(Box::new(destructor));
+    let callbacks = BNObjectDestructionCallbacks {
+        context: destructor as *mut _ as *mut c_void,
+        destructBinaryView: Some(cb_destruct_binary_view::<D>),
+        destructFileMetadata: Some(cb_destruct_file_metadata::<D>),
+        destructFunction: Some(cb_destruct_function::<D>),
+    };
+    let mut handle = ObjectDestructorHandle {
+        callbacks,
+        _life: std::marker::PhantomData,
+    };
+    unsafe { BNRegisterObjectDestructionCallbacks(&mut handle.callbacks) };
+    handle
+}
+
+/// The handle for the [`ObjectDestructor`].
+///
+/// Once this handle is dropped, the destructor will be unregistered and the associated resources will be cleaned up.
+pub struct ObjectDestructorHandle<'a, D: ObjectDestructor> {
+    callbacks: BNObjectDestructionCallbacks,
+    _life: std::marker::PhantomData<&'a D>,
+}
+
+impl<D: ObjectDestructor> Drop for ObjectDestructorHandle<'_, D> {
+    fn drop(&mut self) {
+        unsafe { BNUnregisterObjectDestructionCallbacks(&mut self.callbacks) };
+        let _ = unsafe { Box::from_raw(self.callbacks.context as *mut D) };
+    }
+}
+
+/// The trait required for receiving core object destruction callbacks.
+///
+/// This is useful for cleaning up resources which are associated with a given core object.
+pub trait ObjectDestructor: 'static + Sync + Sized {
+    /// Called when a [`BinaryView`] is about to be destroyed.
+    fn destruct_view(&self, _view: &BinaryView) {}
+
+    /// Called when a [`FileMetadata`] is about to be destroyed.
+    fn destruct_file_metadata(&self, _metadata: &FileMetadata) {}
+
+    /// Called when a [`Function`] is about to be destroyed.
+    fn destruct_function(&self, _func: &Function) {}
+}
+
+unsafe extern "C" fn cb_destruct_binary_view<D: ObjectDestructor>(
+    ctxt: *mut c_void,
+    view: *mut BNBinaryView,
+) {
+    ffi_wrap!("ObjectDestructor::destruct_view", {
+        let destructor = &*(ctxt as *mut D);
+        let view = BinaryView { handle: view };
+        destructor.destruct_view(&view);
+    })
+}
+
+unsafe extern "C" fn cb_destruct_file_metadata<D: ObjectDestructor>(
+    ctxt: *mut c_void,
+    file: *mut BNFileMetadata,
+) {
+    ffi_wrap!("ObjectDestructor::destruct_file_metadata", {
+        let destructor = &*(ctxt as *mut D);
+        let file = FileMetadata::from_raw(file);
+        destructor.destruct_file_metadata(&file);
+    })
+}
+
+unsafe extern "C" fn cb_destruct_function<D: ObjectDestructor>(
+    ctxt: *mut c_void,
+    func: *mut BNFunction,
+) {
+    ffi_wrap!("ObjectDestructor::destruct_function", {
+        let destructor = &*(ctxt as *mut D);
+        let func = Function { handle: func };
+        destructor.destruct_function(&func);
+    })
+}