Models CoreDataRenderer after CoreTypePrinter to prevent crashes

The previous approach of reference counting did lead to crashes as the renderer was freed after registration :/

Author: lukbukkit

rust/src/data_renderer.rs

@@ -1,26 +1,31 @@
 use core::ffi;
-
+use ffi::c_void;
+use std::ptr::NonNull;
+use log::debug;
 use binaryninjacore_sys::*;
 
 use crate::binary_view::BinaryView;
 use crate::disassembly::{DisassemblyTextLine, InstructionTextToken};
-use crate::rc::{Ref, RefCountable};
 use crate::types::Type;
 
 // NOTE the type_ inside the context can be both owned or borrowed, because
 // this type only exist as a reference and is never created by itself (AKA
 // don't have a *from_raw function, it don't need to worry about drop it.
 #[repr(transparent)]
-pub struct TypeContext(BNTypeContext);
+pub struct TypeContext {
+    handle: BNTypeContext
+}
 
 impl TypeContext {
     pub fn type_(&self) -> &Type {
+        // debug!("TypeContext type_");
         // SAFETY Type and `*mut BNType` are transparent
-        unsafe { core::mem::transmute::<&*mut BNType, &Type>(&self.0.type_) }
+        unsafe { core::mem::transmute::<&*mut BNType, &Type>(&self.handle.type_) }
     }
 
     pub fn offset(&self) -> usize {
-        self.0.offset
+        // debug!("TypeContext offset");
+        self.handle.offset
     }
 }
 
