diff options
Diffstat (limited to 'rust/src/database')
| -rw-r--r-- | rust/src/database/kvs.rs | 135 | ||||
| -rw-r--r-- | rust/src/database/snapshot.rs | 277 | ||||
| -rw-r--r-- | rust/src/database/undo.rs | 168 |
3 files changed, 580 insertions, 0 deletions
diff --git a/rust/src/database/kvs.rs b/rust/src/database/kvs.rs new file mode 100644 index 00000000..4b77bbdb --- /dev/null +++ b/rust/src/database/kvs.rs @@ -0,0 +1,135 @@ +use crate::data_buffer::DataBuffer; +use crate::rc::{Array, Ref, RefCountable}; +use crate::string::{BnStrCompatible, BnString}; +use binaryninjacore_sys::{ + BNBeginKeyValueStoreNamespace, BNEndKeyValueStoreNamespace, BNFreeKeyValueStore, + BNGetKeyValueStoreBuffer, BNGetKeyValueStoreDataSize, BNGetKeyValueStoreKeys, + BNGetKeyValueStoreNamespaceSize, BNGetKeyValueStoreSerializedData, BNGetKeyValueStoreValueSize, + BNGetKeyValueStoreValueStorageSize, BNIsKeyValueStoreEmpty, BNKeyValueStore, + BNNewKeyValueStoreReference, BNSetKeyValueStoreBuffer, +}; +use std::collections::HashMap; +use std::ffi::c_char; +use std::fmt::Debug; +use std::ptr::NonNull; + +#[repr(transparent)] +pub struct KeyValueStore { + pub(crate) handle: NonNull<BNKeyValueStore>, +} + +impl KeyValueStore { + pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNKeyValueStore>) -> Ref<Self> { + Ref::new(Self { handle }) + } + + pub fn to_hashmap(&self) -> HashMap<String, DataBuffer> { + let mut hashmap = HashMap::with_capacity(self.keys().len()); + for key in self.keys().iter() { + if let Some(value) = self.value(key) { + hashmap.insert(key.to_string(), value); + } + } + hashmap + } + + /// Get a list of all keys stored in the kvs + pub fn keys(&self) -> Array<BnString> { + let mut count = 0; + let result = unsafe { BNGetKeyValueStoreKeys(self.handle.as_ptr(), &mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + /// Get the value for a single key + pub fn value<S: BnStrCompatible>(&self, key: S) -> Option<DataBuffer> { + let key_raw = key.into_bytes_with_nul(); + let key_ptr = key_raw.as_ref().as_ptr() as *const c_char; + let result = unsafe { BNGetKeyValueStoreBuffer(self.handle.as_ptr(), key_ptr) }; + NonNull::new(result).map(|_| DataBuffer::from_raw(result)) + } + + /// Set the value for a single key + pub fn set_value<S: BnStrCompatible>(&self, key: S, value: &DataBuffer) -> bool { + let key_raw = key.into_bytes_with_nul(); + let key_ptr = key_raw.as_ref().as_ptr() as *const c_char; + unsafe { BNSetKeyValueStoreBuffer(self.handle.as_ptr(), key_ptr, value.as_raw()) } + } + + /// Get the stored representation of the kvs + pub fn serialized_data(&self) -> DataBuffer { + let result = unsafe { BNGetKeyValueStoreSerializedData(self.handle.as_ptr()) }; + assert!(!result.is_null()); + DataBuffer::from_raw(result) + } + + /// Begin storing new keys into a namespace + pub fn begin_namespace<S: BnStrCompatible>(&self, name: S) { + let name_raw = name.into_bytes_with_nul(); + let name_ptr = name_raw.as_ref().as_ptr() as *const c_char; + unsafe { BNBeginKeyValueStoreNamespace(self.handle.as_ptr(), name_ptr) } + } + + /// End storing new keys into a namespace + pub fn end_namespace(&self) { + unsafe { BNEndKeyValueStoreNamespace(self.handle.as_ptr()) } + } + + /// If the kvs is empty + pub fn is_empty(&self) -> bool { + unsafe { BNIsKeyValueStoreEmpty(self.handle.as_ptr()) } + } + + /// Number of values in the kvs + pub fn value_size(&self) -> usize { + unsafe { BNGetKeyValueStoreValueSize(self.handle.as_ptr()) } + } + + /// Length of serialized data + pub fn data_size(&self) -> usize { + unsafe { BNGetKeyValueStoreDataSize(self.handle.as_ptr()) } + } + + /// Size of all data in storage + pub fn value_storage_size(&self) -> usize { + unsafe { BNGetKeyValueStoreValueStorageSize(self.handle.as_ptr()) } + } + + /// Number of namespaces pushed with begin_namespace + pub fn namespace_size(&self) -> usize { + unsafe { BNGetKeyValueStoreNamespaceSize(self.handle.as_ptr()) } + } +} + +impl ToOwned for KeyValueStore { + type Owned = Ref<Self>; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for KeyValueStore { + unsafe fn inc_ref(handle: &Self) -> Ref<Self> { + Ref::new(Self { + handle: NonNull::new(BNNewKeyValueStoreReference(handle.handle.as_ptr())).unwrap(), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeKeyValueStore(handle.handle.as_ptr()); + } +} + +impl Debug for KeyValueStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KeyValueStore") + .field("keys", &self.keys().to_vec()) + .field("is_empty", &self.is_empty()) + .field("value_size", &self.value_size()) + .field("data_size", &self.data_size()) + .field("value_storage_size", &self.value_storage_size()) + .field("namespace_size", &self.namespace_size()) + .finish() + } +} diff --git a/rust/src/database/snapshot.rs b/rust/src/database/snapshot.rs new file mode 100644 index 00000000..35b0dc46 --- /dev/null +++ b/rust/src/database/snapshot.rs @@ -0,0 +1,277 @@ +use crate::data_buffer::DataBuffer; +use crate::database::kvs::KeyValueStore; +use crate::database::undo::UndoEntry; +use crate::database::Database; +use crate::progress::ProgressCallback; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use crate::string::{BnStrCompatible, BnString}; +use binaryninjacore_sys::{ + BNCollaborationFreeSnapshotIdList, BNFreeSnapshot, BNFreeSnapshotList, BNGetSnapshotChildren, + BNGetSnapshotDatabase, BNGetSnapshotFileContents, BNGetSnapshotFileContentsHash, + BNGetSnapshotFirstParent, BNGetSnapshotId, BNGetSnapshotName, BNGetSnapshotParents, + BNGetSnapshotUndoData, BNGetSnapshotUndoEntries, BNGetSnapshotUndoEntriesWithProgress, + BNIsSnapshotAutoSave, BNNewSnapshotReference, BNReadSnapshotData, + BNReadSnapshotDataWithProgress, BNSetSnapshotName, BNSnapshot, BNSnapshotHasAncestor, + BNSnapshotHasContents, BNSnapshotHasUndo, BNSnapshotStoreData, +}; +use std::ffi::{c_char, c_void}; +use std::fmt; +use std::fmt::{Debug, Display, Formatter}; +use std::ptr::NonNull; + +pub struct Snapshot { + pub(crate) handle: NonNull<BNSnapshot>, +} + +impl Snapshot { + pub(crate) unsafe fn from_raw(handle: NonNull<BNSnapshot>) -> Self { + Self { handle } + } + + pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNSnapshot>) -> Ref<Self> { + Ref::new(Self { handle }) + } + + /// Get the owning database + pub fn database(&self) -> Database { + unsafe { + Database::from_raw(NonNull::new(BNGetSnapshotDatabase(self.handle.as_ptr())).unwrap()) + } + } + + /// Get the numerical id + pub fn id(&self) -> SnapshotId { + SnapshotId(unsafe { BNGetSnapshotId(self.handle.as_ptr()) }) + } + + /// Get the displayed snapshot name + pub fn name(&self) -> BnString { + unsafe { BnString::from_raw(BNGetSnapshotName(self.handle.as_ptr())) } + } + + /// Set the displayed snapshot name + pub fn set_name<S: BnStrCompatible>(&self, value: S) { + let value_raw = value.into_bytes_with_nul(); + let value_ptr = value_raw.as_ref().as_ptr() as *const c_char; + unsafe { BNSetSnapshotName(self.handle.as_ptr(), value_ptr) } + } + + /// If the snapshot was the result of an auto-save + pub fn is_auto_save(&self) -> bool { + unsafe { BNIsSnapshotAutoSave(self.handle.as_ptr()) } + } + + /// If the snapshot has contents, and has not been trimmed + pub fn has_contents(&self) -> bool { + unsafe { BNSnapshotHasContents(self.handle.as_ptr()) } + } + + /// If the snapshot has undo data + pub fn has_undo(&self) -> bool { + unsafe { BNSnapshotHasUndo(self.handle.as_ptr()) } + } + + /// Get the first parent of the snapshot, or None if it has no parents + pub fn first_parent(&self) -> Option<Snapshot> { + let result = unsafe { BNGetSnapshotFirstParent(self.handle.as_ptr()) }; + NonNull::new(result).map(|s| unsafe { Snapshot::from_raw(s) }) + } + + /// Get a list of all parent snapshots of the snapshot + pub fn parents(&self) -> Array<Snapshot> { + let mut count = 0; + let result = unsafe { BNGetSnapshotParents(self.handle.as_ptr(), &mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + /// Get a list of all child snapshots of the snapshot + pub fn children(&self) -> Array<Snapshot> { + let mut count = 0; + let result = unsafe { BNGetSnapshotChildren(self.handle.as_ptr(), &mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + /// Get a buffer of the raw data at the time of the snapshot + pub fn file_contents(&self) -> Option<DataBuffer> { + self.has_contents().then(|| unsafe { + let result = BNGetSnapshotFileContents(self.handle.as_ptr()); + assert!(!result.is_null()); + DataBuffer::from_raw(result) + }) + } + + /// Get a hash of the data at the time of the snapshot + pub fn file_contents_hash(&self) -> Option<DataBuffer> { + self.has_contents().then(|| unsafe { + let result = BNGetSnapshotFileContentsHash(self.handle.as_ptr()); + assert!(!result.is_null()); + DataBuffer::from_raw(result) + }) + } + + /// Get a list of undo entries at the time of the snapshot + pub fn undo_entries(&self) -> Array<UndoEntry> { + assert!(self.has_undo()); + let mut count = 0; + let result = unsafe { BNGetSnapshotUndoEntries(self.handle.as_ptr(), &mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + pub fn undo_entries_with_progress<P: ProgressCallback>( + &self, + mut progress: P, + ) -> Array<UndoEntry> { + assert!(self.has_undo()); + let mut count = 0; + + let result = unsafe { + BNGetSnapshotUndoEntriesWithProgress( + self.handle.as_ptr(), + &mut progress as *mut P as *mut c_void, + Some(P::cb_progress_callback), + &mut count, + ) + }; + + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + /// Get the backing kvs data with snapshot fields + pub fn read_data(&self) -> Ref<KeyValueStore> { + let result = unsafe { BNReadSnapshotData(self.handle.as_ptr()) }; + unsafe { KeyValueStore::ref_from_raw(NonNull::new(result).unwrap()) } + } + + pub fn read_data_with_progress<P: ProgressCallback>( + &self, + mut progress: P, + ) -> Ref<KeyValueStore> { + let result = unsafe { + BNReadSnapshotDataWithProgress( + self.handle.as_ptr(), + &mut progress as *mut P as *mut c_void, + Some(P::cb_progress_callback), + ) + }; + + unsafe { KeyValueStore::ref_from_raw(NonNull::new(result).unwrap()) } + } + + pub fn undo_data(&self) -> DataBuffer { + let result = unsafe { BNGetSnapshotUndoData(self.handle.as_ptr()) }; + assert!(!result.is_null()); + DataBuffer::from_raw(result) + } + + pub fn store_data(&self, data: &KeyValueStore) -> bool { + unsafe { + BNSnapshotStoreData( + self.handle.as_ptr(), + data.handle.as_ptr(), + std::ptr::null_mut(), + None, + ) + } + } + + pub fn store_data_with_progress<P: ProgressCallback>( + &self, + data: &KeyValueStore, + mut progress: P, + ) -> bool { + unsafe { + BNSnapshotStoreData( + self.handle.as_ptr(), + data.handle.as_ptr(), + &mut progress as *mut P as *mut c_void, + Some(P::cb_progress_callback), + ) + } + } + + /// Determine if this snapshot has another as an ancestor + pub fn has_ancestor(self, other: &Snapshot) -> bool { + unsafe { BNSnapshotHasAncestor(self.handle.as_ptr(), other.handle.as_ptr()) } + } +} + +impl Debug for Snapshot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Snapshot") + .field("id", &self.id()) + .field("name", &self.name()) + .field("is_auto_save", &self.is_auto_save()) + .field("has_contents", &self.has_contents()) + .field("has_undo", &self.has_undo()) + // TODO: This might be too much. + .field("children", &self.children().to_vec()) + .finish() + } +} + +impl ToOwned for Snapshot { + type Owned = Ref<Self>; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for Snapshot { + unsafe fn inc_ref(handle: &Self) -> Ref<Self> { + Ref::new(Self { + handle: NonNull::new(BNNewSnapshotReference(handle.handle.as_ptr())).unwrap(), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeSnapshot(handle.handle.as_ptr()); + } +} + +impl CoreArrayProvider for Snapshot { + type Raw = *mut BNSnapshot; + type Context = (); + type Wrapped<'a> = Guard<'a, Snapshot>; +} + +unsafe impl CoreArrayProviderInner for Snapshot { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeSnapshotList(raw, count); + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + let raw_ptr = NonNull::new(*raw).unwrap(); + Guard::new(Self::from_raw(raw_ptr), context) + } +} + +#[repr(transparent)] +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct SnapshotId(pub i64); + +impl Display for SnapshotId { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.write_fmt(format_args!("{}", self.0)) + } +} + +impl CoreArrayProvider for SnapshotId { + type Raw = i64; + type Context = (); + type Wrapped<'a> = SnapshotId; +} + +unsafe impl CoreArrayProviderInner for SnapshotId { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNCollaborationFreeSnapshotIdList(raw, count) + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped<'a> { + SnapshotId(*raw) + } +} diff --git a/rust/src/database/undo.rs b/rust/src/database/undo.rs new file mode 100644 index 00000000..727a777b --- /dev/null +++ b/rust/src/database/undo.rs @@ -0,0 +1,168 @@ +use crate::disassembly::InstructionTextToken; +use crate::rc::{Array, CoreArrayProvider, CoreArrayProviderInner, Guard, Ref, RefCountable}; +use crate::string::BnString; +use binaryninjacore_sys::{ + BNFreeUndoAction, BNFreeUndoActionList, BNFreeUndoEntry, BNFreeUndoEntryList, + BNNewUndoActionReference, BNNewUndoEntryReference, BNUndoAction, BNUndoActionGetSummary, + BNUndoActionGetSummaryText, BNUndoEntry, BNUndoEntryGetActions, BNUndoEntryGetId, + BNUndoEntryGetTimestamp, +}; +use std::fmt::Debug; +use std::ptr::NonNull; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +#[repr(transparent)] +pub struct UndoEntry { + handle: NonNull<BNUndoEntry>, +} + +impl UndoEntry { + pub(crate) unsafe fn from_raw(handle: NonNull<BNUndoEntry>) -> Self { + Self { handle } + } + + #[allow(dead_code)] + pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNUndoEntry>) -> Ref<Self> { + Ref::new(Self { handle }) + } + + pub fn id(&self) -> BnString { + let result = unsafe { BNUndoEntryGetId(self.handle.as_ptr()) }; + assert!(!result.is_null()); + unsafe { BnString::from_raw(result) } + } + + pub fn actions(&self) -> Array<UndoAction> { + let mut count = 0; + let result = unsafe { BNUndoEntryGetActions(self.handle.as_ptr(), &mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + pub fn time(&self) -> SystemTime { + let m = Duration::from_secs(unsafe { BNUndoEntryGetTimestamp(self.handle.as_ptr()) }); + UNIX_EPOCH + m + } +} + +impl Debug for UndoEntry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UndoEntry") + .field("id", &self.id()) + .field("time", &self.time()) + .field("actions", &self.actions().to_vec()) + .finish() + } +} + +impl ToOwned for UndoEntry { + type Owned = Ref<Self>; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for UndoEntry { + unsafe fn inc_ref(handle: &Self) -> Ref<Self> { + Ref::new(Self { + handle: NonNull::new(BNNewUndoEntryReference(handle.handle.as_ptr())).unwrap(), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeUndoEntry(handle.handle.as_ptr()); + } +} + +impl CoreArrayProvider for UndoEntry { + type Raw = *mut BNUndoEntry; + type Context = (); + type Wrapped<'a> = Guard<'a, UndoEntry>; +} + +unsafe impl CoreArrayProviderInner for UndoEntry { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeUndoEntryList(raw, count); + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + let raw_ptr = NonNull::new(*raw).unwrap(); + Guard::new(Self::from_raw(raw_ptr), context) + } +} + +#[repr(transparent)] +pub struct UndoAction { + handle: NonNull<BNUndoAction>, +} + +impl UndoAction { + pub(crate) unsafe fn from_raw(handle: NonNull<BNUndoAction>) -> Self { + Self { handle } + } + + #[allow(dead_code)] + pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNUndoAction>) -> Ref<Self> { + Ref::new(Self { handle }) + } + + pub fn summary(&self) -> Array<InstructionTextToken> { + let mut count = 0; + let result = unsafe { BNUndoActionGetSummary(self.handle.as_ptr(), &mut count) }; + assert!(!result.is_null()); + unsafe { Array::new(result, count, ()) } + } + + /// Gets the [`UndoAction`] summary as text rather than [`InstructionTextToken`]'s. + pub fn summary_as_string(&self) -> BnString { + let result = unsafe { BNUndoActionGetSummaryText(self.handle.as_ptr()) }; + assert!(!result.is_null()); + unsafe { BnString::from_raw(result) } + } +} + +impl Debug for UndoAction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("UndoAction") + .field("summary", &self.summary_as_string()) + .finish() + } +} + +impl ToOwned for UndoAction { + type Owned = Ref<Self>; + + fn to_owned(&self) -> Self::Owned { + unsafe { RefCountable::inc_ref(self) } + } +} + +unsafe impl RefCountable for UndoAction { + unsafe fn inc_ref(handle: &Self) -> Ref<Self> { + Ref::new(Self { + handle: NonNull::new(BNNewUndoActionReference(handle.handle.as_ptr())).unwrap(), + }) + } + + unsafe fn dec_ref(handle: &Self) { + BNFreeUndoAction(handle.handle.as_ptr()); + } +} + +impl CoreArrayProvider for UndoAction { + type Raw = *mut BNUndoAction; + type Context = (); + type Wrapped<'a> = Guard<'a, UndoAction>; +} + +unsafe impl CoreArrayProviderInner for UndoAction { + unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) { + BNFreeUndoActionList(raw, count); + } + + unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> { + let raw_ptr = NonNull::new(*raw).unwrap(); + Guard::new(Self::from_raw(raw_ptr), context) + } +} |
