diff --git a/rust/src/lib.rs b/rust/src/lib.rs index e24a3f784..d8e8c2cde 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -149,6 +149,7 @@ pub mod interaction; pub mod linearview; pub mod llil; pub mod logger; +pub mod mainthread; pub mod metadata; pub mod mlil; pub mod platform; diff --git a/rust/src/mainthread.rs b/rust/src/mainthread.rs new file mode 100644 index 000000000..545ec2638 --- /dev/null +++ b/rust/src/mainthread.rs @@ -0,0 +1,204 @@ +//! This module provides two ways to execute "jobs": +//! +//! 1. On the Binary Ninja main thread (the UI event thread when running in the GUI application): +//! * [execute_on_main_thread] +//! * [execute_on_main_thread_and_wait] +//! 2. On a worker thread +//! +//! Any manipulation of the GUI should be performed on the main thread, but any +//! non-GUI work is generally better to be performed using a worker. This is +//! especially true for any longer-running work, as the user interface will +//! be unable to update itself while a job is executing on the main thread. +//! +//! There are three worker queues, in order of decreasing priority: +//! +//! 1. The Interactive Queue ([worker_interactive_enqueue]) +//! 2. The Priority Queue ([worker_priority_enqueue]) +//! 3. The Worker Queue ([worker_enqueue]) +//! +//! All of these queues are serviced by the same pool of worker threads. The +//! difference between the queues is basically one of priority: one queue must +//! be empty of jobs before a worker thread will execute a job from a lower +//! priority queue. +//! +//! The default maximum number of concurrent worker threads is controlled by the +//! `analysis.limits.workerThreadCount` setting but can be adjusted at runtime via +//! [set_worker_thread_count]. +//! +//! The worker threads are native threads, managed by the Binary Ninja core. If +//! more control over the thread is required, consider using the +//! [crate::backgroundtask::BackgroundTask] class. + +use core::{ffi, mem, ptr}; + +use binaryninjacore_sys::*; + +use crate::string::BnStrCompatible; + +#[repr(transparent)] +pub struct MainThreadAction { + handle: ptr::NonNull, +} + +impl Clone for MainThreadAction { + fn clone(&self) -> Self { + let result = unsafe { BNNewMainThreadActionReference(self.as_raw()) }; + unsafe { Self::from_raw(ptr::NonNull::new(result).unwrap()) } + } +} + +impl Drop for MainThreadAction { + fn drop(&mut self) { + unsafe { BNFreeMainThreadAction(self.as_raw()) } + } +} + +impl MainThreadAction { + pub(crate) unsafe fn from_raw(handle: ptr::NonNull) -> Self { + Self { handle } + } + + pub(crate) unsafe fn ref_from_raw(handle: &*mut BNMainThreadAction) -> &Self { + assert!(!handle.is_null()); + mem::transmute(handle) + } + + #[allow(clippy::mut_from_ref)] + pub(crate) unsafe fn as_raw(&self) -> &mut BNMainThreadAction { + &mut *self.handle.as_ptr() + } + + pub fn execute(&self) { + unsafe { BNExecuteMainThreadAction(self.as_raw()) } + } + + pub fn is_done(&self) -> bool { + unsafe { BNIsMainThreadActionDone(self.as_raw()) } + } + + pub fn wait(&self) { + unsafe { BNWaitForMainThreadAction(self.as_raw()) } + } +} + +/// Takes a single parameter which is a function that will be executed +/// on the main Binary Ninja thread. +/// +///
+/// +/// May be required for some GUI operations, but should be used sparingly as it can block the UI. +/// +///
+pub fn execute_on_main_thread(action: F) -> Option { + let action = Box::leak(Box::new(action)); + let result = unsafe { + BNExecuteOnMainThread( + action as *mut F as *mut ffi::c_void, + Some(cb_execute_with_drop::), + ) + }; + ptr::NonNull::new(result).map(|h| unsafe { MainThreadAction::from_raw(h) }) +} + +/// Takes a single parameter which is a function that will be +/// executed on the main Binary Ninja thread and will block execution of further python until the function returns. +/// +///
+/// +/// May be required for some GUI operations, but should be used sparingly as it can block the UI. +/// +///
+pub fn execute_on_main_thread_and_wait(mut action: F) { + unsafe { + BNExecuteOnMainThreadAndWait( + &mut action as *mut F as *mut ffi::c_void, + Some(cb_execute_without_drop::), + ) + } +} + +pub fn worker_enqueue(action: F, name: S) { + let action = Box::leak(Box::new(action)); + let name = name.into_bytes_with_nul(); + unsafe { + BNWorkerEnqueueNamed( + action as *mut F as *mut ffi::c_void, + Some(cb_execute_with_drop::), + name.as_ref().as_ptr() as *const ffi::c_char, + ) + } +} + +pub fn worker_priority_enqueue(action: F, name: S) { + let action = Box::leak(Box::new(action)); + let name = name.into_bytes_with_nul(); + unsafe { + BNWorkerPriorityEnqueueNamed( + action as *mut F as *mut ffi::c_void, + Some(cb_execute_with_drop::), + name.as_ref().as_ptr() as *const ffi::c_char, + ) + } +} + +pub fn worker_interactive_enqueue(action: F, name: S) { + let action = Box::leak(Box::new(action)); + let name = name.into_bytes_with_nul(); + unsafe { + BNWorkerInteractiveEnqueueNamed( + action as *mut F as *mut ffi::c_void, + Some(cb_execute_with_drop::), + name.as_ref().as_ptr() as *const ffi::c_char, + ) + } +} + +/// Returns the number of worker threads that are currently running. +/// By default, this is the number of cores on the system minus one, however this can be changed with +/// [set_worker_thread_count]. +pub fn get_worker_thread_count() -> usize { + unsafe { BNGetWorkerThreadCount() } +} + +/// Sets the number of worker threads that are currently running. +/// By default, this is the number of cores on the system minus one. +pub fn set_worker_thread_count(count: usize) { + unsafe { BNSetWorkerThreadCount(count) } +} + +pub fn is_main_thread() -> bool { + unsafe { BNIsMainThread() } +} + +pub trait MainThreadProvider { + fn add_action(&self, action: &MainThreadAction); +} + +pub fn register_main_thread(action: T) { + // SAFETY new main thread provider is never freed + let context = Box::leak(Box::new(action)); + let mut callback = BNMainThreadCallbacks { + context: context as *mut T as *mut ffi::c_void, + addAction: Some(cb_add_action::), + }; + unsafe { BNRegisterMainThread(&mut callback) } +} + +unsafe extern "C" fn cb_add_action( + ctxt: *mut ffi::c_void, + action: *mut BNMainThreadAction, +) { + let ctxt = &mut *(ctxt as *mut T); + ctxt.add_action(MainThreadAction::ref_from_raw(&action)) +} + +unsafe extern "C" fn cb_execute_without_drop(ctxt: *mut ffi::c_void) { + let ctxt = &mut *(ctxt as *mut F); + ctxt() +} + +unsafe extern "C" fn cb_execute_with_drop(ctxt: *mut ffi::c_void) { + let mut ctxt = Box::from_raw(ctxt as *mut F); + ctxt(); + drop(ctxt); +}