[Rust API] DownloadProvider

2 files changed, 314 insertions, 0 deletions

diff --git a/rust/src/downloadprovider.rs b/rust/src/downloadprovider.rs
new file mode 100644
index 00000000..65a11dd0
--- /dev/null
+++ b/rust/src/downloadprovider.rs
@@ -0,0 +1,313 @@
+use crate::rc::{
+    Array, CoreArrayProvider, CoreArrayWrapper, CoreOwnedArrayProvider, Ref, RefCountable,
+};
+use crate::settings::Settings;
+use crate::string::{BnStr, BnStrCompatible, BnString};
+use binaryninjacore_sys::*;
+use std::collections::HashMap;
+use std::ffi::c_void;
+use std::os::raw::c_char;
+use std::ptr::null_mut;
+use std::slice;
+
+pub struct DownloadProvider {
+    handle: *mut BNDownloadProvider,
+}
+
+impl DownloadProvider {
+    pub fn get<S: BnStrCompatible>(name: S) -> Option<DownloadProvider> {
+        let result = unsafe {
+            BNGetDownloadProviderByName(
+                name.into_bytes_with_nul().as_ref().as_ptr() as *const c_char
+            )
+        };
+        if result.is_null() {
+            return None;
+        }
+        Some(DownloadProvider { handle: result })
+    }
+
+    pub fn list() -> Result<Array<DownloadProvider>, ()> {
+        let mut count = 0;
+        let list: *mut *mut BNDownloadProvider = unsafe { BNGetDownloadProviderList(&mut count) };
+
+        if list.is_null() {
+            return Err(());
+        }
+
+        Ok(unsafe { Array::new(list, count, ()) })
+    }
+
+    pub fn default() -> Result<DownloadProvider, ()> {
+        let s = Settings::new("");
+        let dp_name = s.get_string("network.downloadProviderName", None, None);
+        Self::get(dp_name).ok_or(())
+    }
+
+    pub(crate) fn from_raw(handle: *mut BNDownloadProvider) -> DownloadProvider {
+        Self { handle }
+    }
+
+    pub fn create_instance(&self) -> Result<Ref<DownloadInstance>, ()> {
+        let result: *mut BNDownloadInstance =
+            unsafe { BNCreateDownloadProviderInstance(self.handle) };
+        if result.is_null() {
+            return Err(());
+        }
+
+        Ok(unsafe { DownloadInstance::ref_from_raw(result) })
+    }
+}
+
+impl CoreArrayProvider for DownloadProvider {
+    type Raw = *mut BNDownloadProvider;
+    type Context = ();
+}
+
+unsafe impl CoreOwnedArrayProvider for DownloadProvider {
+    unsafe fn free(raw: *mut Self::Raw, _count: usize, _context: &Self::Context) {
+        BNFreeDownloadProviderList(raw);
+    }
+}
+
+unsafe impl<'a> CoreArrayWrapper<'a> for DownloadProvider {
+    type Wrapped = DownloadProvider;
+
+    unsafe fn wrap_raw(raw: &'a Self::Raw, _context: &'a Self::Context) -> Self::Wrapped {
+        DownloadProvider::from_raw(*raw)
+    }
+}
+
+impl AsRef<DownloadProvider> for DownloadProvider {
+    fn as_ref(&self) -> &Self {
+        self
+    }
+}
+
+pub struct DownloadInstanceOutputCallbacks {
+    pub write: Option<Box<dyn FnMut(&[u8]) -> usize>>,
+    pub progress: Option<Box<dyn FnMut(u64, u64) -> bool>>,
+}
+
+pub struct DownloadInstanceInputOutputCallbacks {
+    pub read: Option<Box<dyn FnMut(&mut [u8]) -> Option<isize>>>,
+    pub write: Option<Box<dyn FnMut(&[u8]) -> usize>>,
+    pub progress: Option<Box<dyn FnMut(u64, u64) -> bool>>,
+}
+
+pub struct DownloadResponse {
+    pub status_code: u16,
+    pub headers: HashMap<String, String>,
+}
+
+pub struct DownloadInstance {
+    handle: *mut BNDownloadInstance,
+}
+
+impl DownloadInstance {
+    pub(crate) unsafe fn from_raw(handle: *mut BNDownloadInstance) -> Self {
+        debug_assert!(!handle.is_null());
+
+        Self { handle }
+    }
+
+    pub(crate) unsafe fn ref_from_raw(handle: *mut BNDownloadInstance) -> Ref<Self> {
+        Ref::new(Self::from_raw(handle))
+    }
+
+    fn get_error(&self) -> BnString {
+        let err: *mut c_char = unsafe { BNGetErrorForDownloadInstance(self.handle) };
+        unsafe { BnString::from_raw(err) }
+    }
+
+    unsafe extern "C" fn o_write_callback(data: *mut u8, len: u64, ctxt: *mut c_void) -> u64 {
+        let callbacks = ctxt as *mut DownloadInstanceOutputCallbacks;
+        if let Some(func) = &mut (*callbacks).write {
+            let slice = slice::from_raw_parts(data, len as usize);
+            let result = (func)(slice);
+            result as u64
+        } else {
+            0u64
+        }
+    }
+
+    unsafe extern "C" fn o_progress_callback(ctxt: *mut c_void, progress: u64, total: u64) -> bool {
+        let callbacks = ctxt as *mut DownloadInstanceOutputCallbacks;
+        if let Some(func) = &mut (*callbacks).progress {
+            (func)(progress, total)
+        } else {
+            true
+        }
+    }
+
+    pub fn perform_request<S: BnStrCompatible>(
+        &mut self,
+        url: S,
+        callbacks: DownloadInstanceOutputCallbacks,
+    ) -> Result<(), BnString> {
+        let callbacks = Box::into_raw(Box::new(callbacks));
+        let mut cbs = BNDownloadInstanceOutputCallbacks {
+            writeCallback: Some(Self::o_write_callback),
+            writeContext: callbacks as *mut c_void,
+            progressCallback: Some(Self::o_progress_callback),
+            progressContext: callbacks as *mut c_void,
+        };
+
+        let result = unsafe {
+            BNPerformDownloadRequest(
+                self.handle,
+                url.into_bytes_with_nul().as_ref().as_ptr() as *const c_char,
+                &mut cbs as *mut BNDownloadInstanceOutputCallbacks,
+            )
+        };
+
+        // Drop it
+        unsafe { Box::from_raw(callbacks) };
+        if result < 0 {
+            return Err(self.get_error());
+        }
+
+        return Ok(());
+    }
+
+    unsafe extern "C" fn i_read_callback(data: *mut u8, len: u64, ctxt: *mut c_void) -> i64 {
+        let callbacks = ctxt as *mut DownloadInstanceInputOutputCallbacks;
+        if let Some(func) = &mut (*callbacks).read {
+            let slice = slice::from_raw_parts_mut(data, len as usize);
+            let result = (func)(slice);
+            if let Some(count) = result {
+                count as i64
+            } else {
+                -1
+            }
+        } else {
+            0
+        }
+    }
+
+    unsafe extern "C" fn i_write_callback(data: *mut u8, len: u64, ctxt: *mut c_void) -> u64 {
+        let callbacks = ctxt as *mut DownloadInstanceInputOutputCallbacks;
+        if let Some(func) = &mut (*callbacks).write {
+            let slice = slice::from_raw_parts(data, len as usize);
+            let result = (func)(slice);
+            result as u64
+        } else {
+            0
+        }
+    }
+
+    unsafe extern "C" fn i_progress_callback(ctxt: *mut c_void, progress: u64, total: u64) -> bool {
+        let callbacks = ctxt as *mut DownloadInstanceInputOutputCallbacks;
+        if let Some(func) = &mut (*callbacks).progress {
+            (func)(progress, total)
+        } else {
+            true
+        }
+    }
+
+    pub fn perform_custom_request<
+        M: BnStrCompatible,
+        U: BnStrCompatible,
+        HK: BnStrCompatible,
+        HV: BnStrCompatible,
+        I: IntoIterator<Item = (HK, HV)>,
+    >(
+        &mut self,
+        method: M,
+        url: U,
+        headers: I,
+        callbacks: DownloadInstanceInputOutputCallbacks,
+    ) -> Result<DownloadResponse, BnString> {
+        let mut header_keys = vec![];
+        let mut header_values = vec![];
+        for (key, value) in headers {
+            header_keys.push(key.into_bytes_with_nul());
+            header_values.push(value.into_bytes_with_nul());
+        }
+
+        let mut header_key_ptrs = vec![];
+        let mut header_value_ptrs = vec![];
+
+        for (key, value) in header_keys.iter().zip(header_values.iter()) {
+            header_key_ptrs.push(key.as_ref().as_ptr() as *const c_char);
+            header_value_ptrs.push(value.as_ref().as_ptr() as *const c_char);
+        }
+
+        let callbacks = Box::into_raw(Box::new(callbacks));
+        let mut cbs = BNDownloadInstanceInputOutputCallbacks {
+            readCallback: Some(Self::i_read_callback),
+            readContext: callbacks as *mut c_void,
+            writeCallback: Some(Self::i_write_callback),
+            writeContext: callbacks as *mut c_void,
+            progressCallback: Some(Self::i_progress_callback),
+            progressContext: callbacks as *mut c_void,
+        };
+
+        let mut response: *mut BNDownloadInstanceResponse = null_mut();
+
+        let result = unsafe {
+            BNPerformCustomRequest(
+                self.handle,
+                method.into_bytes_with_nul().as_ref().as_ptr() as *const c_char,
+                url.into_bytes_with_nul().as_ref().as_ptr() as *const c_char,
+                header_key_ptrs.len() as u64,
+                header_key_ptrs.as_ptr(),
+                header_value_ptrs.as_ptr(),
+                &mut response as *mut *mut BNDownloadInstanceResponse,
+                &mut cbs as *mut BNDownloadInstanceInputOutputCallbacks,
+            )
+        };
+
+        if result < 0 {
+            unsafe { BNFreeDownloadInstanceResponse(response) };
+            return Err(self.get_error());
+        }
+
+        let mut response_headers = HashMap::new();
+        unsafe {
+            let response_header_keys: &[*mut c_char] =
+                slice::from_raw_parts((*response).headerKeys, (*response).headerCount as usize);
+            let response_header_values: &[*mut c_char] =
+                slice::from_raw_parts((*response).headerValues, (*response).headerCount as usize);
+
+            for (key, value) in response_header_keys
+                .iter()
+                .zip(response_header_values.iter())
+            {
+                response_headers.insert(
+                    BnStr::from_raw(*key).to_string(),
+                    BnStr::from_raw(*value).to_string(),
+                );
+            }
+        }
+
+        let r = DownloadResponse {
+            status_code: unsafe { (*response).statusCode },
+            headers: response_headers,
+        };
+
+        unsafe { BNFreeDownloadInstanceResponse(response) };
+
+        Ok(r)
+    }
+}
+
+impl ToOwned for DownloadInstance {
+    type Owned = Ref<Self>;
+
+    fn to_owned(&self) -> Self::Owned {
+        unsafe { RefCountable::inc_ref(self) }
+    }
+}
+
+unsafe impl RefCountable for DownloadInstance {
+    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
+        Ref::new(Self {
+            handle: BNNewDownloadInstanceReference(handle.handle),
+        })
+    }
+
+    unsafe fn dec_ref(handle: &Self) {
+        BNFreeDownloadInstance(handle.handle);
+    }
+}
diff --git a/rust/src/lib.rs b/rust/src/lib.rs
index c91c9792..c13d2b22 100644
--- a/rust/src/lib.rs
+++ b/rust/src/lib.rs
@@ -138,6 +138,7 @@ pub mod databuffer;
 pub mod debuginfo;
 pub mod demangle;
 pub mod disassembly;
+pub mod downloadprovider;
 pub mod fileaccessor;
 pub mod filemetadata;
 pub mod flowgraph;