From 1174e544b0283bd30962d0a6a528af0aca43367c Mon Sep 17 00:00:00 2001
From: Brian Potchik <brian@vector35.com>
Date: Thu, 17 Aug 2023 17:47:55 -0400
Subject: Ensure consistent lifetimes for UI-accessible BinaryView instances in
 Python.

---
 python/binaryview.py        | 26 ++++++++++++++++++++++++++
 python/scriptingprovider.py |  8 ++++++++
 2 files changed, 34 insertions(+)

(limited to 'python')

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)
-- 
cgit v1.3.1