DebugInfo: make parseInfo have a progress callback

This will enable us to have progress bars and cancellation of debuginfo application

7 files changed, 101 insertions, 59 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 3d29520f..70f8d45b 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -1112,6 +1112,13 @@ namespace BinaryNinja {
 	std::function<bool(size_t, size_t)> SplitProgress(
 	    std::function<bool(size_t, size_t)> originalFn, size_t subpart, std::vector<double> subpartWeights);
 
+	struct ProgressContext
+	{
+		std::function<bool(size_t, size_t)> callback;
+	};
+
+	bool ProgressCallback(void* ctxt, size_t current, size_t total);
+
 	std::string GetUniqueIdentifierString();
 
 	std::map<std::string, uint64_t> GetMemoryUsageInfo();
@@ -11984,7 +11991,7 @@ namespace BinaryNinja {
 		static std::vector<Ref<DebugInfoParser>> GetListForView(const Ref<BinaryView> data);
 
 		std::string GetName() const;
-		Ref<DebugInfo> Parse(Ref<BinaryView> view, Ref<DebugInfo> existingDebugInfo = nullptr) const;
+		Ref<DebugInfo> Parse(Ref<BinaryView> view, Ref<DebugInfo> existingDebugInfo = nullptr, std::function<bool(size_t, size_t)> progress = {}) const;
 
 		bool IsValidForView(const Ref<BinaryView> view) const;
 	};
@@ -11992,7 +11999,7 @@ namespace BinaryNinja {
 	class CustomDebugInfoParser : public DebugInfoParser
 	{
 		static bool IsValidCallback(void* ctxt, BNBinaryView* view);
-		static bool ParseCallback(void* ctxt, BNDebugInfo* debugInfo, BNBinaryView* view);
+		static bool ParseCallback(void* ctxt, BNDebugInfo* debugInfo, BNBinaryView* view, bool (*progress)(void*, size_t, size_t), void* progressCtxt);
 		BNDebugInfoParser* Register(const std::string& name);
 
 	  public:
@@ -12000,7 +12007,7 @@ namespace BinaryNinja {
 		virtual ~CustomDebugInfoParser() {}
 
 		virtual bool IsValid(Ref<BinaryView>) = 0;
-		virtual bool ParseInfo(Ref<DebugInfo>, Ref<BinaryView>) = 0;
+		virtual bool ParseInfo(Ref<DebugInfo>, Ref<BinaryView>, std::function<bool(size_t, size_t)>) = 0;
 	};
 
 	/*!
diff --git a/binaryninjacore.h b/binaryninjacore.h
index f9571a4f..bc4b9bfd 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -6271,7 +6271,7 @@ extern "C"
 	BINARYNINJACOREAPI char* BNRustSimplifyStrToStr(const char* const);
 
 	BINARYNINJACOREAPI BNDebugInfoParser* BNRegisterDebugInfoParser(const char* name,
-	    bool (*isValid)(void*, BNBinaryView*), bool (*parseInfo)(void*, BNDebugInfo*, BNBinaryView*), void* context);
+	    bool (*isValid)(void*, BNBinaryView*), bool (*parseInfo)(void*, BNDebugInfo*, BNBinaryView*, bool(*)(void*, size_t, size_t), void*), void* context);
 	BINARYNINJACOREAPI void BNUnregisterDebugInfoParser(const char* rawName);
 	BINARYNINJACOREAPI BNDebugInfoParser* BNGetDebugInfoParserByName(const char* name);
 	BINARYNINJACOREAPI BNDebugInfoParser** BNGetDebugInfoParsers(size_t* count);
@@ -6279,7 +6279,8 @@ extern "C"
 	BINARYNINJACOREAPI char* BNGetDebugInfoParserName(BNDebugInfoParser* parser);
 	BINARYNINJACOREAPI bool BNIsDebugInfoParserValidForView(BNDebugInfoParser* parser, BNBinaryView* view);
 	BINARYNINJACOREAPI BNDebugInfo* BNParseDebugInfo(
-	    BNDebugInfoParser* parser, BNBinaryView* view, BNDebugInfo* existingDebugInfo);
+	    BNDebugInfoParser* parser, BNBinaryView* view, BNDebugInfo* existingDebugInfo,
+	    bool (*progress)(void*, size_t, size_t), void* progressCtxt);
 	BINARYNINJACOREAPI BNDebugInfoParser* BNNewDebugInfoParserReference(BNDebugInfoParser* parser);
 	BINARYNINJACOREAPI void BNFreeDebugInfoParserReference(BNDebugInfoParser* parser);
 	BINARYNINJACOREAPI void BNFreeDebugInfoParserList(BNDebugInfoParser** parsers, size_t count);
diff --git a/binaryview.cpp b/binaryview.cpp
index 8c6080d0..e8c080df 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -3494,12 +3494,6 @@ void BinaryView::DefineUserType(const QualifiedName& name, Ref<Type> type)
 }
 
 
-struct ProgressCallback
-{
-	std::function<bool(size_t, size_t)> func;
-};
-
-
 void BinaryView::DefineTypes(const vector<pair<string, QualifiedNameAndType>>& types, std::function<bool(size_t, size_t)> progress)
 {
 	BNQualifiedNameTypeAndId* apiTypes = new BNQualifiedNameTypeAndId[types.size()];
@@ -3510,14 +3504,9 @@ void BinaryView::DefineTypes(const vector<pair<string, QualifiedNameAndType>>& t
 		apiTypes[i].id = BNAllocString(types[i].first.c_str());
 	}
 
-	ProgressCallback cb;
-	cb.func = progress;
-	BNDefineAnalysisTypes(m_object, apiTypes, types.size(), [](void* ctxt, size_t cur, size_t total) {
-		ProgressCallback* cb = (ProgressCallback*)ctxt;
-		if (cb->func)
-			return cb->func(cur, total);
-		return true;
-	}, &cb);
+	ProgressContext cb;
+	cb.callback = progress;
+	BNDefineAnalysisTypes(m_object, apiTypes, types.size(), ProgressCallback, &cb);
 
 	for (size_t i = 0; i < types.size(); i++)
 	{
@@ -3537,14 +3526,9 @@ void BinaryView::DefineUserTypes(const vector<QualifiedNameAndType>& types, std:
 		apiTypes[i].type = types[i].type->GetObject();
 	}
 
-	ProgressCallback cb;
-	cb.func = progress;
-	BNDefineUserAnalysisTypes(m_object, apiTypes, types.size(), [](void* ctxt, size_t cur, size_t total) {
-		ProgressCallback* cb = (ProgressCallback*)ctxt;
-		if (cb->func)
-			return cb->func(cur, total);
-		return true;
-	}, &cb);
+	ProgressContext cb;
+	cb.callback = progress;
+	BNDefineUserAnalysisTypes(m_object, apiTypes, types.size(), ProgressCallback, &cb);
 
 	for (size_t i = 0; i < types.size(); i++)
 	{
diff --git a/database.cpp b/database.cpp
index 326b57f2..ba20e02b 100644
--- a/database.cpp
+++ b/database.cpp
@@ -25,19 +25,6 @@ using namespace Json;
 using namespace std;
 
 
-struct ProgressContext
-{
-	std::function<bool(size_t, size_t)> callback;
-};
-
-
-bool ProgressCallback(void* ctxt, size_t current, size_t total)
-{
-	ProgressContext* pctxt = reinterpret_cast<ProgressContext*>(ctxt);
-	return pctxt->callback(current, total);
-}
-
-
 KeyValueStore::KeyValueStore()
 {
 	m_object = BNCreateKeyValueStore();
diff --git a/debuginfo.cpp b/debuginfo.cpp
index c9a6f319..9783cba8 100644
--- a/debuginfo.cpp
+++ b/debuginfo.cpp
@@ -337,19 +337,25 @@ string DebugInfoParser::GetName() const
 }
 
 
-Ref<DebugInfo> DebugInfoParser::Parse(Ref<BinaryView> view, Ref<DebugInfo> existingDebugInfo) const
+Ref<DebugInfo> DebugInfoParser::Parse(Ref<BinaryView> view, Ref<DebugInfo> existingDebugInfo, std::function<bool(size_t, size_t)> progress) const
 {
+	ProgressContext ctxt;
+	if (progress)
+		ctxt.callback = progress;
+	else
+		ctxt.callback = [=](size_t, size_t) { return true; };
+
 	BNDebugInfo* info = nullptr;
 	if (existingDebugInfo)
 	{
-		info = BNParseDebugInfo(m_object, view->GetObject(), existingDebugInfo->GetObject());
+		info = BNParseDebugInfo(m_object, view->GetObject(), existingDebugInfo->GetObject(), ProgressCallback, &ctxt);
 		if (!info)
 			return nullptr;
 		info = BNNewDebugInfoReference(info);
 	}
 	else
 	{
-		info = BNParseDebugInfo(m_object, view->GetObject(), nullptr);
+		info = BNParseDebugInfo(m_object, view->GetObject(), nullptr, ProgressCallback, &ctxt);
 		if (!info)
 			return nullptr;
 	}
@@ -375,10 +381,12 @@ bool CustomDebugInfoParser::IsValidCallback(void* ctxt, BNBinaryView* view)
 }
 
 
-bool CustomDebugInfoParser::ParseCallback(void* ctxt, BNDebugInfo* debugInfo, BNBinaryView* view)
+bool CustomDebugInfoParser::ParseCallback(void* ctxt, BNDebugInfo* debugInfo, BNBinaryView* view, bool (*progress)(void*, size_t, size_t), void* progressCtxt)
 {
 	CustomDebugInfoParser* parser = (CustomDebugInfoParser*)ctxt;
-	return parser->ParseInfo(new DebugInfo(debugInfo), new BinaryView(view));
+	return parser->ParseInfo(new DebugInfo(debugInfo), new BinaryView(view), [=](size_t cur, size_t max) {
+		return progress(progressCtxt, cur, max);
+	});
 }
 
 
diff --git a/python/debuginfo.py b/python/debuginfo.py
index 4a809643..fc434bfd 100644
--- a/python/debuginfo.py
+++ b/python/debuginfo.py
@@ -34,6 +34,7 @@ from . import binaryview
 from . import filemetadata
 
 _debug_info_parsers = {}
+ProgressFuncType = Callable[[int, int], bool]
 
 
 class _DebugInfoParserMetaClass(type):
@@ -106,15 +107,15 @@ class _DebugInfoParserMetaClass(type):
 
 	@staticmethod
 	def _parse_info(
-	    debug_info: core.BNDebugInfo, view: core.BNBinaryView,
-	    callback: Callable[["DebugInfo", 'binaryview.BinaryView'], bool]
+	    debug_info: core.BNDebugInfo, view: core.BNBinaryView, progress: ProgressFuncType,
+	    callback: Callable[["DebugInfo", 'binaryview.BinaryView'], bool],
 	) -> bool:
 		try:
 			file_metadata = filemetadata.FileMetadata(handle=core.BNGetFileForView(view))
 			view_obj = binaryview.BinaryView(file_metadata=file_metadata, handle=core.BNNewViewReference(view))
 			parser_ref = core.BNNewDebugInfoReference(debug_info)
 			assert parser_ref is not None, "core.BNNewDebugInfoReference returned None"
-			return callback(DebugInfo(parser_ref), view_obj)
+			return callback(DebugInfo(parser_ref), view_obj, progress)
 		except:
 			log_error(traceback.format_exc())
 			return False
@@ -122,7 +123,7 @@ class _DebugInfoParserMetaClass(type):
 	@classmethod
 	def register(
 	    cls, name: str, is_valid: Callable[['binaryview.BinaryView'], bool],
-	    parse_info: Callable[["DebugInfo", 'binaryview.BinaryView'], None]
+	    parse_info: Callable[["DebugInfo", 'binaryview.BinaryView', ProgressFuncType], bool]
 	) -> "DebugInfoParser":
 		"""Registers a DebugInfoParser. See ``debuginfo.DebugInfoParser`` for more details."""
 		binaryninja._init_plugins()
@@ -131,8 +132,9 @@ class _DebugInfoParserMetaClass(type):
 		                               ctypes.POINTER(core.BNBinaryView
 		                                              ))(lambda ctxt, view: cls._is_valid(view, is_valid))
 		parse_info_cb = ctypes.CFUNCTYPE(
-		    ctypes.c_bool, ctypes.c_void_p, ctypes.POINTER(core.BNDebugInfo), ctypes.POINTER(core.BNBinaryView)
-		)(lambda ctxt, debug_info, view: cls._parse_info(debug_info, view, parse_info))
+		    ctypes.c_bool, ctypes.c_void_p, ctypes.POINTER(core.BNDebugInfo), ctypes.POINTER(core.BNBinaryView),
+			ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_size_t), ctypes.c_void_p,
+		)(lambda ctxt, debug_info, view, progress, progress_ctxt: cls._parse_info(debug_info, view, lambda cur, max: progress(progress_ctxt, cur, max), parse_info))
 
 		# Don't let our callbacks get garbage collected
 		global _debug_info_parsers
@@ -217,17 +219,21 @@ class DebugInfoParser(object, metaclass=_DebugInfoParserMetaClass):
 		"""Returns whether this debug-info parser is valid for the provided binary view"""
 		return core.BNIsDebugInfoParserValidForView(self.handle, view.handle)
 
-	def parse_debug_info(self, view: 'binaryview.BinaryView', debug_info: Optional["DebugInfo"] = None) -> Optional["DebugInfo"]:
+	def parse_debug_info(self, view: 'binaryview.BinaryView', debug_info: Optional["DebugInfo"] = None, progress: ProgressFuncType = None) -> Optional["DebugInfo"]:
 		"""Returns a ``DebugInfo`` object populated with debug info by this debug-info parser. Only provide a ``DebugInfo`` object if you wish to append to the existing debug info"""
+		if progress is None:
+			progress = lambda cur, max: True
+		progress_c = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_size_t)(lambda ctxt, cur, max: progress(cur, max))
+
 		if isinstance(debug_info, DebugInfo):
-			parser = core.BNParseDebugInfo(self.handle, view.handle, debug_info.handle)
+			parser = core.BNParseDebugInfo(self.handle, view.handle, debug_info.handle, progress_c, None)
 			if parser is None:
 				return None
 			parser_ref = core.BNNewDebugInfoReference(parser)
 			assert parser_ref is not None, "core.BNNewDebugInfoReference returned None"
 			return DebugInfo(parser_ref)
 		else:
-			parser = core.BNParseDebugInfo(self.handle, view.handle, None)
+			parser = core.BNParseDebugInfo(self.handle, view.handle, None, progress_c, None)
 			if parser is None:
 				return None
 			return DebugInfo(parser)
diff --git a/rust/src/debuginfo.rs b/rust/src/debuginfo.rs
index 4e33e009..57e40f71 100644
--- a/rust/src/debuginfo.rs
+++ b/rust/src/debuginfo.rs
@@ -40,7 +40,7 @@
 //!         true
 //!     }
 //!
-//!     fn parse_info(&self, _debug_info: &mut DebugInfo, _view: &BinaryView) {
+//!     fn parse_info(&self, _debug_info: &mut DebugInfo, _view: &BinaryView, _progress: Box<dyn Fn(usize, usize) -> bool>) {
 //!         println!("Parsing info");
 //!     }
 //! }
@@ -76,8 +76,11 @@ use crate::{
     types::{DataVariableAndName, NameAndType, Type},
 };
 
+use std::mem::transmute;
 use std::{hash::Hash, mem, os::raw::c_void, ptr, slice};
 
+struct ProgressContext(Option<Box<dyn Fn(usize, usize) -> Result<(), ()>>>);
+
 //////////////////////
 //  DebugInfoParser
 
@@ -131,17 +134,43 @@ impl DebugInfoParser {
         unsafe { BNIsDebugInfoParserValidForView(self.handle, view.handle) }
     }
 
+    extern "C" fn cb_progress(ctxt: *mut c_void, cur: usize, max: usize) -> bool {
+        ffi_wrap!("DebugInfoParser::cb_progress", unsafe {
+            let progress = ctxt as *mut ProgressContext;
+            match &(*progress).0 {
+                Some(func) => (func)(cur, max).is_ok(),
+                None => true,
+            }
+        })
+    }
+
     /// Returns a `DebugInfo` object populated with debug info by this debug-info parser. Only provide a `DebugInfo` object if you wish to append to the existing debug info
     pub fn parse_debug_info(
         &self,
         view: &BinaryView,
         existing_debug_info: Option<&DebugInfo>,
+        progress: Option<Box<dyn Fn(usize, usize) -> Result<(), ()>>>,
     ) -> Option<Ref<DebugInfo>> {
+        let mut progress_raw = ProgressContext(progress);
         let info: *mut BNDebugInfo = match existing_debug_info {
             Some(debug_info) => unsafe {
-                BNParseDebugInfo(self.handle, view.handle, debug_info.handle)
+                BNParseDebugInfo(
+                    self.handle,
+                    view.handle,
+                    debug_info.handle,
+                    Some(Self::cb_progress),
+                    &mut progress_raw as *mut _ as *mut c_void,
+                )
+            },
+            None => unsafe {
+                BNParseDebugInfo(
+                    self.handle,
+                    view.handle,
+                    ptr::null_mut(),
+                    Some(Self::cb_progress),
+                    &mut progress_raw as *mut _ as *mut c_void,
+                )
             },
-            None => unsafe { BNParseDebugInfo(self.handle, view.handle, ptr::null_mut()) },
         };
         if info.is_null() {
             return None;
@@ -171,6 +200,8 @@ impl DebugInfoParser {
             ctxt: *mut c_void,
             debug_info: *mut BNDebugInfo,
             view: *mut BNBinaryView,
+            progress: Option<unsafe extern "C" fn(*mut c_void, usize, usize) -> bool>,
+            progress_ctxt: *mut c_void,
         ) -> bool
         where
             C: CustomDebugInfoParser,
@@ -180,7 +211,20 @@ impl DebugInfoParser {
                 let view = BinaryView::from_raw(view);
                 let mut debug_info = DebugInfo::from_raw(debug_info);
 
-                cmd.parse_info(&mut debug_info, &view)
+                cmd.parse_info(
+                    &mut debug_info,
+                    &view,
+                    Box::new(move |cur: usize, max: usize| match progress {
+                        Some(func) => {
+                            if func(progress_ctxt, cur, max) {
+                                Ok(())
+                            } else {
+                                Err(())
+                            }
+                        }
+                        _ => Ok(()),
+                    }),
+                )
             })
         }
 
@@ -787,5 +831,10 @@ impl ToOwned for DebugInfo {
 /// Implement this trait to implement a debug info parser.  See `DebugInfoParser` for more details.
 pub trait CustomDebugInfoParser: 'static + Sync {
     fn is_valid(&self, view: &BinaryView) -> bool;
-    fn parse_info(&self, debug_info: &mut DebugInfo, view: &BinaryView) -> bool;
+    fn parse_info(
+        &self,
+        debug_info: &mut DebugInfo,
+        view: &BinaryView,
+        progress: Box<dyn Fn(usize, usize) -> Result<(), ()>>,
+    ) -> bool;
 }