@@ -46,18 +51,20 @@ pub trait CustomDataRenderer: Sized + Sync + Send + 'static {
 }
 
 trait CustomDataRendererFFI: CustomDataRenderer {
-    unsafe extern "C" fn free_object_ffi(ctxt: *mut ffi::c_void) {
+    unsafe extern "C" fn free_object_ffi(ctxt: *mut c_void) {
+        // debug!("free_object_ffi");
         drop(Box::from_raw(ctxt as *mut Self))
     }
 
     unsafe extern "C" fn is_valid_for_data_ffi(
-        ctxt: *mut ffi::c_void,
+        ctxt: *mut c_void,
         view: *mut BNBinaryView,
         addr: u64,
         type_: *mut BNType,
         type_ctx: *mut BNTypeContext,
         ctx_count: usize,
     ) -> bool {
+        // debug!("is_valid_for_data_ffi");
         let ctxt = ctxt as *mut Self;
         // SAFETY BNTypeContext and TypeContext are transparent
         let types = core::slice::from_raw_parts(type_ctx as *mut TypeContext, ctx_count);
@@ -70,7 +77,7 @@ trait CustomDataRendererFFI: CustomDataRenderer {
     }
 
     unsafe extern "C" fn get_lines_for_data_ffi(
-        ctxt: *mut ffi::c_void,
+        ctxt: *mut c_void,
         view: *mut BNBinaryView,
         addr: u64,
         type_: *mut BNType,
@@ -82,6 +89,7 @@ trait CustomDataRendererFFI: CustomDataRenderer {
         ctx_count: usize,
         language: *const ffi::c_char,
     ) -> *mut BNDisassemblyTextLine {
+        // debug!("get_lines_for_data_ffi");
         let ctxt = ctxt as *mut Self;
         // SAFETY BNTypeContext and TypeContext are transparent
         let types = core::slice::from_raw_parts(type_ctx as *mut TypeContext, ctx_count);
@@ -104,10 +112,11 @@ trait CustomDataRendererFFI: CustomDataRenderer {
     }
 
     unsafe extern "C" fn free_lines_ffi(
-        _ctx: *mut ffi::c_void,
+        _ctx: *mut c_void,
         lines: *mut BNDisassemblyTextLine,
         count: usize,
     ) {
+        // debug!("free_lines_ffi");
         let lines = Box::from_raw(core::slice::from_raw_parts_mut(lines, count));
         drop(
             lines
@@ -120,78 +129,55 @@ trait CustomDataRendererFFI: CustomDataRenderer {
 
 impl<C: CustomDataRenderer> CustomDataRendererFFI for C {}
 
-pub struct CoreDataRenderer(*mut BNDataRenderer);
-
-impl CoreDataRenderer {
-    pub(crate) unsafe fn ref_from_raw(raw: *mut BNDataRenderer) -> Ref<Self> {
-        Ref::new(Self(raw))
-    }
-    pub(crate) fn as_raw(&self) -> *mut BNDataRenderer {
-        self.0
-    }
+pub struct CoreDataRenderer {
+    pub(crate) handle: NonNull<BNDataRenderer>,
 }
 
-unsafe impl RefCountable for CoreDataRenderer {
-    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
-        Self::ref_from_raw(BNNewDataRendererReference(handle.0))
-    }
-
-    unsafe fn dec_ref(handle: &Self) {
-        BNFreeDataRenderer(handle.0);
-    }
-}
-
-impl ToOwned for CoreDataRenderer {
-    type Owned = Ref<Self>;
-
-    fn to_owned(&self) -> Self::Owned {
-        unsafe { <Self as RefCountable>::inc_ref(self) }
+impl CoreDataRenderer {
+    pub(crate) unsafe fn from_raw(handle: NonNull<BNDataRenderer>) -> CoreDataRenderer {
+        Self { handle }
     }
 }
 
-fn create_custom_data_renderer<C: CustomDataRenderer>(custom: C) -> Ref<CoreDataRenderer> {
-    let custom = Box::leak(Box::new(custom));
+fn create_custom_data_renderer<T: CustomDataRenderer>(renderer: T) -> (&'static mut T, CoreDataRenderer) {
+    let renderer = Box::leak(Box::new(renderer));
     let mut callbacks = BNCustomDataRenderer {
-        context: custom as *mut C as *mut ffi::c_void,
-        freeObject: Some(<C as CustomDataRendererFFI>::free_object_ffi),
-        isValidForData: Some(<C as CustomDataRendererFFI>::is_valid_for_data_ffi),
-        getLinesForData: Some(<C as CustomDataRendererFFI>::get_lines_for_data_ffi),
-        freeLines: Some(<C as CustomDataRendererFFI>::free_lines_ffi),
+        context: renderer as *mut _ as *mut c_void,
+        freeObject: Some(<T as CustomDataRendererFFI>::free_object_ffi),
+        isValidForData: Some(<T as CustomDataRendererFFI>::is_valid_for_data_ffi),
+        getLinesForData: Some(<T as CustomDataRendererFFI>::get_lines_for_data_ffi),
+        freeLines: Some(<T as CustomDataRendererFFI>::free_lines_ffi),
     };
-    unsafe { CoreDataRenderer::ref_from_raw(BNCreateDataRenderer(&mut callbacks)) }
+    let result = unsafe { BNCreateDataRenderer(&mut callbacks) };
+    let core = unsafe { CoreDataRenderer::from_raw(NonNull::new(result).unwrap()) };
+    (renderer, core)
 }
 
-pub fn register_generic_data_renderer<C: CustomDataRenderer>(custom: C) -> Ref<CoreDataRenderer> {
-    let renderer = create_custom_data_renderer(custom);
-    register_generic_renderer(&renderer);
-    renderer
+pub fn register_generic_data_renderer<T: CustomDataRenderer>(custom: T) -> (&'static mut T, CoreDataRenderer) {
+    let (renderer, core) = create_custom_data_renderer(custom);
+    // debug!("register_generic_data_renderer: core={:?}", core.handle);
+    let container = DataRendererContainer::get();
+    unsafe { BNRegisterGenericDataRenderer(container.handle, core.handle.as_ptr()) }
+    (renderer, core)
 }
 
-pub fn register_specific_data_renderer<C: CustomDataRenderer>(custom: C) -> Ref<CoreDataRenderer> {
-    let renderer = create_custom_data_renderer(custom);
-    register_specific_renderer(&renderer);
-    renderer
+pub fn register_specific_data_renderer<C: CustomDataRenderer>(custom: C) -> (&'static mut C, CoreDataRenderer) {
+    let (renderer, core) = create_custom_data_renderer(custom);
+    // debug!("register_specific_data_renderer: core={:?}", core.handle);
+    let container = DataRendererContainer::get();
+    unsafe { BNRegisterTypeSpecificDataRenderer(container.handle, core.handle.as_ptr()) }
+    (renderer, core)
 }
 
 #[derive(Clone, Copy)]
-struct DataRendererContainer(*mut BNDataRendererContainer);
+struct DataRendererContainer {
+    pub(crate) handle: *mut BNDataRendererContainer
+}
 
 impl DataRendererContainer {
-    pub(crate) fn as_raw(&self) -> *mut BNDataRendererContainer {
-        self.0
-    }
-
     pub fn get() -> Self {
-        Self(unsafe { BNGetDataRendererContainer() })
+        Self {
+            handle: unsafe { BNGetDataRendererContainer() }
+        }
     }
-}
-
-fn register_generic_renderer(renderer: &CoreDataRenderer) {
-    let container = DataRendererContainer::get();
-    unsafe { BNRegisterGenericDataRenderer(container.as_raw(), renderer.as_raw()) }
-}
-
-fn register_specific_renderer(renderer: &CoreDataRenderer) {
-    let container = DataRendererContainer::get();
-    unsafe { BNRegisterTypeSpecificDataRenderer(container.as_raw(), renderer.as_raw()) }
-}
+}