Improvements for advanced binary search.

4 files changed, 45 insertions, 11 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index baa2c5fc..ceb42937 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -6490,7 +6490,9 @@ namespace BinaryNinja {
 			const FunctionViewType& viewType, const std::function<bool(size_t current, size_t total)>& progress,
 		    const std::function<bool(uint64_t addr, const LinearDisassemblyLine& line)>& matchCallback);
 
-		bool Search(const std::string& query, const std::function<bool(uint64_t offset, const DataBuffer& buffer)>& otherCallback);
+		bool Search(const std::string& query,
+			const std::function<bool(size_t current, size_t total)>& progressCallback,
+			const std::function<bool(uint64_t addr, const DataBuffer& buffer)>& matchCallback);
 
 		void Reanalyze();
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 6487c534..97e7c52a 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,14 +37,14 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 96
+#define BN_CURRENT_CORE_ABI_VERSION 97
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
 // will require rebuilding. The minimum version is increased when there are
 // incompatible changes that break binary compatibility, such as changes to
 // existing types or functions.
-#define BN_MINIMUM_CORE_ABI_VERSION 94
+#define BN_MINIMUM_CORE_ABI_VERSION 97
 
 #ifdef __GNUC__
 	#ifdef BINARYNINJACORE_LIBRARY
@@ -4214,7 +4214,7 @@ extern "C"
 	    BNProgressFunction progress, void* matchCtxt,
 	    bool (*matchCallback)(void* matchCtxt, uint64_t addr, BNLinearDisassemblyLine* line));
 
-	BINARYNINJACOREAPI bool BNSearch(BNBinaryView* view, const char* query, void* context, bool (*callback)(void*, uint64_t, BNDataBuffer*));
+	BINARYNINJACOREAPI bool BNSearch(BNBinaryView* view, const char* query, void* context, BNProgressFunction progressCallback, void* matchContext, bool (*callback)(void*, uint64_t, BNDataBuffer*));
 	BINARYNINJACOREAPI bool BNPerformSearch(const char* query, const uint8_t* buffer, size_t size, bool(*callback)(void*, size_t, size_t), void* context);
 
 	BINARYNINJACOREAPI void BNBeginBulkAddSegments(BNBinaryView* view);
diff --git a/binaryview.cpp b/binaryview.cpp
index a6546951..981e8329 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -4795,11 +4795,13 @@ bool BinaryView::FindAllConstant(uint64_t start, uint64_t end, uint64_t constant
 }
 
 
-bool BinaryView::Search(const string& query, const std::function<bool(uint64_t offset, const DataBuffer& buffer)>& otherCallback)
+bool BinaryView::Search(const string& query, const std::function<bool(size_t current, size_t total)>& progressCallback, const std::function<bool(uint64_t offset, const DataBuffer& buffer)>& matchCallback)
 {
+	ProgressContext fp;
+	fp.callback = progressCallback;
 	MatchCallbackContextForDataBuffer mc;
-	mc.func = otherCallback;
-	return BNSearch(m_object, query.c_str(), &mc, MatchCallbackForDataBuffer);
+	mc.func = matchCallback;
+	return BNSearch(m_object, query.c_str(), &fp, ProgressCallback, &mc, MatchCallbackForDataBuffer);
 }
 
 
diff --git a/python/binaryview.py b/python/binaryview.py
index dcdfe2fe..c6a1c096 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -9211,7 +9211,8 @@ to a the type "tagRECT" found in the typelibrary "winX64common"
 
 			return self.QueueGenerator(t, results)
 
-	def search(self, pattern: str, start: int = None, end: int = None, raw: bool = False, ignore_case: bool = False, overlap: bool = False, align: int = 1) -> QueueGenerator:
+	def search(self, pattern: str, start: int = None, end: int = None, raw: bool = False, ignore_case: bool = False, overlap: bool = False, align: int = 1,
+		limit: int = None, progress_callback: Optional[ProgressFuncType] = None, match_callback: Optional[DataMatchCallbackType] = None) -> QueueGenerator:
 		r"""
 		Searches for matches of the specified `pattern` within this BinaryView with an optionally provided address range specified by `start` and `end`.
 		The search pattern can be interpreted in various ways:
@@ -9230,6 +9231,12 @@ to a the type "tagRECT" found in the typelibrary "winX64common"
 		:param bool ignore_case: Whether to perform case-insensitive matching (default: False).
 		:param bool overlap: Whether to allow matches to overlap (default: False).
 		:param int align: The alignment of matches, must be a power of 2 (default: 1).
+		:param int limit: The maximum number of matches to return (default: None).
+		:param callback progress_callback: An optional function to be called with the current progress and total count. \
+		This function should return a boolean value that decides whether the search should continue or stop.
+		:param callback match_callback: A function that gets called when a match is found. The callback takes two parameters: \
+		the address of the match, and the actual DataBuffer that satisfies the search. This function can return a boolean value \
+		that decides whether the search should continue or stop.
 
 		:return: A generator object that yields the offset and matched DataBuffer for each match found.
 		:rtype: QueueGenerator
@@ -9260,10 +9267,33 @@ to a the type "tagRECT" found in the typelibrary "winX64common"
 			"overlap": overlap,
 			"align": align
 		}
+
+		if progress_callback:
+			progress_callback_obj = ctypes.CFUNCTYPE(
+				ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong
+			)(lambda ctxt, cur, total: progress_callback(cur, total))
+		else:
+			progress_callback_obj = ctypes.CFUNCTYPE(
+				ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_ulonglong
+			)(lambda ctxt, cur, total: True)
+
+		match_count = 0
+		def internal_match_callback(ctxt, offset, match):
+			nonlocal match_count
+			match_count += 1
+			data_buffer = databuffer.DataBuffer(handle=match)
+			results.put((offset, data_buffer))
+			if match_callback is not None:
+				should_continue = match_callback(offset, data_buffer)
+				if not should_continue:
+					return False
+			if limit is not None and match_count >= limit:
+				return False
+			return True
+
+		match_callback_obj = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.POINTER(core.BNDataBuffer))(internal_match_callback)
 		results = queue.Queue()
-		match_callback_obj = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_ulonglong, ctypes.POINTER(core.BNDataBuffer)
-		)(lambda ctxt, offset, match: results.put((offset, databuffer.DataBuffer(handle=match))) or True)
-		t = threading.Thread(target=lambda: core.BNSearch(self.handle, json.dumps(query), None, match_callback_obj))
+		t = threading.Thread(target=lambda: core.BNSearch(self.handle, json.dumps(query), None, progress_callback_obj, None, match_callback_obj))
 		return self.QueueGenerator(t, results)
 
 	def reanalyze(self) -> None: