[Core] Add optional callback for when the type archive closes for BNTypeArchiveNotification

Update the rust bindings accordingly

Author: emesare

binaryninjacore.h

@@ -3561,6 +3561,7 @@ extern "C"
 		void (*typeUpdated)(void* ctxt, BNTypeArchive* archive, const char* id, BNType* oldDefinition, BNType* newDefinition);
 		void (*typeRenamed)(void* ctxt, BNTypeArchive* archive, const char* id, const BNQualifiedName* oldName, const BNQualifiedName* newName);
 		void (*typeDeleted)(void* ctxt, BNTypeArchive* archive, const char* id, BNType* definition);
+		void (*closed)(void* ctxt, BNTypeArchive* archive);
 	} BNTypeArchiveNotification;
 
 	BN_ENUM(uint8_t, BNTypeContainerType)

rust/src/types/archive.rs

@@ -616,36 +616,40 @@ impl TypeArchive {
     }
 
     /// Register a notification listener
-    pub fn register_notification_callback<T: TypeArchiveNotificationCallback>(
+    #[must_use]
+    pub fn register_notification_callback<'a, T: TypeArchiveNotificationCallback + 'a>(
         &self,
         callback: T,
-    ) -> TypeArchiveCallbackHandle<T> {
+    ) -> TypeArchiveCallbackHandle<'a, T> {
         // SAFETY free on [TypeArchiveCallbackHandle::Drop]
         let callback = Box::leak(Box::new(callback));
-        let mut notification = BNTypeArchiveNotification {
+        let notification = BNTypeArchiveNotification {
             context: callback as *mut T as *mut c_void,
             typeAdded: Some(cb_type_added::<T>),
             typeUpdated: Some(cb_type_updated::<T>),
             typeRenamed: Some(cb_type_renamed::<T>),
             typeDeleted: Some(cb_type_deleted::<T>),
+            closed: Some(cb_type_archive_closed::<T>),
         };
-        unsafe { BNRegisterTypeArchiveNotification(self.handle.as_ptr(), &mut notification) }
+        let mut boxed_notification = Box::new(notification);
+        unsafe {
+            BNRegisterTypeArchiveNotification(self.handle.as_ptr(), boxed_notification.as_mut())
+        }
         TypeArchiveCallbackHandle {
-            callback,
             type_archive: self.to_owned(),
+            handle: boxed_notification,
+            _lifetime: Default::default(),
         }
     }
 
-    // NOTE NotificationClosure is left private, there is no need for the user
-    // to know or use it.
-    #[allow(private_interfaces)]
-    pub fn register_notification_closure<A, U, R, D>(
+    #[must_use]
+    pub fn register_notification_closure<'a, A, U, R, D>(
         &self,
         type_added: A,
         type_updated: U,
         type_renamed: R,
         type_deleted: D,
-    ) -> TypeArchiveCallbackHandle<NotificationClosure<A, U, R, D>>
+    ) -> TypeArchiveCallbackHandle<'a, NotificationClosure<A, U, R, D>>
     where
         A: FnMut(&TypeArchive, &str, &Type),
         U: FnMut(&TypeArchive, &str, &Type, &Type),
@@ -661,7 +665,9 @@ impl TypeArchive {
     }
 
     /// Close a type archive, disconnecting it from any active views and closing
-    /// any open file handles
+    /// any open file handles.
+    ///
+    /// This will signal all registered notification callbacks with [`TypeArchiveNotificationCallback::closed`].
     pub fn close(&self) {
         unsafe { BNCloseTypeArchive(self.handle.as_ptr()) }
     }
@@ -874,29 +880,35 @@ unsafe impl CoreArrayProviderInner for TypeArchive {
     }
 }
 
-pub struct TypeArchiveCallbackHandle<T: TypeArchiveNotificationCallback> {
-    callback: *mut T,
+pub struct TypeArchiveCallbackHandle<'a, T: TypeArchiveNotificationCallback>
+where
+    T: 'a,
+{
     type_archive: Ref<TypeArchive>,
+    handle: Box<BNTypeArchiveNotification>,
+    _lifetime: std::marker::PhantomData<&'a T>,
 }
 
-impl<T: TypeArchiveNotificationCallback> Drop for TypeArchiveCallbackHandle<T> {
+impl<T: TypeArchiveNotificationCallback> TypeArchiveCallbackHandle<'_, T> {
+    unsafe fn unregister_for_archive(&mut self) {
+        BNUnregisterTypeArchiveNotification(self.type_archive.handle.as_ptr(), self.handle.as_mut())
+    }
+
+    unsafe fn extract_context(&mut self) -> Box<T> {
+        Box::from_raw(self.handle.context as *mut T)
+    }
+
+    pub fn unregister(mut self) -> T {
+        unsafe { self.unregister_for_archive() };
+        unsafe { *self.extract_context() }
+    }
+}
+
+impl<T: TypeArchiveNotificationCallback> Drop for TypeArchiveCallbackHandle<'_, T> {
     fn drop(&mut self) {
-        let mut notification = BNTypeArchiveNotification {
-            context: self.callback as *mut c_void,
-            typeAdded: Some(cb_type_added::<T>),
-            typeUpdated: Some(cb_type_updated::<T>),
-            typeRenamed: Some(cb_type_renamed::<T>),
-            typeDeleted: Some(cb_type_deleted::<T>),
-        };
-        // unregister the notification callback
-        unsafe {
-            BNUnregisterTypeArchiveNotification(
-                self.type_archive.handle.as_ptr(),
-                &mut notification,
-            )
-        }
-        // free the context created at [TypeArchive::register_notification_callback]
-        drop(unsafe { Box::from_raw(self.callback) });
+        unsafe { self.unregister_for_archive() };
+        // drop context, avoid memory leak
+        let _ctxt = unsafe { self.extract_context() };
     }
 }
 
@@ -944,9 +956,14 @@ pub trait TypeArchiveNotificationCallback {
     /// * `id` - Id of type deleted
     /// * `definition` - Definition of type deleted
     fn type_deleted(&mut self, _archive: &TypeArchive, _id: &str, _definition: &Type) {}
+
+    /// Called when the associated type archive is closed.
+    ///
+    /// Useful to free up data in this context associated with the type archive.
+    fn closed(&mut self, _archive: &TypeArchive) {}
 }
 
-struct NotificationClosure<A, U, R, D>
+pub struct NotificationClosure<A, U, R, D>
 where
     A: FnMut(&TypeArchive, &str, &Type),
     U: FnMut(&TypeArchive, &str, &Type, &Type),
@@ -1068,6 +1085,16 @@ unsafe extern "C" fn cb_type_deleted<T: TypeArchiveNotificationCallback>(
     )
 }
 
+unsafe extern "C" fn cb_type_archive_closed<T: TypeArchiveNotificationCallback>(
+    ctxt: *mut ::std::os::raw::c_void,
+    archive: *mut BNTypeArchive,
+) {
+    let ctxt: &mut T = &mut *(ctxt as *mut T);
+    // `archive` is owned by the caller.
+    let archive = unsafe { TypeArchive::from_raw(NonNull::new(archive).unwrap()) };
+    ctxt.closed(&archive)
+}
+
 #[repr(transparent)]
 pub struct TypeArchiveMergeConflict {
     handle: NonNull<BNTypeArchiveMergeConflict>,