From 067148afd1a8f9153bb5bfe6b5361b83e9f89ab6 Mon Sep 17 00:00:00 2001
From: Mason Reed <mason@vector35.com>
Date: Wed, 7 May 2025 23:23:44 -0400
Subject: [Rust] Improve `FileAccessor`

- Add unit tests
- Expose read/write functionality
- Add some much needed documentation when passing to memory map
---
 rust/src/file_accessor.rs | 94 +++++++++++++++++++++++++++++++----------------
 1 file changed, 62 insertions(+), 32 deletions(-)

(limited to 'rust/src/file_accessor.rs')

diff --git a/rust/src/file_accessor.rs b/rust/src/file_accessor.rs
index 12d25ff0..2ba7938f 100644
--- a/rust/src/file_accessor.rs
+++ b/rust/src/file_accessor.rs
@@ -13,41 +13,36 @@
 // limitations under the License.
 
 use binaryninjacore_sys::BNFileAccessor;
-use std::io::{Read, Seek, SeekFrom, Write};
+use std::io::{ErrorKind, Read, Seek, SeekFrom, Write};
 use std::marker::PhantomData;
 use std::slice;
 
-pub struct FileAccessor<'a> {
-    pub(crate) api_object: BNFileAccessor,
-    _ref: PhantomData<&'a mut ()>,
+pub trait Accessor: Read + Write + Seek + Sized {}
+
+impl<T: Read + Write + Seek + Sized> Accessor for T {}
+
+pub struct FileAccessor<A: Accessor> {
+    pub(crate) raw: BNFileAccessor,
+    accessor: PhantomData<A>,
 }
 
-impl<'a> FileAccessor<'a> {
-    pub fn new<F>(f: &'a mut F) -> Self
-    where
-        F: 'a + Read + Write + Seek + Sized,
-    {
+impl<A: Accessor> FileAccessor<A> {
+    pub fn new(accessor: A) -> Self {
         use std::os::raw::c_void;
 
-        extern "C" fn cb_get_length<F>(ctxt: *mut c_void) -> u64
-        where
-            F: Read + Write + Seek + Sized,
-        {
-            let f = unsafe { &mut *(ctxt as *mut F) };
+        extern "C" fn cb_get_length<A: Accessor>(ctxt: *mut c_void) -> u64 {
+            let f = unsafe { &mut *(ctxt as *mut A) };
 
             f.seek(SeekFrom::End(0)).unwrap_or(0)
         }
 
-        extern "C" fn cb_read<F>(
+        extern "C" fn cb_read<A: Accessor>(
             ctxt: *mut c_void,
             dest: *mut c_void,
             offset: u64,
             len: usize,
-        ) -> usize
-        where
-            F: Read + Write + Seek + Sized,
-        {
-            let f = unsafe { &mut *(ctxt as *mut F) };
+        ) -> usize {
+            let f = unsafe { &mut *(ctxt as *mut A) };
             let dest = unsafe { slice::from_raw_parts_mut(dest as *mut u8, len) };
 
             if f.seek(SeekFrom::Start(offset)).is_err() {
@@ -58,16 +53,13 @@ impl<'a> FileAccessor<'a> {
             }
         }
 
-        extern "C" fn cb_write<F>(
+        extern "C" fn cb_write<A: Accessor>(
             ctxt: *mut c_void,
             offset: u64,
             src: *const c_void,
             len: usize,
-        ) -> usize
-        where
-            F: Read + Write + Seek + Sized,
-        {
-            let f = unsafe { &mut *(ctxt as *mut F) };
+        ) -> usize {
+            let f = unsafe { &mut *(ctxt as *mut A) };
             let src = unsafe { slice::from_raw_parts(src as *const u8, len) };
 
             if f.seek(SeekFrom::Start(offset)).is_err() {
@@ -77,14 +69,52 @@ impl<'a> FileAccessor<'a> {
             }
         }
 
+        let boxed_accessor = Box::new(accessor);
+        let leaked_accessor = Box::leak(boxed_accessor);
+
         Self {
-            api_object: BNFileAccessor {
-                context: f as *mut F as *mut _,
-                getLength: Some(cb_get_length::<F>),
-                read: Some(cb_read::<F>),
-                write: Some(cb_write::<F>),
+            raw: BNFileAccessor {
+                context: leaked_accessor as *mut A as *mut _,
+                getLength: Some(cb_get_length::<A>),
+                read: Some(cb_read::<A>),
+                write: Some(cb_write::<A>),
             },
-            _ref: PhantomData,
+            accessor: PhantomData,
+        }
+    }
+
+    pub fn read(&self, addr: u64, len: usize) -> Result<Vec<u8>, ErrorKind> {
+        let cb_read = self.raw.read.unwrap();
+        let mut buf = vec![0; len];
+        let read_len = unsafe { cb_read(self.raw.context, buf.as_mut_ptr() as *mut _, addr, len) };
+        if read_len != len {
+            return Err(ErrorKind::UnexpectedEof);
+        }
+        Ok(buf)
+    }
+
+    pub fn write(&self, addr: u64, data: &[u8]) -> usize {
+        let cb_write = self.raw.write.unwrap();
+        unsafe {
+            cb_write(
+                self.raw.context,
+                addr,
+                data.as_ptr() as *const _,
+                data.len(),
+            )
+        }
+    }
+
+    pub fn length(&self) -> u64 {
+        let cb_get_length = self.raw.getLength.unwrap();
+        unsafe { cb_get_length(self.raw.context) }
+    }
+}
+
+impl<A: Accessor> Drop for FileAccessor<A> {
+    fn drop(&mut self) {
+        unsafe {
+            let _ = Box::from_raw(self.raw.context as *mut A);
         }
     }
 }
-- 
cgit v1.3.1