Ensure consistent lifetimes for UI-accessible BinaryView instances in Python.

5 files changed, 58 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index efe9a22a..d0d86bbd 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -14385,6 +14385,7 @@ namespace BinaryNinja {
 		static BNScriptingProviderExecuteResult ExecuteScriptInputCallback(void* ctxt, const char* input);
 		static BNScriptingProviderExecuteResult ExecuteScriptFromFilenameCallback(void *ctxt, const char* filename);
 		static void CancelScriptInputCallback(void* ctxt);
+		static void ReleaseBinaryViewCallback(void* ctxt, BNBinaryView* view);
 		static void SetCurrentBinaryViewCallback(void* ctxt, BNBinaryView* view);
 		static void SetCurrentFunctionCallback(void* ctxt, BNFunction* func);
 		static void SetCurrentBasicBlockCallback(void* ctxt, BNBasicBlock* block);
@@ -14399,6 +14400,7 @@ namespace BinaryNinja {
 		virtual BNScriptingProviderExecuteResult ExecuteScriptInput(const std::string& input) = 0;
 		virtual BNScriptingProviderExecuteResult ExecuteScriptInputFromFilename(const std::string& filename) = 0;
 		virtual void CancelScriptInput();
+		virtual void ReleaseBinaryView(BinaryView* view);
 		virtual void SetCurrentBinaryView(BinaryView* view);
 		virtual void SetCurrentFunction(Function* func);
 		virtual void SetCurrentBasicBlock(BasicBlock* block);
@@ -14432,6 +14434,7 @@ namespace BinaryNinja {
 		virtual BNScriptingProviderExecuteResult ExecuteScriptInput(const std::string& input) override;
 		virtual BNScriptingProviderExecuteResult ExecuteScriptInputFromFilename(const std::string& filename) override;
 		virtual void CancelScriptInput() override;
+		virtual void ReleaseBinaryView(BinaryView* view) override;
 		virtual void SetCurrentBinaryView(BinaryView* view) override;
 		virtual void SetCurrentFunction(Function* func) override;
 		virtual void SetCurrentBasicBlock(BasicBlock* block) override;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 3120e4a7..2f053802 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -2540,6 +2540,7 @@ extern "C"
 		BNScriptingProviderExecuteResult (*executeScriptInput)(void* ctxt, const char* input);
 		BNScriptingProviderExecuteResult (*executeScriptInputFromFilename)(void *ctxt, const char* input);
 		void (*cancelScriptInput)(void* ctxt);
+		void (*releaseBinaryView)(void* ctxt, BNBinaryView* view);
 		void (*setCurrentBinaryView)(void* ctxt, BNBinaryView* view);
 		void (*setCurrentFunction)(void* ctxt, BNFunction* func);
 		void (*setCurrentBasicBlock)(void* ctxt, BNBasicBlock* block);
@@ -6001,6 +6002,7 @@ extern "C"
 	BINARYNINJACOREAPI BNScriptingProviderExecuteResult BNExecuteScriptInputFromFilename(
 		BNScriptingInstance* instance, const char* filename);
 	BINARYNINJACOREAPI void BNCancelScriptInput(BNScriptingInstance* instance);
+	BINARYNINJACOREAPI void BNScriptingInstanceReleaseBinaryView(BNScriptingInstance* instance, BNBinaryView* view);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentBinaryView(BNScriptingInstance* instance, BNBinaryView* view);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentFunction(BNScriptingInstance* instance, BNFunction* func);
 	BINARYNINJACOREAPI void BNSetScriptingInstanceCurrentBasicBlock(BNScriptingInstance* instance, BNBasicBlock* block);
diff --git a/python/binaryview.py b/python/binaryview.py
index 6ba8d33f..4b1ca035 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -1911,6 +1911,30 @@ class BinaryView:
 	registered_view_type = None
 	_associated_data = {}
 	_registered_instances = []
+	_cached_instances = {}
+
+	@classmethod
+	def _cache_insert(cls, instance):
+		key = ctypes.addressof(instance.handle.contents)
+		if key not in cls._cached_instances:
+			cls._cached_instances[ctypes.addressof(instance.handle.contents)] = instance
+
+	@classmethod
+	def _cache_remove(cls, handle):
+		key = ctypes.addressof(handle.contents)
+		if key in cls._cached_instances:
+			cls._cached_instances.pop(key)
+
+	@classmethod
+	def _cache_contains(cls, handle):
+		return ctypes.addressof(handle.contents) in cls._cached_instances
+
+	def __new__(cls, file_metadata=None, parent_view=None, handle=None):
+		if handle:
+			key = ctypes.addressof(handle.contents)
+			if key in cls._cached_instances:
+				return cls._cached_instances[key]
+		return super().__new__(cls)
 
 	def __init__(
 	    self, file_metadata: Optional['filemetadata.FileMetadata'] = None, parent_view: Optional['BinaryView'] = None,
@@ -1918,6 +1942,8 @@ class BinaryView:
 	):
 		if handle is not None:
 			_handle = handle
+			if self.__class__._cache_contains(handle):
+				return
 			if file_metadata is None:
 				self._file = filemetadata.FileMetadata(handle=core.BNGetFileForView(handle))
 			else:
diff --git a/python/scriptingprovider.py b/python/scriptingprovider.py
index 10f97f0e..5e746558 100644
--- a/python/scriptingprovider.py
+++ b/python/scriptingprovider.py
@@ -138,6 +138,7 @@ class ScriptingInstance:
 			self._cb.executeScriptInput = self._cb.executeScriptInput.__class__(self._execute_script_input)
 			self._cb.executeScriptInputFromFilename = self._cb.executeScriptInputFromFilename.__class__(self._execute_script_input_from_filename)
 			self._cb.cancelScriptInput = self._cb.cancelScriptInput.__class__(self._cancel_script_input)
+			self._cb.releaseBinaryView = self._cb.releaseBinaryView.__class__(self._release_binary_view)
 			self._cb.setCurrentBinaryView = self._cb.setCurrentBinaryView.__class__(self._set_current_binary_view)
 			self._cb.setCurrentFunction = self._cb.setCurrentFunction.__class__(self._set_current_function)
 			self._cb.setCurrentBasicBlock = self._cb.setCurrentBasicBlock.__class__(self._set_current_basic_block)
@@ -189,10 +190,17 @@ class ScriptingInstance:
 			log_error(traceback.format_exc())
 			return ScriptingProviderExecuteResult.ScriptExecutionCancelled
 
+	def _release_binary_view(self, ctxt, view):
+		try:
+			binaryview.BinaryView._cache_remove(view)
+		except:
+			log_error(traceback.format_exc())
+
 	def _set_current_binary_view(self, ctxt, view):
 		try:
 			if view:
 				view = binaryview.BinaryView(handle=core.BNNewViewReference(view))
+				binaryview.BinaryView._cache_insert(view)
 			else:
 				view = None
 			self.perform_set_current_binary_view(view)
diff --git a/scriptingprovider.cpp b/scriptingprovider.cpp
index c3c4d615..1aefd4f7 100644
--- a/scriptingprovider.cpp
+++ b/scriptingprovider.cpp
@@ -63,6 +63,7 @@ ScriptingInstance::ScriptingInstance(ScriptingProvider* provider)
 	cb.externalRefReleased = nullptr;
 	cb.executeScriptInput = ExecuteScriptInputCallback;
 	cb.cancelScriptInput = CancelScriptInputCallback;
+	cb.releaseBinaryView = ReleaseBinaryViewCallback;
 	cb.setCurrentBinaryView = SetCurrentBinaryViewCallback;
 	cb.setCurrentFunction = SetCurrentFunctionCallback;
 	cb.setCurrentBasicBlock = SetCurrentBasicBlockCallback;
@@ -102,6 +103,14 @@ void ScriptingInstance::CancelScriptInputCallback(void* ctxt)
 }
 
 
+void ScriptingInstance::ReleaseBinaryViewCallback(void* ctxt, BNBinaryView* view)
+{
+	CallbackRef<ScriptingInstance> instance(ctxt);
+	Ref<BinaryView> object = view ? new BinaryView(BNNewViewReference(view)) : nullptr;
+	instance->ReleaseBinaryView(object);
+}
+
+
 void ScriptingInstance::SetCurrentBinaryViewCallback(void* ctxt, BNBinaryView* view)
 {
 	CallbackRef<ScriptingInstance> instance(ctxt);
@@ -168,6 +177,10 @@ void ScriptingInstance::DestroyInstance()
 
 void ScriptingInstance::CancelScriptInput() {}
 
+
+void ScriptingInstance::ReleaseBinaryView(BinaryView*) {}
+
+
 void ScriptingInstance::SetCurrentBinaryView(BinaryView*) {}
 
 
@@ -265,6 +278,12 @@ void CoreScriptingInstance::CancelScriptInput()
 }
 
 
+void CoreScriptingInstance::ReleaseBinaryView(BinaryView* view)
+{
+	BNScriptingInstanceReleaseBinaryView(m_object, view ? view->GetObject() : nullptr);
+}
+
+
 void CoreScriptingInstance::SetCurrentBinaryView(BinaryView* view)
 {
 	BNSetScriptingInstanceCurrentBinaryView(m_object, view ? view->GetObject() : nullptr);