fix scriptingprovider inconsistencies

Author: rbran

rust/src/scriptingprovider.rs

@@ -1,5 +1,7 @@
 use core::{ffi, mem, ptr};
+use std::mem::MaybeUninit;
 use std::pin::Pin;
+use std::sync::Arc;
 
 use binaryninjacore_sys::*;
 
@@ -18,6 +20,9 @@ pub struct ScriptingProvider {
     handle: ptr::NonNull<BNScriptingProvider>,
 }
 
+unsafe impl Sync for ScriptingProvider {}
+unsafe impl Send for ScriptingProvider {}
+
 impl ScriptingProvider {
     pub(crate) unsafe fn from_raw(handle: ptr::NonNull<BNScriptingProvider>) -> Self {
         Self { handle }
@@ -95,30 +100,9 @@ impl ScriptingProvider {
         }
     }
 
-    pub fn new_instance<S: ScriptingInstanceCallbacks>(&self) -> ScriptingInstance {
-        // SAFETY freed by cb_destroy_instance
-        let uninit = Box::leak(Box::new(mem::MaybeUninit::zeroed()));
-        let mut callbacks = BNScriptingInstanceCallbacks {
-            context: unsafe { uninit.assume_init_mut() as *mut S as *mut ffi::c_void },
-            destroyInstance: Some(cb_destroy_instance::<S>),
-            externalRefTaken: Some(cb_external_ref_taken::<S>),
-            externalRefReleased: Some(cb_external_ref_released::<S>),
-            executeScriptInput: Some(cb_execute_script_input::<S>),
-            executeScriptInputFromFilename: Some(cb_execute_script_input_from_filename::<S>),
-            cancelScriptInput: Some(cb_cancel_script_input::<S>),
-            releaseBinaryView: Some(cb_release_binary_view::<S>),
-            setCurrentBinaryView: Some(cb_set_current_binary_view::<S>),
-            setCurrentFunction: Some(cb_set_current_function::<S>),
-            setCurrentBasicBlock: Some(cb_set_current_basic_block::<S>),
-            setCurrentAddress: Some(cb_set_current_address::<S>),
-            setCurrentSelection: Some(cb_set_current_selection::<S>),
-            completeInput: Some(cb_complete_input::<S>),
-            stop: Some(cb_stop::<S>),
-        };
-        let result = unsafe { BNInitScriptingInstance(self.as_raw(), &mut callbacks) };
-        let instance = unsafe { ScriptingInstance::from_raw(ptr::NonNull::new(result).unwrap()) };
-        uninit.write(S::new(*self, &instance));
-        instance
+    pub fn new_instance(&self) -> Ref<ScriptingInstance> {
+        let instance = unsafe { BNCreateScriptingProviderInstance(self.as_raw()) };
+        unsafe { ScriptingInstance::ref_from_raw(ptr::NonNull::new(instance).unwrap()) }
     }
 }
 
@@ -143,6 +127,9 @@ pub struct ScriptingInstance {
     handle: ptr::NonNull<BNScriptingInstance>,
 }
 
+unsafe impl Sync for ScriptingInstance {}
+unsafe impl Send for ScriptingInstance {}
+
 impl ToOwned for ScriptingInstance {
     type Owned = Ref<Self>;
 
@@ -163,18 +150,10 @@ unsafe impl RefCountable for ScriptingInstance {
 }
 
 impl ScriptingInstance {
-    pub(crate) unsafe fn from_raw(handle: ptr::NonNull<BNScriptingInstance>) -> Self {
-        Self { handle }
-    }
-
     pub(crate) unsafe fn ref_from_raw(handle: ptr::NonNull<BNScriptingInstance>) -> Ref<Self> {
         Ref::new(Self { handle })
     }
 
-    pub(crate) unsafe fn into_raw(self) -> *mut BNScriptingInstance {
-        mem::ManuallyDrop::new(self).handle.as_ptr()
-    }
-
     #[allow(clippy::mut_from_ref)]
     pub(crate) unsafe fn as_raw(&self) -> &mut BNScriptingInstance {
         &mut *self.handle.as_ptr()
@@ -214,10 +193,10 @@ impl ScriptingInstance {
         unsafe { BNNotifyInputReadyStateForScriptingInstance(self.as_raw(), state) }
     }
 
-    pub fn register_output_listener<L: ScriptingOutputListener>(
-        self,
+    pub fn register_output_listener<'a, L: ScriptingOutputListener>(
+        &'a self,
         listener: L,
-    ) -> ScriptingInstanceWithListener<L> {
+    ) -> ScriptingInstanceWithListener<'a, L> {
         let mut listener = Box::pin(listener);
         let mut callbacks = BNScriptingOutputListener {
             context: unsafe { listener.as_mut().get_unchecked_mut() } as *mut _ as *mut ffi::c_void,
@@ -323,26 +302,26 @@ impl ScriptingInstance {
     }
 }
 
-pub struct ScriptingInstanceWithListener<L: ScriptingOutputListener> {
-    handle: ScriptingInstance,
+pub struct ScriptingInstanceWithListener<'a, L: ScriptingOutputListener> {
+    handle: &'a ScriptingInstance,
     listener: Pin<Box<L>>,
 }
 
-impl<L: ScriptingOutputListener> AsRef<ScriptingInstance> for ScriptingInstanceWithListener<L> {
-    fn as_ref(&self) -> &ScriptingInstance {
-        &self.handle
+impl<L: ScriptingOutputListener> AsRef<L> for ScriptingInstanceWithListener<'_, L> {
+    fn as_ref(&self) -> &L {
+        &self.listener
     }
 }
 
-impl<L: ScriptingOutputListener> core::ops::Deref for ScriptingInstanceWithListener<L> {
-    type Target = ScriptingInstance;
+impl<L: ScriptingOutputListener> core::ops::Deref for ScriptingInstanceWithListener<'_, L> {
+    type Target = L;
     fn deref(&self) -> &Self::Target {
-        &self.handle
+        &self.listener
     }
 }
 
-impl<L: ScriptingOutputListener> ScriptingInstanceWithListener<L> {
-    pub fn unregister(mut self) -> ScriptingInstance {
+impl<L: ScriptingOutputListener> ScriptingInstanceWithListener<'_, L> {
+    pub fn unregister(mut self) {
         let mut callbacks = BNScriptingOutputListener {
             context: unsafe { self.listener.as_mut().get_unchecked_mut() } as *mut _
                 as *mut ffi::c_void,
@@ -354,28 +333,29 @@ impl<L: ScriptingOutputListener> ScriptingInstanceWithListener<L> {
         unsafe {
             BNUnregisterScriptingInstanceOutputListener(self.handle.as_raw(), &mut callbacks)
         };
-        // drop the listener
-        let Self {
-            handle: instance,
-            listener: _,
-        } = self;
-        // return the inner instance
-        instance
     }
 }
 
 pub trait ScriptingCustomProvider: Sync + Send {
+    type Instance: ScriptingInstanceCallbacks;
     fn new(core: ScriptingProvider) -> Self;
-    fn create_instance(&self) -> ScriptingInstance;
+    fn init_instance(&self, handle: CoreScriptingInstance) -> Self::Instance;
+    fn get_core(&self) -> &ScriptingProvider;
     fn load_module(&self, repo_path: &str, plugin_path: &str, force: bool) -> bool;
     fn install_modules(&self, modules: &str) -> bool;
+
+    fn new_instance(&self) -> (Arc<Self::Instance>, Ref<ScriptingInstance>)
+    where
+        Self: Sized,
+    {
+        register_instance::<Self>(self)
+    }
 }
 
+pub struct CoreScriptingInstance(Ref<ScriptingInstance>);
+
 pub trait ScriptingInstanceCallbacks: Sync + Send {
-    fn new(provider: ScriptingProvider, handle: &ScriptingInstance) -> Self;
-    fn destroy_instance(&self);
-    fn external_ref_taken(&self);
-    fn external_ref_released(&self);
+    fn get_core(&self) -> &CoreScriptingInstance;
     fn execute_script_input(&self, input: &str) -> ScriptingProviderExecuteResult;
     fn execute_script_input_from_filename(&self, input: &str) -> ScriptingProviderExecuteResult;
     fn cancel_script_input(&self);
@@ -387,6 +367,22 @@ pub trait ScriptingInstanceCallbacks: Sync + Send {
     fn set_current_selection(&self, begin: u64, end: u64);
     fn complete_input(&self, text: &str, state: u64) -> String;
     fn stop(&self);
+
+    fn notify_output(&self, text: &str) {
+        self.get_core().0.notify_output(text);
+    }
+    fn notify_warning(&self, text: &str) {
+        self.get_core().0.notify_warning(text);
+    }
+    fn notify_error(&self, text: &str) {
+        self.get_core().0.notify_error(text);
+    }
+    fn register_output_listener<'a, L: ScriptingOutputListener>(
+        &'a self,
+        listener: L,
+    ) -> ScriptingInstanceWithListener<'a, L> {
+        self.get_core().0.register_output_listener(listener)
+    }
 }
 
 pub trait ScriptingOutputListener: Sync + Send {
@@ -404,7 +400,6 @@ where
 {
     let name = name.into_bytes_with_nul();
     let api_name = api_name.into_bytes_with_nul();
-    // SAFETY: Websocket provider is never freed
     let provider_uinit = Box::leak(Box::new(mem::MaybeUninit::zeroed()));
     let result = unsafe {
         BNRegisterScriptingProvider(
@@ -423,11 +418,59 @@ where
     (unsafe { provider_uinit.assume_init_ref() }, provider_core)
 }
 
+fn register_instance<S: ScriptingCustomProvider>(
+    provider: &S,
+) -> (Arc<S::Instance>, Ref<ScriptingInstance>) {
+    let instance_uninit: *const MaybeUninit<S::Instance> = Arc::into_raw(Arc::new_uninit());
+    let mut callbacks = BNScriptingInstanceCallbacks {
+        context: instance_uninit as *mut ffi::c_void,
+        destroyInstance: Some(cb_destroy_instance::<S::Instance>),
+        externalRefTaken: Some(cb_external_ref_taken::<S::Instance>),
+        externalRefReleased: Some(cb_external_ref_released::<S::Instance>),
+        executeScriptInput: Some(cb_execute_script_input::<S::Instance>),
+        executeScriptInputFromFilename: Some(cb_execute_script_input_from_filename::<S::Instance>),
+        cancelScriptInput: Some(cb_cancel_script_input::<S::Instance>),
+        releaseBinaryView: Some(cb_release_binary_view::<S::Instance>),
+        setCurrentBinaryView: Some(cb_set_current_binary_view::<S::Instance>),
+        setCurrentFunction: Some(cb_set_current_function::<S::Instance>),
+        setCurrentBasicBlock: Some(cb_set_current_basic_block::<S::Instance>),
+        setCurrentAddress: Some(cb_set_current_address::<S::Instance>),
+        setCurrentSelection: Some(cb_set_current_selection::<S::Instance>),
+        completeInput: Some(cb_complete_input::<S::Instance>),
+        stop: Some(cb_stop::<S::Instance>),
+    };
+    let handle = unsafe { BNInitScriptingInstance(provider.get_core().as_raw(), &mut callbacks) };
+    let core_instance = unsafe {
+        CoreScriptingInstance(ScriptingInstance::ref_from_raw(
+            ptr::NonNull::new(handle).unwrap(),
+        ))
+    };
+    let new_instance = provider.init_instance(core_instance);
+
+    // recreate the Arc, initialize it, clone it, then leak it again
+    let instance: Arc<_> = unsafe {
+        let mut instance_uninit = Arc::from_raw(instance_uninit);
+        Arc::get_mut(&mut instance_uninit)
+            .unwrap()
+            .write(new_instance);
+        let instance = Arc::<MaybeUninit<_>>::assume_init(instance_uninit);
+        let result = Arc::clone(&instance);
+        // leak, this will be owned by the [ScriptingInstance]
+        let _ptr = Arc::into_raw(instance);
+        result
+    };
+
+    let core_instance =
+        unsafe { ScriptingInstance::ref_from_raw(ptr::NonNull::new(handle).unwrap()) };
+    (instance, core_instance)
+}
+
 unsafe extern "C" fn cb_create_instance<S: ScriptingCustomProvider>(
     ctxt: *mut ffi::c_void,
 ) -> *mut BNScriptingInstance {
     let ctxt = &mut *(ctxt as *mut S);
-    ctxt.create_instance().into_raw()
+    let (_rust_instance, core_instance) = register_instance(ctxt);
+    core_instance.as_raw()
 }
 
 unsafe extern "C" fn cb_load_module<S: ScriptingCustomProvider>(
@@ -456,19 +499,17 @@ unsafe extern "C" fn cb_install_modules<S: ScriptingCustomProvider>(
 }
 
 unsafe extern "C" fn cb_destroy_instance<S: ScriptingInstanceCallbacks>(ctxt: *mut ffi::c_void) {
-    drop(Box::from_raw(ctxt as *mut S))
+    Arc::from_raw(ctxt as *mut S);
 }
 
 unsafe extern "C" fn cb_external_ref_taken<S: ScriptingInstanceCallbacks>(ctxt: *mut ffi::c_void) {
-    let ctxt = &mut *(ctxt as *mut S);
-    ctxt.external_ref_taken()
+    Arc::increment_strong_count(ctxt as *mut S);
 }
 
 unsafe extern "C" fn cb_external_ref_released<S: ScriptingInstanceCallbacks>(
     ctxt: *mut ffi::c_void,
 ) {
-    let ctxt = &mut *(ctxt as *mut S);
-    ctxt.external_ref_released()
+    Arc::decrement_strong_count(ctxt as *mut S);
 }
 
 unsafe extern "C" fn cb_execute_script_input<S: ScriptingInstanceCallbacks>(
@@ -477,7 +518,8 @@ unsafe extern "C" fn cb_execute_script_input<S: ScriptingInstanceCallbacks>(
 ) -> BNScriptingProviderExecuteResult {
     let input = ffi::CStr::from_ptr(input);
     let ctxt = &mut *(ctxt as *mut S);
-    ctxt.execute_script_input(&input.to_string_lossy())
+    let result = ctxt.execute_script_input(&input.to_string_lossy());
+    result
 }
 
 unsafe extern "C" fn cb_execute_script_input_from_filename<S: ScriptingInstanceCallbacks>(