diff --git a/rust/src/data_notification.rs b/rust/src/data_notification.rs new file mode 100644 index 000000000..464e4c2a4 --- /dev/null +++ b/rust/src/data_notification.rs @@ -0,0 +1,437 @@ +use core::ffi; +use core::ptr; +use std::ffi::CStr; + +use binaryninjacore_sys::{ + BNBinaryDataNotification, BNBinaryView, BNComponent, BNDataVariable, BNExternalLibrary, + BNExternalLocation, BNFunction, BNQualifiedName, BNRegisterDataNotification, BNSection, + BNSegment, BNStringType, BNSymbol, BNTagReference, BNTagType, BNType, BNTypeArchive, + BNUndoEntry, BNUnregisterDataNotification, +}; + +use crate::binary_view::BinaryView; +use crate::component::Component; +use crate::database::undo::UndoEntry; +use crate::disassembly::StringType; +use crate::external_library::{ExternalLibrary, ExternalLocation}; +use crate::function::Function; +use crate::rc::Ref; +use crate::section::Section; +use crate::segment::Segment; +use crate::symbol::Symbol; +use crate::tags::{TagReference, TagType}; +use crate::type_archive::TypeArchive; +use crate::types::{QualifiedName, Type}; +use crate::variable::DataVariable; + +macro_rules! trait_handler { +( + $( + $ffi_param_name:ident => $fun_name:ident( + $( + $arg_name:ident: + $raw_arg_type:ty: + $arg_type:ty = + $value_calculated:expr + ),* $(,)? + ) $(-> $ret_type:ty)? + ),* $(,)? +) => { + /// Used to describe which call should be triggered by BinaryNinja. + /// By default all calls are disabled. + /// + /// Used by [CustomDataNotification::register] + #[derive(Default)] + pub struct DataNotificationTriggers { + $($fun_name: bool,)* + } + impl DataNotificationTriggers { + $( + pub fn $fun_name(mut self) -> Self { + self.$fun_name = true; + self + } + )* + } + pub trait CustomDataNotification: Sync + Send { + $( + #[inline] + #[allow(unused_variables)] + fn $fun_name(&mut self, $($arg_name: $arg_type),*) $(-> $ret_type)* { + $( <$ret_type as Default>::default() )* + } + )* + + fn register<'a>( + self, + view: &BinaryView, + triggers: DataNotificationTriggers, + ) -> DataNotificationHandle<'a, Self> where Self: 'a + Sized { + register_data_notification(view, self, triggers) + } + } + $( + unsafe extern "C" fn $fun_name( + ctxt: *mut ::std::os::raw::c_void, + $($arg_name: $raw_arg_type),* + ) $(-> $ret_type)* { + let handle: &mut H = &mut *(ctxt as *mut H); + handle.$fun_name($($value_calculated),*) + } + )* + fn register_data_notification<'a, H: CustomDataNotification + 'a>( + view: &BinaryView, + notify: H, + triggers: DataNotificationTriggers, + ) -> DataNotificationHandle<'a, H> { + // SAFETY: this leak is undone on drop + let leak_notify = Box::leak(Box::new(notify)); + let mut handle = BNBinaryDataNotification { + context: leak_notify as *mut _ as *mut ffi::c_void, + $($ffi_param_name: triggers.$fun_name.then_some($fun_name::)),* + }; + unsafe { BNRegisterDataNotification(view.handle, &mut handle) }; + DataNotificationHandle { + bv: view.to_owned(), + handle, + _life: std::marker::PhantomData, + } + } + + /// Implement closures that will be called by BinaryNinja on the event of + /// data modification. + /// + /// example: + /// ```no_run + /// # use binaryninja::data_notification::DataNotificationClosure; + /// # use binaryninja::function::Function; + /// # use binaryninja::binary_view::BinaryView; + /// # let bv: BinaryView = todo!(); + /// let custom = DataNotificationClosure::default() + /// .function_updated(|_bv: &BinaryView, _func: &Function| todo!() ) + /// // other calls should be added here + /// .register(&bv); + /// ``` + pub struct DataNotificationClosure<'a> { + $( + $fun_name: Option $ret_type)* + Sync + Send + 'a + >>, + )* + } + + impl Default for DataNotificationClosure<'_> { + fn default() -> Self { + Self { $($fun_name: None,)* } + } + } + + impl<'a> DataNotificationClosure<'a> { + pub fn new() -> Self { + Default::default() + } + $( + pub fn $fun_name $ret_type)* + Sync + Send + 'a>(mut self, param: F) -> Self { + self.$fun_name = Some(Box::new(param)); + self + } + )* + + pub fn register( + self, + view: &BinaryView, + ) -> DataNotificationHandle<'a, Self> { + let mut triggers = DataNotificationTriggers::default(); + $( + if self.$fun_name.is_some() { + triggers = triggers.$fun_name(); + } + )* + register_data_notification(view, self, triggers) + } + } + + impl CustomDataNotification for DataNotificationClosure<'_> { + $( + fn $fun_name(&mut self, $($arg_name: $arg_type),*) $(-> $ret_type)* { + let Some(handle) = self.$fun_name.as_mut() else { + unreachable!(); + }; + handle($($arg_name),*) + } + )* + } +}; +} + +trait_handler! { + notificationBarrier => notification_barrier( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + ) -> u64, + dataWritten => data_written( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + offset: u64: u64 = offset, + len: usize: usize = len, + ), + dataInserted => data_inserted( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + offset: u64: u64 = offset, + len: usize: usize = len, + ), + dataRemoved => data_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + offset: u64: u64 = offset, + // TODO why the len is u64 here? Maybe is a bug on CoreAPI + len: u64: u64 = len, + ), + functionAdded => function_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + func: *mut BNFunction: &Function = &Function::from_raw(func), + ), + functionRemoved => function_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + func: *mut BNFunction: &Function = &Function::from_raw(func), + ), + functionUpdated => function_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + func: *mut BNFunction: &Function = &Function::from_raw(func), + ), + functionUpdateRequested => function_update_requested( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + func: *mut BNFunction: &Function = &Function::from_raw(func), + ), + dataVariableAdded => data_variable_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + var: *mut BNDataVariable: &DataVariable = &DataVariable::from_raw(&*var), + ), + dataVariableRemoved => data_variable_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + var: *mut BNDataVariable: &DataVariable = &DataVariable::from_raw(&*var), + ), + dataVariableUpdated => data_variable_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + var: *mut BNDataVariable: &DataVariable = &DataVariable::from_raw(&*var), + ), + dataMetadataUpdated => data_metadata_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + offset: u64: u64 = offset, + ), + tagTypeUpdated => tag_type_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + tag_type: *mut BNTagType: &TagType = &TagType{ handle: tag_type }, + ), + tagAdded => tag_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + tag_ref: *mut BNTagReference: &TagReference = &TagReference::from(&*tag_ref), + ), + tagRemoved => tag_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + tag_ref: *mut BNTagReference: &TagReference = &TagReference::from(&*tag_ref), + ), + tagUpdated => tag_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + tag_ref: *mut BNTagReference: &TagReference = &TagReference::from(&*tag_ref), + ), + symbolAdded => symbol_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + sym: *mut BNSymbol: &Symbol = &Symbol::from_raw(sym), + ), + symbolRemoved => symbol_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + sym: *mut BNSymbol: &Symbol = &Symbol::from_raw(sym), + ), + symbolUpdated => symbol_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + sym: *mut BNSymbol: &Symbol = &Symbol::from_raw(sym), + ), + stringFound => string_found( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + type_: BNStringType: StringType = type_, + offset: u64: u64 = offset, + len: usize: usize = len, + ), + stringRemoved => string_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + type_: BNStringType: StringType = type_, + offset: u64: u64 = offset, + len: usize: usize = len, + ), + typeDefined => type_defined( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + name: *mut BNQualifiedName: &QualifiedName = &QualifiedName::from_raw(&*name), + type_: *mut BNType: &Type = &Type::from_raw(type_), + ), + typeUndefined => type_undefined( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + name: *mut BNQualifiedName: &QualifiedName = &QualifiedName::from_raw(&*name), + type_: *mut BNType: &Type = &Type::from_raw(type_), + ), + typeReferenceChanged => type_reference_changed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + name: *mut BNQualifiedName: &QualifiedName = &QualifiedName::from_raw(&*name), + type_: *mut BNType: &Type = &Type::from_raw(type_), + ), + typeFieldReferenceChanged => type_field_reference_changed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + name: *mut BNQualifiedName: &QualifiedName = &QualifiedName::from_raw(&*name), + offset: u64: u64 = offset, + ), + segmentAdded => segment_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + segment: *mut BNSegment: &Segment = &Segment::from_raw(segment), + ), + segmentRemoved => segment_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + segment: *mut BNSegment: &Segment = &Segment::from_raw(segment), + ), + segmentUpdated => segment_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + segment: *mut BNSegment: &Segment = &Segment::from_raw(segment), + ), + sectionAdded => section_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + section: *mut BNSection: &Section = &Section::from_raw(section), + ), + sectionRemoved => section_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + section: *mut BNSection: &Section = &Section::from_raw(section), + ), + sectionUpdated => section_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + section: *mut BNSection: &Section = &Section::from_raw(section), + ), + componentNameUpdated => component_name_updated( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + previous_name: *mut ::std::os::raw::c_char: &str = CStr::from_ptr(previous_name).to_str().unwrap(), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + ), + componentAdded => component_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + ), + componentMoved => component_moved( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + former_parent: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(former_parent).unwrap()), + new_parent: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(new_parent).unwrap()), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + ), + componentRemoved => component_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + former_parent: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(former_parent).unwrap()), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + ), + componentFunctionAdded => component_function_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + function: *mut BNFunction: &Function = &Function::from_raw(function), + ), + componentFunctionRemoved => component_function_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + function: *mut BNFunction: &Function = &Function::from_raw(function), + ), + componentDataVariableAdded => component_data_variable_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + var: *mut BNDataVariable: &DataVariable = &DataVariable::from_raw(&*var), + ), + componentDataVariableRemoved => component_data_variable_removed( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + component: *mut BNComponent: &Component = &Component::from_raw(ptr::NonNull::new(component).unwrap()), + var: *mut BNDataVariable: &DataVariable = &DataVariable::from_raw(&*var), + ), + externalLibraryAdded => external_library_added( + data: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(data), + library: *mut BNExternalLibrary: &ExternalLibrary = &ExternalLibrary::from_raw(ptr::NonNull::new(library).unwrap()), + ), + externalLibraryUpdated => external_library_updated( + data: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(data), + library: *mut BNExternalLibrary: &ExternalLibrary = &ExternalLibrary::from_raw(ptr::NonNull::new(library).unwrap()), + ), + externalLibraryRemoved => external_library_removed( + data: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(data), + library: *mut BNExternalLibrary: &ExternalLibrary = &ExternalLibrary::from_raw(ptr::NonNull::new(library).unwrap()), + ), + externalLocationAdded => external_location_added( + data: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(data), + location: *mut BNExternalLocation: &ExternalLocation = &ExternalLocation::from_raw(ptr::NonNull::new(location).unwrap()), + ), + externalLocationUpdated => external_location_updated( + data: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(data), + location: *mut BNExternalLocation: &ExternalLocation = &ExternalLocation::from_raw(ptr::NonNull::new(location).unwrap()), + ), + externalLocationRemoved => external_location_removed( + data: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(data), + location: *mut BNExternalLocation: &ExternalLocation = &ExternalLocation::from_raw(ptr::NonNull::new(location).unwrap()), + ), + typeArchiveAttached => type_archive_attached( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + id: *const ::std::os::raw::c_char: &str = CStr::from_ptr(id).to_str().unwrap(), + path: *const ::std::os::raw::c_char: &[u8] = CStr::from_ptr(path).to_bytes(), + ), + typeArchiveDetached => type_archive_detached( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + id: *const ::std::os::raw::c_char: &str = CStr::from_ptr(id).to_str().unwrap(), + path: *const ::std::os::raw::c_char: &[u8] = CStr::from_ptr(path).to_bytes(), + ), + typeArchiveConnected => type_archive_connected( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + archive: *mut BNTypeArchive: &TypeArchive = &TypeArchive::from_raw(ptr::NonNull::new(archive).unwrap()), + ), + typeArchiveDisconnected => type_archive_disconnected( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + archive: *mut BNTypeArchive: &TypeArchive = &TypeArchive::from_raw(ptr::NonNull::new(archive).unwrap()), + ), + undoEntryAdded => undo_entry_added( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + entry: *mut BNUndoEntry: &UndoEntry = &UndoEntry::from_raw(ptr::NonNull::new(entry).unwrap()), + ), + undoEntryTaken => undo_entry_taken( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + entry: *mut BNUndoEntry: &UndoEntry = &UndoEntry::from_raw(ptr::NonNull::new(entry).unwrap()), + ), + redoEntryTaken => redo_entry_taken( + view: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(view), + entry: *mut BNUndoEntry: &UndoEntry = &UndoEntry::from_raw(ptr::NonNull::new(entry).unwrap()), + ), + rebased => rebased( + oldview: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(oldview), + newview: *mut BNBinaryView: &BinaryView = &BinaryView::from_raw(newview), + ), +} + +pub struct DataNotificationHandle<'a, T: CustomDataNotification> +where + T: 'a, +{ + bv: Ref, + handle: BNBinaryDataNotification, + _life: std::marker::PhantomData<&'a T>, +} + +impl DataNotificationHandle<'_, T> { + unsafe fn unregister_bv(&mut self) { + BNUnregisterDataNotification(self.bv.handle, &mut self.handle) + } + + unsafe fn extract_context(&mut self) -> Box { + Box::from_raw(self.handle.context as *mut T) + } + + pub fn unregister(self) -> T { + // NOTE don't drop the ctxt, return it + let mut slf = core::mem::ManuallyDrop::new(self); + unsafe { slf.unregister_bv() }; + unsafe { *slf.extract_context() } + } +} + +impl Drop for DataNotificationHandle<'_, T> { + fn drop(&mut self) { + unsafe { self.unregister_bv() }; + // drop context, avoid memory leak + let _ctxt = unsafe { self.extract_context() }; + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index a3df0cd35..a3d9c6bfe 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -42,6 +42,7 @@ pub mod component; pub mod confidence; pub mod custom_binary_view; pub mod data_buffer; +pub mod data_notification; pub mod database; pub mod debuginfo; pub mod demangle; diff --git a/rust/src/section.rs b/rust/src/section.rs index c88ef9b0e..67f3e8653 100644 --- a/rust/src/section.rs +++ b/rust/src/section.rs @@ -68,7 +68,7 @@ pub struct Section { } impl Section { - unsafe fn from_raw(handle: *mut BNSection) -> Self { + pub(crate) unsafe fn from_raw(handle: *mut BNSection) -> Self { debug_assert!(!handle.is_null()); Self { handle } } diff --git a/rust/tests/data_notification.rs b/rust/tests/data_notification.rs new file mode 100644 index 000000000..2ba67b7c3 --- /dev/null +++ b/rust/tests/data_notification.rs @@ -0,0 +1,67 @@ +use std::path::PathBuf; + +use binaryninja::binary_view::{BinaryView, BinaryViewExt}; +use binaryninja::data_notification::*; +use binaryninja::function::Function; +use binaryninja::headless::Session; +use binaryninja::tags::TagReference; + +#[test] +fn test_data_notification_dyn_closure() { + let _session = Session::new().expect("Failed to initialize session"); + let out_dir = env!("OUT_DIR").parse::().unwrap(); + let bv = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view"); + + let mut func_updated_count = 0usize; + let custom = DataNotificationClosure::default() + .function_updated(|_bv: &BinaryView, _func: &Function| { + func_updated_count += 1; + }) + .register(&bv); + + let funcs = bv.functions(); + for func in &funcs { + let crash = bv.create_tag_type("Test", "🚧"); + func.add_tag( + &crash, + "Dummy tag", + Some(10.try_into().unwrap()), + true, + None, + ); + } + custom.unregister(); + + assert_eq!(funcs.len(), func_updated_count); +} + +#[test] +fn test_data_notification_impl() { + let _session = Session::new().expect("Failed to initialize session"); + let out_dir = env!("OUT_DIR").parse::().unwrap(); + let bv = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view"); + + #[derive(Default)] + struct Tag { + tags: usize, + } + + impl CustomDataNotification for Tag { + fn tag_added(&mut self, _view: &BinaryView, _tag_ref: &TagReference) { + self.tags += 1; + } + } + + let triggers = DataNotificationTriggers::default().tag_added(); + let tags_lock = Tag::default().register(&bv, triggers); + + let funcs = bv.functions(); + for (i, func) in funcs.iter().enumerate() { + let crash = bv.create_tag_type("Test", "🚧"); + func.add_tag(&crash, "Dummy tag", Some(i.try_into().unwrap()), true, None); + } + + let tags = tags_lock.unregister(); + + assert_eq!(funcs.len(), tags.tags); +}