From e28e89f107d11e0add00b22247377a14a1767bbd Mon Sep 17 00:00:00 2001 From: Rubens Brandao Date: Thu, 27 Jun 2024 09:54:12 -0300 Subject: [PATCH 1/4] implement rust Scripting --- rust/src/lib.rs | 1 + rust/src/scriptingprovider.rs | 589 ++++++++++++++++++++++++++++++++++ 2 files changed, 590 insertions(+) create mode 100644 rust/src/scriptingprovider.rs diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 0cb0484e6..ca44c217e 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -69,6 +69,7 @@ pub mod rc; pub mod references; pub mod relocation; pub mod render_layer; +pub mod scriptingprovider; pub mod section; pub mod segment; pub mod settings; diff --git a/rust/src/scriptingprovider.rs b/rust/src/scriptingprovider.rs new file mode 100644 index 000000000..d02b089a1 --- /dev/null +++ b/rust/src/scriptingprovider.rs @@ -0,0 +1,589 @@ +use core::{ffi, mem, ptr}; +use std::pin::Pin; + +use binaryninjacore_sys::*; + +use crate::basic_block::BasicBlock; +use crate::binary_view::BinaryView; +use crate::function::{Function, NativeBlock}; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner}; +use crate::string::{BnStrCompatible, BnString}; + +pub type ScriptingProviderExecuteResult = BNScriptingProviderExecuteResult; +pub type ScriptingProviderInputReadyState = BNScriptingProviderInputReadyState; + +#[derive(Clone, Copy)] +#[repr(transparent)] +pub struct ScriptingProvider { + handle: ptr::NonNull, +} + +impl ScriptingProvider { + pub(crate) unsafe fn from_raw(handle: ptr::NonNull) -> Self { + Self { handle } + } + + pub(crate) unsafe fn ref_from_raw(handle: &*mut BNScriptingProvider) -> &Self { + assert!(!handle.is_null()); + mem::transmute(handle) + } + + #[allow(clippy::mut_from_ref)] + pub(crate) unsafe fn as_raw(&self) -> &mut BNScriptingProvider { + &mut *self.handle.as_ptr() + } + + pub fn all() -> Array { + let mut count = 0; + let result = unsafe { BNGetScriptingProviderList(&mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + pub fn by_name(name: S) -> Option { + let name = name.into_bytes_with_nul(); + let result = + unsafe { BNGetScriptingProviderByName(name.as_ref().as_ptr() as *const ffi::c_char) }; + ptr::NonNull::new(result).map(|h| unsafe { Self::from_raw(h) }) + } + + pub fn by_api_name(name: S) -> Option { + let name = name.into_bytes_with_nul(); + let result = unsafe { + BNGetScriptingProviderByAPIName(name.as_ref().as_ptr() as *const ffi::c_char) + }; + ptr::NonNull::new(result).map(|h| unsafe { Self::from_raw(h) }) + } + + pub fn name(&self) -> BnString { + let result = unsafe { BNGetScriptingProviderName(self.as_raw()) }; + assert!(!result.is_null()); + unsafe { BnString::from_raw(result) } + } + + pub fn api_name(&self) -> BnString { + let result = unsafe { BNGetScriptingProviderAPIName(self.as_raw()) }; + assert!(!result.is_null()); + unsafe { BnString::from_raw(result) } + } + + pub fn load_module( + &self, + repository: R, + module: M, + force: bool, + ) -> bool { + let repository = repository.into_bytes_with_nul(); + let module = module.into_bytes_with_nul(); + unsafe { + BNLoadScriptingProviderModule( + self.as_raw(), + repository.as_ref().as_ptr() as *const ffi::c_char, + module.as_ref().as_ptr() as *const ffi::c_char, + force, + ) + } + } + + pub fn install_modules(&self, modules: M) -> bool { + let modules = modules.into_bytes_with_nul(); + unsafe { + BNInstallScriptingProviderModules( + self.as_raw(), + modules.as_ref().as_ptr() as *const ffi::c_char, + ) + } + } + + pub fn new_instance(&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::), + externalRefTaken: Some(cb_external_ref_taken::), + externalRefReleased: Some(cb_external_ref_released::), + executeScriptInput: Some(cb_execute_script_input::), + executeScriptInputFromFilename: Some(cb_execute_script_input_from_filename::), + cancelScriptInput: Some(cb_cancel_script_input::), + releaseBinaryView: Some(cb_release_binary_view::), + setCurrentBinaryView: Some(cb_set_current_binary_view::), + setCurrentFunction: Some(cb_set_current_function::), + setCurrentBasicBlock: Some(cb_set_current_basic_block::), + setCurrentAddress: Some(cb_set_current_address::), + setCurrentSelection: Some(cb_set_current_selection::), + completeInput: Some(cb_complete_input::), + stop: Some(cb_stop::), + }; + 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 + } +} + +impl CoreArrayProvider for ScriptingProvider { + type Raw = *mut BNScriptingProvider; + type Context = (); + type Wrapped<'a> = &'a Self; +} + +unsafe impl CoreArrayProviderInner for ScriptingProvider { + unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) { + BNFreeScriptingProviderList(raw) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> { + Self::ref_from_raw(raw) + } +} + +#[repr(transparent)] +pub struct ScriptingInstance { + handle: ptr::NonNull, +} + +impl Clone for ScriptingInstance { + fn clone(&self) -> Self { + let result = unsafe { BNNewScriptingInstanceReference(self.as_raw()) }; + unsafe { Self::from_raw(ptr::NonNull::new(result).unwrap()) } + } +} + +impl Drop for ScriptingInstance { + fn drop(&mut self) { + unsafe { BNFreeScriptingInstance(self.as_raw()) } + } +} + +impl ScriptingInstance { + pub(crate) unsafe fn from_raw(handle: ptr::NonNull) -> Self { + 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() + } + + pub fn notify_output(&self, text: S) { + let text = text.into_bytes_with_nul(); + unsafe { + BNNotifyOutputForScriptingInstance( + self.as_raw(), + text.as_ref().as_ptr() as *const ffi::c_char, + ) + } + } + + pub fn notify_warning(&self, text: S) { + let text = text.into_bytes_with_nul(); + unsafe { + BNNotifyOutputForScriptingInstance( + self.as_raw(), + text.as_ref().as_ptr() as *const ffi::c_char, + ) + } + } + + pub fn notify_error(&self, text: S) { + let text = text.into_bytes_with_nul(); + unsafe { + BNNotifyOutputForScriptingInstance( + self.as_raw(), + text.as_ref().as_ptr() as *const ffi::c_char, + ) + } + } + + pub fn notify_input_ready_state(&self, state: ScriptingProviderInputReadyState) { + unsafe { BNNotifyInputReadyStateForScriptingInstance(self.as_raw(), state) } + } + + pub fn register_output_listener( + self, + listener: L, + ) -> ScriptingInstanceWithListener { + 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, + output: Some(cb_output::), + warning: Some(cb_warning::), + error: Some(cb_error::), + inputReadyStateChanged: Some(cb_input_ready_state_changed::), + }; + unsafe { BNRegisterScriptingInstanceOutputListener(self.as_raw(), &mut callbacks) } + + ScriptingInstanceWithListener { + handle: self, + listener, + } + } + + pub fn delimiters(&self) -> BnString { + let result = unsafe { BNGetScriptingInstanceDelimiters(self.as_raw()) }; + assert!(!result.is_null()); + unsafe { BnString::from_raw(result as *mut ffi::c_char) } + } + + pub fn set_delimiters(&self, delimiters: S) { + let delimiters = delimiters.into_bytes_with_nul(); + unsafe { + BNSetScriptingInstanceDelimiters( + self.as_raw(), + delimiters.as_ref().as_ptr() as *const ffi::c_char, + ) + } + } + + pub fn input_ready_state(&self) -> ScriptingProviderInputReadyState { + unsafe { BNGetScriptingInstanceInputReadyState(self.as_raw()) } + } + + pub fn execute_script_input( + &self, + input: S, + ) -> ScriptingProviderExecuteResult { + let input = input.into_bytes_with_nul(); + unsafe { + BNExecuteScriptInput(self.as_raw(), input.as_ref().as_ptr() as *const ffi::c_char) + } + } + + pub fn execute_script_input_from_filename( + &self, + filename: S, + ) -> ScriptingProviderExecuteResult { + let filename = filename.into_bytes_with_nul(); + unsafe { + BNExecuteScriptInputFromFilename( + self.as_raw(), + filename.as_ref().as_ptr() as *const ffi::c_char, + ) + } + } + + pub fn cancel_script_input(&self) { + unsafe { BNCancelScriptInput(self.as_raw()) } + } + + pub fn release_binary_view(&self, view: &BinaryView) { + unsafe { BNScriptingInstanceReleaseBinaryView(self.as_raw(), view.handle) } + } + + pub fn set_current_binary_view(&self, view: &BinaryView) { + unsafe { BNSetScriptingInstanceCurrentBinaryView(self.as_raw(), view.handle) } + } + + pub fn set_current_function(&self, view: &Function) { + unsafe { BNSetScriptingInstanceCurrentFunction(self.as_raw(), view.handle) } + } + + pub fn set_current_basic_block(&self, view: &BasicBlock) { + unsafe { BNSetScriptingInstanceCurrentBasicBlock(self.as_raw(), view.handle) } + } + + pub fn set_current_address(&self, address: u64) { + unsafe { BNSetScriptingInstanceCurrentAddress(self.as_raw(), address) } + } + + pub fn set_current_selection(&self, begin: u64, end: u64) { + unsafe { BNSetScriptingInstanceCurrentSelection(self.as_raw(), begin, end) } + } + + pub fn complete_input(&self, text: S, state: u64) -> BnString { + let text = text.into_bytes_with_nul(); + let result = unsafe { + BNScriptingInstanceCompleteInput( + self.as_raw(), + text.as_ref().as_ptr() as *const ffi::c_char, + state, + ) + }; + assert!(!result.is_null()); + unsafe { BnString::from_raw(result) } + } + + pub fn stop(&self) { + unsafe { BNStopScriptingInstance(self.as_raw()) } + } +} + +pub struct ScriptingInstanceWithListener { + handle: ScriptingInstance, + listener: Pin>, +} + +impl AsRef for ScriptingInstanceWithListener { + fn as_ref(&self) -> &ScriptingInstance { + &self.handle + } +} + +impl core::ops::Deref for ScriptingInstanceWithListener { + type Target = ScriptingInstance; + fn deref(&self) -> &Self::Target { + &self.handle + } +} + +impl ScriptingInstanceWithListener { + pub fn unregister(mut self) -> ScriptingInstance { + let mut callbacks = BNScriptingOutputListener { + context: unsafe { self.listener.as_mut().get_unchecked_mut() } as *mut _ + as *mut ffi::c_void, + output: Some(cb_output::), + warning: Some(cb_warning::), + error: Some(cb_error::), + inputReadyStateChanged: Some(cb_input_ready_state_changed::), + }; + 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 { + fn new(core: ScriptingProvider) -> Self; + fn create_instance(&self) -> ScriptingInstance; + fn load_module(&self, repo_path: &str, plugin_path: &str, force: bool) -> bool; + fn install_modules(&self, modules: &str) -> bool; +} + +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 execute_script_input(&self, input: &str) -> ScriptingProviderExecuteResult; + fn execute_script_input_from_filename(&self, input: &str) -> ScriptingProviderExecuteResult; + fn cancel_script_input(&self); + fn release_binary_view(&self, view: &BinaryView); + fn set_current_binary_view(&self, view: &BinaryView); + fn set_current_function(&self, func: &Function); + fn set_current_basic_block(&self, block: &BasicBlock); + fn set_current_address(&self, addr: u64); + fn set_current_selection(&self, begin: u64, end: u64); + fn complete_input(&self, text: &str, state: u64) -> String; + fn stop(&self); +} + +pub trait ScriptingOutputListener: Sync + Send { + fn output(&self, text: &str); + fn warning(&self, text: &str); + fn error(&self, text: &str); + fn input_ready_state_changed(&self, state: ScriptingProviderInputReadyState); +} + +pub fn register_scripting_provider(name: N, api_name: A) -> (&'static S, ScriptingProvider) +where + N: BnStrCompatible, + A: BnStrCompatible, + S: ScriptingCustomProvider + 'static, +{ + 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( + name.as_ref().as_ptr() as *const ffi::c_char, + api_name.as_ref().as_ptr() as *const ffi::c_char, + &mut BNScriptingProviderCallbacks { + context: provider_uinit as *mut _ as *mut ffi::c_void, + createInstance: Some(cb_create_instance::), + loadModule: Some(cb_load_module::), + installModules: Some(cb_install_modules::), + }, + ) + }; + let provider_core = unsafe { ScriptingProvider::from_raw(ptr::NonNull::new(result).unwrap()) }; + provider_uinit.write(S::new(provider_core)); + (unsafe { provider_uinit.assume_init_ref() }, provider_core) +} + +unsafe extern "C" fn cb_create_instance( + ctxt: *mut ffi::c_void, +) -> *mut BNScriptingInstance { + let ctxt = &mut *(ctxt as *mut S); + ctxt.create_instance().into_raw() +} + +unsafe extern "C" fn cb_load_module( + ctxt: *mut ffi::c_void, + repo_path: *const ffi::c_char, + plugin_path: *const ffi::c_char, + force: bool, +) -> bool { + let ctxt = &mut *(ctxt as *mut S); + let repo_path = ffi::CStr::from_ptr(repo_path); + let plugin_path = ffi::CStr::from_ptr(plugin_path); + ctxt.load_module( + &repo_path.to_string_lossy(), + &plugin_path.to_string_lossy(), + force, + ) +} + +unsafe extern "C" fn cb_install_modules( + ctxt: *mut ffi::c_void, + modules: *const ffi::c_char, +) -> bool { + let ctxt = &mut *(ctxt as *mut S); + let modules = ffi::CStr::from_ptr(modules); + ctxt.install_modules(&modules.to_string_lossy()) +} + +unsafe extern "C" fn cb_destroy_instance(ctxt: *mut ffi::c_void) { + drop(Box::from_raw(ctxt as *mut S)) +} + +unsafe extern "C" fn cb_external_ref_taken(ctxt: *mut ffi::c_void) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.external_ref_taken() +} + +unsafe extern "C" fn cb_external_ref_released( + ctxt: *mut ffi::c_void, +) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.external_ref_released() +} + +unsafe extern "C" fn cb_execute_script_input( + ctxt: *mut ffi::c_void, + input: *const ffi::c_char, +) -> BNScriptingProviderExecuteResult { + let input = ffi::CStr::from_ptr(input); + let ctxt = &mut *(ctxt as *mut S); + ctxt.execute_script_input(&input.to_string_lossy()) +} + +unsafe extern "C" fn cb_execute_script_input_from_filename( + ctxt: *mut ffi::c_void, + input: *const ffi::c_char, +) -> BNScriptingProviderExecuteResult { + let input = ffi::CStr::from_ptr(input); + let ctxt = &mut *(ctxt as *mut S); + ctxt.execute_script_input(&input.to_string_lossy()) +} + +unsafe extern "C" fn cb_cancel_script_input(ctxt: *mut ffi::c_void) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.cancel_script_input() +} + +unsafe extern "C" fn cb_release_binary_view( + ctxt: *mut ffi::c_void, + view: *mut BNBinaryView, +) { + let view = BinaryView { handle: view }; + let ctxt = &mut *(ctxt as *mut S); + ctxt.release_binary_view(&view) +} + +unsafe extern "C" fn cb_set_current_binary_view( + ctxt: *mut ffi::c_void, + view: *mut BNBinaryView, +) { + let view = BinaryView { handle: view }; + let ctxt = &mut *(ctxt as *mut S); + ctxt.set_current_binary_view(&view) +} + +unsafe extern "C" fn cb_set_current_function( + ctxt: *mut ffi::c_void, + func: *mut BNFunction, +) { + let func = Function { handle: func }; + let ctxt = &mut *(ctxt as *mut S); + ctxt.set_current_function(&func) +} + +unsafe extern "C" fn cb_set_current_basic_block( + ctxt: *mut ffi::c_void, + block: *mut BNBasicBlock, +) { + let block = BasicBlock::from_raw(block, NativeBlock::new()); + let ctxt = &mut *(ctxt as *mut S); + ctxt.set_current_basic_block(&block) +} + +unsafe extern "C" fn cb_set_current_address( + ctxt: *mut ffi::c_void, + addr: u64, +) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.set_current_address(addr) +} + +unsafe extern "C" fn cb_set_current_selection( + ctxt: *mut ffi::c_void, + begin: u64, + end: u64, +) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.set_current_selection(begin, end) +} + +unsafe extern "C" fn cb_complete_input( + ctxt: *mut ffi::c_void, + text: *const ffi::c_char, + state: u64, +) -> *mut ffi::c_char { + let ctxt = &mut *(ctxt as *mut S); + let text = ffi::CStr::from_ptr(text); + let result = ctxt.complete_input(&text.to_string_lossy(), state); + BnString::into_raw(BnString::new(result)) +} + +unsafe extern "C" fn cb_stop(ctxt: *mut ffi::c_void) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.stop() +} + +unsafe extern "C" fn cb_output( + ctxt: *mut ffi::c_void, + text: *const ffi::c_char, +) { + let ctxt = &mut *(ctxt as *mut S); + let text = ffi::CStr::from_ptr(text); + ctxt.output(&text.to_string_lossy()) +} + +unsafe extern "C" fn cb_warning( + ctxt: *mut ffi::c_void, + text: *const ffi::c_char, +) { + let ctxt = &mut *(ctxt as *mut S); + let text = ffi::CStr::from_ptr(text); + ctxt.warning(&text.to_string_lossy()) +} + +unsafe extern "C" fn cb_error( + ctxt: *mut ffi::c_void, + text: *const ffi::c_char, +) { + let ctxt = &mut *(ctxt as *mut S); + let text = ffi::CStr::from_ptr(text); + ctxt.error(&text.to_string_lossy()) +} + +unsafe extern "C" fn cb_input_ready_state_changed( + ctxt: *mut ffi::c_void, + state: BNScriptingProviderInputReadyState, +) { + let ctxt = &mut *(ctxt as *mut S); + ctxt.input_ready_state_changed(state) +} From eb5683e02bd1aa8b61d39f60e870aef5334939e7 Mon Sep 17 00:00:00 2001 From: rbran Date: Mon, 3 Feb 2025 15:58:09 -0300 Subject: [PATCH 2/4] impl Ref for ScriptingInstance --- rust/src/scriptingprovider.rs | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/rust/src/scriptingprovider.rs b/rust/src/scriptingprovider.rs index d02b089a1..9ca6b4115 100644 --- a/rust/src/scriptingprovider.rs +++ b/rust/src/scriptingprovider.rs @@ -6,13 +6,13 @@ use binaryninjacore_sys::*; use crate::basic_block::BasicBlock; use crate::binary_view::BinaryView; use crate::function::{Function, NativeBlock}; -use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner}; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Ref, RefCountable}; use crate::string::{BnStrCompatible, BnString}; pub type ScriptingProviderExecuteResult = BNScriptingProviderExecuteResult; pub type ScriptingProviderInputReadyState = BNScriptingProviderInputReadyState; -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Eq, PartialEq, Hash)] #[repr(transparent)] pub struct ScriptingProvider { handle: ptr::NonNull, @@ -143,16 +143,22 @@ pub struct ScriptingInstance { handle: ptr::NonNull, } -impl Clone for ScriptingInstance { - fn clone(&self) -> Self { - let result = unsafe { BNNewScriptingInstanceReference(self.as_raw()) }; - unsafe { Self::from_raw(ptr::NonNull::new(result).unwrap()) } +impl ToOwned for ScriptingInstance { + type Owned = Ref; + + fn to_owned(&self) -> Self::Owned { + unsafe { Self::inc_ref(self) } } } -impl Drop for ScriptingInstance { - fn drop(&mut self) { - unsafe { BNFreeScriptingInstance(self.as_raw()) } +unsafe impl RefCountable for ScriptingInstance { + unsafe fn inc_ref(handle: &Self) -> Ref { + let result = unsafe { BNNewScriptingInstanceReference(handle.as_raw()) }; + unsafe { Self::ref_from_raw(ptr::NonNull::new(result).unwrap()) } + } + + unsafe fn dec_ref(handle: &Self) { + unsafe { BNFreeScriptingInstance(handle.as_raw()) } } } @@ -161,6 +167,10 @@ impl ScriptingInstance { Self { handle } } + pub(crate) unsafe fn ref_from_raw(handle: ptr::NonNull) -> Ref { + Ref::new(Self { handle }) + } + pub(crate) unsafe fn into_raw(self) -> *mut BNScriptingInstance { mem::ManuallyDrop::new(self).handle.as_ptr() } From f1890a8678997a643b3e75c8aff0dbfa4c174d70 Mon Sep 17 00:00:00 2001 From: rbran Date: Tue, 4 Feb 2025 15:27:20 +0000 Subject: [PATCH 3/4] fix ScriptingInstance notify warning/error --- rust/src/scriptingprovider.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/src/scriptingprovider.rs b/rust/src/scriptingprovider.rs index 9ca6b4115..1ecf03dc4 100644 --- a/rust/src/scriptingprovider.rs +++ b/rust/src/scriptingprovider.rs @@ -193,7 +193,7 @@ impl ScriptingInstance { pub fn notify_warning(&self, text: S) { let text = text.into_bytes_with_nul(); unsafe { - BNNotifyOutputForScriptingInstance( + BNNotifyWarningForScriptingInstance( self.as_raw(), text.as_ref().as_ptr() as *const ffi::c_char, ) @@ -203,7 +203,7 @@ impl ScriptingInstance { pub fn notify_error(&self, text: S) { let text = text.into_bytes_with_nul(); unsafe { - BNNotifyOutputForScriptingInstance( + BNNotifyErrorForScriptingInstance( self.as_raw(), text.as_ref().as_ptr() as *const ffi::c_char, ) From 934192ad25e643057faa2eebf4733833af2501b6 Mon Sep 17 00:00:00 2001 From: rbran Date: Tue, 4 Feb 2025 19:29:31 +0000 Subject: [PATCH 4/4] fix scriptingprovider inconsistencies --- rust/src/scriptingprovider.rs | 172 +++++++++++++++++----------- rust/tests/scriptingprovider.rs | 195 ++++++++++++++++++++++++++++++++ 2 files changed, 302 insertions(+), 65 deletions(-) create mode 100644 rust/tests/scriptingprovider.rs diff --git a/rust/src/scriptingprovider.rs b/rust/src/scriptingprovider.rs index 1ecf03dc4..6cb2f2cb6 100644 --- a/rust/src/scriptingprovider.rs +++ b/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, } +unsafe impl Sync for ScriptingProvider {} +unsafe impl Send for ScriptingProvider {} + impl ScriptingProvider { pub(crate) unsafe fn from_raw(handle: ptr::NonNull) -> Self { Self { handle } @@ -95,30 +100,9 @@ impl ScriptingProvider { } } - pub fn new_instance(&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::), - externalRefTaken: Some(cb_external_ref_taken::), - externalRefReleased: Some(cb_external_ref_released::), - executeScriptInput: Some(cb_execute_script_input::), - executeScriptInputFromFilename: Some(cb_execute_script_input_from_filename::), - cancelScriptInput: Some(cb_cancel_script_input::), - releaseBinaryView: Some(cb_release_binary_view::), - setCurrentBinaryView: Some(cb_set_current_binary_view::), - setCurrentFunction: Some(cb_set_current_function::), - setCurrentBasicBlock: Some(cb_set_current_basic_block::), - setCurrentAddress: Some(cb_set_current_address::), - setCurrentSelection: Some(cb_set_current_selection::), - completeInput: Some(cb_complete_input::), - stop: Some(cb_stop::), - }; - 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 { + 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, } +unsafe impl Sync for ScriptingInstance {} +unsafe impl Send for ScriptingInstance {} + impl ToOwned for ScriptingInstance { type Owned = Ref; @@ -163,18 +150,10 @@ unsafe impl RefCountable for ScriptingInstance { } impl ScriptingInstance { - pub(crate) unsafe fn from_raw(handle: ptr::NonNull) -> Self { - Self { handle } - } - pub(crate) unsafe fn ref_from_raw(handle: ptr::NonNull) -> Ref { 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( - self, + pub fn register_output_listener<'a, L: ScriptingOutputListener>( + &'a self, listener: L, - ) -> ScriptingInstanceWithListener { + ) -> 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 { - handle: ScriptingInstance, +pub struct ScriptingInstanceWithListener<'a, L: ScriptingOutputListener> { + handle: &'a ScriptingInstance, listener: Pin>, } -impl AsRef for ScriptingInstanceWithListener { - fn as_ref(&self) -> &ScriptingInstance { - &self.handle +impl AsRef for ScriptingInstanceWithListener<'_, L> { + fn as_ref(&self) -> &L { + &self.listener } } -impl core::ops::Deref for ScriptingInstanceWithListener { - type Target = ScriptingInstance; +impl core::ops::Deref for ScriptingInstanceWithListener<'_, L> { + type Target = L; fn deref(&self) -> &Self::Target { - &self.handle + &self.listener } } -impl ScriptingInstanceWithListener { - pub fn unregister(mut self) -> ScriptingInstance { +impl 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 ScriptingInstanceWithListener { 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, Ref) + where + Self: Sized, + { + register_instance::(self) + } } +pub struct CoreScriptingInstance(Ref); + 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( + provider: &S, +) -> (Arc, Ref) { + let instance_uninit: *const MaybeUninit = Arc::into_raw(Arc::new_uninit()); + let mut callbacks = BNScriptingInstanceCallbacks { + context: instance_uninit as *mut ffi::c_void, + destroyInstance: Some(cb_destroy_instance::), + externalRefTaken: Some(cb_external_ref_taken::), + externalRefReleased: Some(cb_external_ref_released::), + executeScriptInput: Some(cb_execute_script_input::), + executeScriptInputFromFilename: Some(cb_execute_script_input_from_filename::), + cancelScriptInput: Some(cb_cancel_script_input::), + releaseBinaryView: Some(cb_release_binary_view::), + setCurrentBinaryView: Some(cb_set_current_binary_view::), + setCurrentFunction: Some(cb_set_current_function::), + setCurrentBasicBlock: Some(cb_set_current_basic_block::), + setCurrentAddress: Some(cb_set_current_address::), + setCurrentSelection: Some(cb_set_current_selection::), + completeInput: Some(cb_complete_input::), + stop: Some(cb_stop::), + }; + 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::>::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( 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( @@ -456,19 +499,17 @@ unsafe extern "C" fn cb_install_modules( } unsafe extern "C" fn cb_destroy_instance(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(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( 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( @@ -477,7 +518,8 @@ unsafe extern "C" fn cb_execute_script_input( ) -> 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( diff --git a/rust/tests/scriptingprovider.rs b/rust/tests/scriptingprovider.rs new file mode 100644 index 000000000..feb8383f1 --- /dev/null +++ b/rust/tests/scriptingprovider.rs @@ -0,0 +1,195 @@ +use std::sync::Mutex; + +use binaryninja::basic_block::BasicBlock; +use binaryninja::binary_view::BinaryView; +use binaryninja::function::{Function, NativeBlock}; +use binaryninja::headless::Session; +use binaryninja::scriptingprovider::{ + register_scripting_provider, CoreScriptingInstance, ScriptingCustomProvider, + ScriptingInstanceCallbacks, ScriptingOutputListener, ScriptingProvider, + ScriptingProviderExecuteResult, ScriptingProviderInputReadyState, +}; +use rstest::*; + +#[fixture] +#[once] +fn session() -> Session { + Session::new().expect("Failed to initialize session") +} + +struct MyScriptingProvider { + core: ScriptingProvider, +} + +impl ScriptingCustomProvider for MyScriptingProvider { + type Instance = MyScriptingProviderInstance; + + fn new(core: ScriptingProvider) -> Self { + Self { core } + } + + fn init_instance(&self, handle: CoreScriptingInstance) -> Self::Instance { + MyScriptingProviderInstance { core: handle } + } + + fn get_core(&self) -> &ScriptingProvider { + &self.core + } + + fn load_module(&self, _repo_path: &str, _plugin_path: &str, _force: bool) -> bool { + todo!() + } + + fn install_modules(&self, _modules: &str) -> bool { + todo!() + } +} + +struct MyScriptingProviderInstance { + core: CoreScriptingInstance, +} + +impl ScriptingInstanceCallbacks for MyScriptingProviderInstance { + fn get_core(&self) -> &CoreScriptingInstance { + &self.core + } + + fn execute_script_input(&self, input: &str) -> ScriptingProviderExecuteResult { + self.notify_output(&format!("execute_script_input({input})")); + ScriptingProviderExecuteResult::SuccessfulScriptExecution + } + + fn execute_script_input_from_filename(&self, input: &str) -> ScriptingProviderExecuteResult { + self.notify_output(&format!("execute_script_input_from_filename({input})")); + ScriptingProviderExecuteResult::SuccessfulScriptExecution + } + + fn cancel_script_input(&self) {} + + fn release_binary_view(&self, _view: &BinaryView) { + todo!() + } + + fn set_current_binary_view(&self, _view: &BinaryView) { + todo!() + } + + fn set_current_function(&self, _func: &Function) { + todo!() + } + + fn set_current_basic_block(&self, _block: &BasicBlock) { + todo!() + } + + fn set_current_address(&self, _addr: u64) { + todo!() + } + + fn set_current_selection(&self, _begin: u64, _end: u64) { + todo!() + } + + fn complete_input(&self, _text: &str, _state: u64) -> String { + todo!() + } + + fn stop(&self) { + todo!() + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OutputType { + Output, + Warning, + Error, +} + +#[derive(Debug, Default)] +struct MyListener { + output: Mutex>, +} + +impl ScriptingOutputListener for MyListener { + fn output(&self, text: &str) { + let mut output = self.output.lock().unwrap(); + output.push((OutputType::Output, text.to_string())) + } + + fn warning(&self, text: &str) { + let mut output = self.output.lock().unwrap(); + output.push((OutputType::Warning, text.to_string())) + } + + fn error(&self, text: &str) { + let mut output = self.output.lock().unwrap(); + output.push((OutputType::Error, text.to_string())) + } + + fn input_ready_state_changed(&self, _state: ScriptingProviderInputReadyState) {} +} + +#[rstest] +fn register_script_provider(_session: &Session) { + let (rust_provider, _core_provider) = register_scripting_provider::<_, _, MyScriptingProvider>( + "RustScriptProvider", + "RustScriptProvider", + ); + let (_rust_instance, _core_instance) = rust_provider.new_instance(); +} + +#[rstest] +fn listen_script_provider(_session: &Session) { + let (rust_provider, core_provider) = register_scripting_provider::<_, _, MyScriptingProvider>( + "RustScriptProvider", + "RustScriptProvider", + ); + let (rust_instance, core_instance) = rust_provider.new_instance(); + + let listener1 = core_instance.register_output_listener(MyListener::default()); + assert_eq!( + rust_instance.execute_script_input("test"), + ScriptingProviderExecuteResult::SuccessfulScriptExecution + ); + let listener2 = rust_instance.register_output_listener(MyListener::default()); + assert_eq!( + rust_instance.execute_script_input_from_filename("test2"), + ScriptingProviderExecuteResult::SuccessfulScriptExecution + ); + + let output1 = listener1.output.lock().unwrap(); + assert_eq!( + &*output1, + &[ + (OutputType::Output, "execute_script_input(test)".to_string()), + ( + OutputType::Output, + "execute_script_input_from_filename(test2)".to_string() + ) + ] + ); + let output2 = listener2.output.lock().unwrap(); + assert_eq!( + &*output2, + &[( + OutputType::Output, + "execute_script_input_from_filename(test2)".to_string() + )] + ); + + let other_core_instance = core_provider.new_instance(); + let listener3 = other_core_instance.register_output_listener(MyListener::default()); + assert_eq!( + other_core_instance.execute_script_input("test3"), + ScriptingProviderExecuteResult::SuccessfulScriptExecution + ); + let output3 = listener3.output.lock().unwrap(); + assert_eq!( + &*output3, + &[( + OutputType::Output, + "execute_script_input(test3)".to_string() + ),] + ); +}