From 0b79039508f3c58d1a74ac9c0078cb223a27955f Mon Sep 17 00:00:00 2001
From: Rubens Brandao <git@rubens.io>
Date: Mon, 22 Apr 2024 09:51:45 -0300
Subject: [PATCH] implement open/create database with progress function

---
 rust/src/filemetadata.rs | 65 +++++++++++++++++++++++++++-------------
 1 file changed, 45 insertions(+), 20 deletions(-)

diff --git a/rust/src/filemetadata.rs b/rust/src/filemetadata.rs
index 6f335c249..3de14a9f3 100644
--- a/rust/src/filemetadata.rs
+++ b/rust/src/filemetadata.rs
@@ -40,6 +40,7 @@ use binaryninjacore_sys::{
     BNSetFilename,
     BNUndo,
 };
+use binaryninjacore_sys::{BNCreateDatabaseWithProgress, BNOpenExistingDatabaseWithProgress};
 
 use crate::binaryview::BinaryView;
 
@@ -204,16 +205,28 @@ impl FileMetadata {
         }
     }
 
-    pub fn create_database<S: BnStrCompatible>(&self, filename: S) -> bool {
+    pub fn create_database<S: BnStrCompatible>(
+        &self,
+        filename: S,
+        progress_func: Option<fn(usize, usize) -> bool>,
+    ) -> bool {
         let filename = filename.into_bytes_with_nul();
+        let filename_ptr = filename.as_ref().as_ptr() as *mut _;
         let raw = "Raw".into_bytes_with_nul();
-
-        unsafe {
-            BNCreateDatabase(
-                BNGetFileViewOfType(self.handle, raw.as_ptr() as *mut _),
-                filename.as_ref().as_ptr() as *mut _,
-                ptr::null_mut() as *mut _,
-            )
+        let raw_ptr = raw.as_ptr() as *mut _;
+
+        let handle = unsafe { BNGetFileViewOfType(self.handle, raw_ptr) };
+        match progress_func {
+            None => unsafe { BNCreateDatabase(handle, filename_ptr, ptr::null_mut()) },
+            Some(func) => unsafe {
+                BNCreateDatabaseWithProgress(
+                    handle,
+                    filename_ptr,
+                    core::mem::transmute(func),
+                    Some(cb_progress_func),
+                    ptr::null_mut(),
+                )
+            },
         }
     }
 
@@ -244,17 +257,25 @@ impl FileMetadata {
         }
     }
 
-    pub fn open_database<S: BnStrCompatible>(&self, filename: S) -> Result<Ref<BinaryView>, ()> {
+    pub fn open_database<S: BnStrCompatible>(
+        &self,
+        filename: S,
+        progress_func: Option<fn(usize, usize) -> bool>,
+    ) -> Result<Ref<BinaryView>, ()> {
         let filename = filename.into_bytes_with_nul();
         let filename_ptr = filename.as_ref().as_ptr() as *mut _;
 
-        let view = unsafe { BNOpenExistingDatabase(self.handle, filename_ptr) };
-
-        // TODO : add optional progress function
-        // let view = match progress_func {
-        //     None => BNOpenExistingDatabase(self.handle, filename_ptr),
-        //     _ => BNOpenExistingDatabaseWithProgress(self.handle, str(filename), None, ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong)(lambda ctxt, cur, total: progress_func(cur, total)))
-        // };
+        let view = match progress_func {
+            None => unsafe { BNOpenExistingDatabase(self.handle, filename_ptr) },
+            Some(func) => unsafe {
+                BNOpenExistingDatabaseWithProgress(
+                    self.handle,
+                    filename_ptr,
+                    core::mem::transmute(func),
+                    Some(cb_progress_func),
+                )
+            },
+        };
 
         if view.is_null() {
             Err(())
@@ -284,7 +305,11 @@ unsafe impl RefCountable for FileMetadata {
     }
 }
 
-/*
-BNCreateDatabase,
-BNCreateDatabaseWithProgress,
-*/
+unsafe extern "C" fn cb_progress_func(
+    ctxt: *mut ::std::os::raw::c_void,
+    progress: usize,
+    total: usize,
+) -> bool {
+    let func: fn(usize, usize) -> bool = core::mem::transmute(ctxt);
+    func(progress, total)
+}