Finish the `DataBuffer` implementation

2 files changed, 254 insertions, 12 deletions

diff --git a/rust/src/databuffer.rs b/rust/src/databuffer.rs
index ebe19848..c3d268cd 100644
--- a/rust/src/databuffer.rs
+++ b/rust/src/databuffer.rs
@@ -17,9 +17,10 @@
 use binaryninjacore_sys::*;
 
 use std::ffi::c_void;
-use std::ptr;
 use std::slice;
 
+use crate::string::BnString;
+
 pub struct DataBuffer(*mut BNDataBuffer);
 
 impl DataBuffer {
@@ -31,10 +32,6 @@ impl DataBuffer {
     }
 
     pub fn get_data(&self) -> &[u8] {
-        if self.0.is_null() {
-            // TODO : Change the default value and remove this
-            return &[];
-        }
         let buffer = unsafe { BNGetDataBufferContents(self.0) };
         if buffer.is_null() {
             &[]
@@ -43,6 +40,65 @@ impl DataBuffer {
         }
     }
 
+    pub fn get_data_at(&self, offset: usize) -> &[u8] {
+        let len = self.len();
+        if offset > len {
+            panic!();
+        }
+        let slice_len = len - offset;
+        let buffer = unsafe { BNGetDataBufferContentsAt(self.0, offset) };
+        if buffer.is_null() {
+            &[]
+        } else {
+            unsafe { slice::from_raw_parts(buffer as *const _, slice_len) }
+        }
+    }
+
+    /// Create a copy of a especified part of the data
+    pub fn get_slice(&self, start: usize, len: usize) -> Option<Self> {
+        if start + len > self.len() {
+            return None;
+        }
+        let ptr = unsafe { BNGetDataBufferSlice(self.0, start, len) };
+        (!ptr.is_null()).then(|| Self(ptr))
+    }
+
+    /// change the size of the allocated data, if new size is bigger data is
+    /// need to be initialized
+    pub unsafe fn set_len(&mut self, len: usize) {
+        unsafe { BNSetDataBufferLength(self.0, len) }
+    }
+
+    /// set the size to 0
+    pub fn clear(&self) {
+        unsafe { BNClearDataBuffer(self.0) }
+    }
+
+    /// Copy the contents of `src` into `dst`
+    pub fn assign(dst: &mut Self, src: &Self) {
+        unsafe { BNAssignDataBuffer(dst.0, src.0) }
+    }
+
+    /// Concat the contents of `src` into `dst`
+    pub fn append(dst: &mut Self, src: &Self) {
+        unsafe { BNAppendDataBuffer(dst.0, src.0) }
+    }
+
+    /// concat the contents of `data` into self
+    pub fn append_data(&self, data: &[u8]) {
+        unsafe { BNAppendDataBufferContents(self.0, data.as_ptr() as *const c_void, data.len()) }
+    }
+
+    /// Return the byte at `offset`
+    pub unsafe fn byte_at(&self, offset: usize) -> u8 {
+        unsafe { BNGetDataBufferByte(self.0, offset) }
+    }
+
+    /// Set the value of the byte at `offset`
+    pub unsafe fn set_byte_at(&mut self, offset: usize, byte: u8) {
+        unsafe { BNSetDataBufferByte(self.0, offset, byte) }
+    }
+
     pub fn set_data(&mut self, data: &[u8]) {
         unsafe {
             BNSetDataBufferContents(
@@ -53,12 +109,48 @@ impl DataBuffer {
         }
     }
 
+    pub fn to_escaped_string(&self, null_terminates: bool) -> BnString {
+        unsafe { BnString::from_raw(BNDataBufferToEscapedString(self.0, null_terminates)) }
+    }
+
+    pub fn from_escaped_string(value: &BnString) -> Self {
+        Self(unsafe { BNDecodeEscapedString(value.as_raw()) })
+    }
+
+    pub fn to_base64(&self) -> BnString {
+        unsafe { BnString::from_raw(BNDataBufferToBase64(self.0)) }
+    }
+
+    pub fn from_base64(value: &BnString) -> Self {
+        Self(unsafe { BNDecodeBase64(value.as_raw()) })
+    }
+
+    pub fn zlib_compress(&self) -> Self {
+        Self(unsafe { BNZlibCompress(self.0) })
+    }
+
+    pub fn zlib_decompress(&self) -> Self {
+        Self(unsafe { BNZlibDecompress(self.0) })
+    }
+
+    pub fn lzma_decompress(&self) -> Self {
+        Self(unsafe { BNLzmaDecompress(self.0) })
+    }
+
+    pub fn lzma2_decompress(&self) -> Self {
+        Self(unsafe { BNLzma2Decompress(self.0) })
+    }
+
+    pub fn xz_decompress(&self) -> Self {
+        Self(unsafe { BNXzDecompress(self.0) })
+    }
+
     pub fn len(&self) -> usize {
         unsafe { BNGetDataBufferLength(self.0) }
     }
 
     pub fn is_empty(&self) -> bool {
-        unsafe { BNGetDataBufferLength(self.0) == 0 }
+        self.len() == 0
     }
 
     pub fn new(data: &[u8]) -> Result<Self, ()> {
@@ -71,19 +163,16 @@ impl DataBuffer {
     }
 }
 
-// TODO : delete this
 impl Default for DataBuffer {
     fn default() -> Self {
-        DataBuffer::from_raw(ptr::null_mut())
+        Self(unsafe { BNCreateDataBuffer([].as_ptr() as *const c_void, 0) })
     }
 }
 
 impl Drop for DataBuffer {
     fn drop(&mut self) {
-        if !self.0.is_null() {
-            unsafe {
-                BNFreeDataBuffer(self.0);
-            }
+        unsafe {
+            BNFreeDataBuffer(self.0);
         }
     }
 }
@@ -93,3 +182,152 @@ impl Clone for DataBuffer {
         Self::from_raw(unsafe { BNDuplicateDataBuffer(self.0) })
     }
 }
+
+impl TryFrom<&[u8]> for DataBuffer {
+    type Error = ();
+
+    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
+        DataBuffer::new(value)
+    }
+}
+
+impl AsRef<[u8]> for DataBuffer {
+    fn as_ref(&self) -> &[u8] {
+        self.get_data()
+    }
+}
+
+impl std::borrow::Borrow<[u8]> for DataBuffer {
+    fn borrow(&self) -> &[u8] {
+        self.as_ref()
+    }
+}
+
+macro_rules! data_buffer_index {
+    ($range:ty, $output:ty) => {
+        impl std::ops::Index<$range> for DataBuffer {
+            type Output = $output;
+
+            fn index(&self, index: $range) -> &Self::Output {
+                &self.get_data()[index]
+            }
+        }
+    };
+}
+
+data_buffer_index!(usize, u8);
+data_buffer_index!(std::ops::Range<usize>, [u8]);
+data_buffer_index!(std::ops::RangeInclusive<usize>, [u8]);
+data_buffer_index!(std::ops::RangeTo<usize>, [u8]);
+data_buffer_index!(std::ops::RangeFull, [u8]);
+
+impl PartialEq for DataBuffer {
+    fn eq(&self, other: &Self) -> bool {
+        self.as_ref() == other.as_ref()
+    }
+}
+impl Eq for DataBuffer {}
+
+impl PartialOrd for DataBuffer {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.as_ref().cmp(other.as_ref()))
+    }
+}
+
+impl Ord for DataBuffer {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.as_ref().cmp(other.as_ref())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::DataBuffer;
+
+    const DUMMY_DATA_0: &[u8] = b"0123456789\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x09\xFF";
+    const DUMMY_DATA_1: &[u8] = b"qwertyuiopasdfghjkl\xE7zxcvbnm\x00\x01\x00";
+
+    #[test]
+    fn get_slice() {
+        let data = DataBuffer::new(DUMMY_DATA_0).unwrap();
+        let slice = data.get_slice(9, 10).unwrap();
+        assert_eq!(slice.get_data(), &DUMMY_DATA_0[9..19]);
+    }
+
+    #[test]
+    fn set_len_write() {
+        let mut data = DataBuffer::default();
+        assert_eq!(data.get_data(), &[]);
+        unsafe { data.set_len(DUMMY_DATA_0.len()) };
+        assert_eq!(data.len(), DUMMY_DATA_0.len());
+        let mut contents = DUMMY_DATA_0.to_vec();
+        data.set_data(&contents);
+        // modify the orinal contents, to make sure DataBuffer copied the data
+        // and is not using the original pointer
+        contents.as_mut_slice().fill(0x55);
+        drop(contents);
+        assert_eq!(data.get_data(), &DUMMY_DATA_0[..]);
+
+        // make sure the new len truncate the original data
+        unsafe { data.set_len(13) };
+        assert_eq!(data.get_data(), &DUMMY_DATA_0[..13]);
+
+        data.clear();
+        assert_eq!(data.get_data(), &[]);
+    }
+
+    #[test]
+    fn assign_append() {
+        let mut dst = DataBuffer::new(DUMMY_DATA_0).unwrap();
+        let mut src = DataBuffer::new(DUMMY_DATA_1).unwrap();
+        DataBuffer::assign(&mut dst, &src);
+
+        assert_eq!(dst.get_data(), DUMMY_DATA_1);
+        assert_eq!(src.get_data(), DUMMY_DATA_1);
+        // overwrite the src, to make sure that src is copied to dst, and not
+        // moved into it
+        src.set_data(DUMMY_DATA_0);
+        assert_eq!(dst.get_data(), DUMMY_DATA_1);
+        assert_eq!(src.get_data(), DUMMY_DATA_0);
+
+        DataBuffer::append(&mut dst, &src);
+        let result: Vec<_> = DUMMY_DATA_1.iter().chain(DUMMY_DATA_0).copied().collect();
+        assert_eq!(dst.get_data(), &result);
+
+        assert_eq!(src.get_data(), DUMMY_DATA_0);
+        src.set_data(DUMMY_DATA_1);
+        assert_eq!(src.get_data(), DUMMY_DATA_1);
+        assert_eq!(dst.get_data(), &result);
+    }
+
+    #[test]
+    fn to_from_formats() {
+        let data = DataBuffer::new(DUMMY_DATA_0).unwrap();
+        let escaped = data.to_escaped_string(false);
+        let unescaped = DataBuffer::from_escaped_string(&escaped);
+        drop(escaped);
+        let escaped_part = data.to_escaped_string(true);
+        let unescaped_part = DataBuffer::from_escaped_string(&escaped_part);
+        drop(escaped_part);
+
+        let part = &DUMMY_DATA_0[0..DUMMY_DATA_0
+            .iter()
+            .position(|x| *x == 0)
+            .unwrap_or(DUMMY_DATA_0.len())];
+        assert_eq!(data.get_data(), DUMMY_DATA_0);
+        assert_eq!(unescaped.get_data(), DUMMY_DATA_0);
+        assert_eq!(unescaped_part.get_data(), part);
+
+        let escaped = data.to_base64();
+        let unescaped = DataBuffer::from_base64(&escaped);
+        drop(escaped);
+        assert_eq!(data.get_data(), DUMMY_DATA_0);
+        assert_eq!(unescaped.get_data(), DUMMY_DATA_0);
+
+        let compressed = data.zlib_compress();
+        let decompressed = compressed.zlib_decompress();
+        drop(compressed);
+        assert_eq!(data.get_data(), DUMMY_DATA_0);
+        assert_eq!(decompressed.get_data(), DUMMY_DATA_0);
+    }
+}
diff --git a/rust/src/string.rs b/rust/src/string.rs
index 936e3033..bfb998d5 100644
--- a/rust/src/string.rs
+++ b/rust/src/string.rs
@@ -79,6 +79,10 @@ impl BnString {
         res
     }
 
+    pub(crate) fn as_raw(&self) -> &raw::c_char {
+        unsafe { &*self.raw }
+    }
+
     pub fn as_str(&self) -> &str {
         unsafe { CStr::from_ptr(self.raw).to_str().unwrap() }
     }