binaryninja/database/
kvs.rs

1use crate::data_buffer::DataBuffer;
2use crate::rc::{Array, Ref, RefCountable};
3use crate::string::{BnString, IntoCStr};
4use binaryninjacore_sys::{
5    BNBeginKeyValueStoreNamespace, BNEndKeyValueStoreNamespace, BNFreeKeyValueStore,
6    BNGetKeyValueStoreBuffer, BNGetKeyValueStoreDataSize, BNGetKeyValueStoreKeys,
7    BNGetKeyValueStoreNamespaceSize, BNGetKeyValueStoreSerializedData, BNGetKeyValueStoreValueSize,
8    BNGetKeyValueStoreValueStorageSize, BNIsKeyValueStoreEmpty, BNKeyValueStore,
9    BNNewKeyValueStoreReference, BNSetKeyValueStoreBuffer,
10};
11use std::collections::HashMap;
12use std::fmt::Debug;
13use std::ptr::NonNull;
14
15#[repr(transparent)]
16pub struct KeyValueStore {
17    pub(crate) handle: NonNull<BNKeyValueStore>,
18}
19
20impl KeyValueStore {
21    pub(crate) unsafe fn ref_from_raw(handle: NonNull<BNKeyValueStore>) -> Ref<Self> {
22        Ref::new(Self { handle })
23    }
24
25    pub fn to_hashmap(&self) -> HashMap<String, DataBuffer> {
26        let mut hashmap = HashMap::with_capacity(self.keys().len());
27        for key in self.keys().iter() {
28            if let Some(value) = self.value(key) {
29                hashmap.insert(key.to_string(), value);
30            }
31        }
32        hashmap
33    }
34
35    /// Get a list of all keys stored in the kvs
36    pub fn keys(&self) -> Array<BnString> {
37        let mut count = 0;
38        let result = unsafe { BNGetKeyValueStoreKeys(self.handle.as_ptr(), &mut count) };
39        assert!(!result.is_null());
40        unsafe { Array::new(result, count, ()) }
41    }
42
43    /// Get the value for a single key
44    pub fn value(&self, key: &str) -> Option<DataBuffer> {
45        let key_raw = key.to_cstr();
46        let key_ptr = key_raw.as_ptr();
47        let result = unsafe { BNGetKeyValueStoreBuffer(self.handle.as_ptr(), key_ptr) };
48        NonNull::new(result).map(|_| DataBuffer::from_raw(result))
49    }
50
51    /// Set the value for a single key
52    pub fn set_value(&self, key: &str, value: &DataBuffer) -> bool {
53        let key_raw = key.to_cstr();
54        let key_ptr = key_raw.as_ptr();
55        unsafe { BNSetKeyValueStoreBuffer(self.handle.as_ptr(), key_ptr, value.as_raw()) }
56    }
57
58    /// Get the stored representation of the kvs
59    pub fn serialized_data(&self) -> DataBuffer {
60        let result = unsafe { BNGetKeyValueStoreSerializedData(self.handle.as_ptr()) };
61        assert!(!result.is_null());
62        DataBuffer::from_raw(result)
63    }
64
65    /// Begin storing new keys into a namespace
66    pub fn begin_namespace(&self, name: &str) {
67        let name_raw = name.to_cstr();
68        let name_ptr = name_raw.as_ptr();
69        unsafe { BNBeginKeyValueStoreNamespace(self.handle.as_ptr(), name_ptr) }
70    }
71
72    /// End storing new keys into a namespace
73    pub fn end_namespace(&self) {
74        unsafe { BNEndKeyValueStoreNamespace(self.handle.as_ptr()) }
75    }
76
77    /// If the kvs is empty
78    pub fn is_empty(&self) -> bool {
79        unsafe { BNIsKeyValueStoreEmpty(self.handle.as_ptr()) }
80    }
81
82    /// Number of values in the kvs
83    pub fn value_size(&self) -> usize {
84        unsafe { BNGetKeyValueStoreValueSize(self.handle.as_ptr()) }
85    }
86
87    /// Length of serialized data
88    pub fn data_size(&self) -> usize {
89        unsafe { BNGetKeyValueStoreDataSize(self.handle.as_ptr()) }
90    }
91
92    /// Size of all data in storage
93    pub fn value_storage_size(&self) -> usize {
94        unsafe { BNGetKeyValueStoreValueStorageSize(self.handle.as_ptr()) }
95    }
96
97    /// Number of namespaces pushed with begin_namespace
98    pub fn namespace_size(&self) -> usize {
99        unsafe { BNGetKeyValueStoreNamespaceSize(self.handle.as_ptr()) }
100    }
101}
102
103impl ToOwned for KeyValueStore {
104    type Owned = Ref<Self>;
105
106    fn to_owned(&self) -> Self::Owned {
107        unsafe { RefCountable::inc_ref(self) }
108    }
109}
110
111unsafe impl RefCountable for KeyValueStore {
112    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
113        Ref::new(Self {
114            handle: NonNull::new(BNNewKeyValueStoreReference(handle.handle.as_ptr())).unwrap(),
115        })
116    }
117
118    unsafe fn dec_ref(handle: &Self) {
119        BNFreeKeyValueStore(handle.handle.as_ptr());
120    }
121}
122
123impl Debug for KeyValueStore {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        f.debug_struct("KeyValueStore")
126            .field("keys", &self.keys().to_vec())
127            .field("is_empty", &self.is_empty())
128            .field("value_size", &self.value_size())
129            .field("data_size", &self.data_size())
130            .field("value_storage_size", &self.value_storage_size())
131            .field("namespace_size", &self.namespace_size())
132            .finish()
133    }
134}