python3: ensure callbacks live long enough

2 files changed, 19 insertions, 3 deletions

diff --git a/python/binaryview.py b/python/binaryview.py
index 84dd0f69..4239a9e8 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -676,8 +676,10 @@ class BinaryView(object):
 	registered_view_type = None
 	next_address = 0
 	_associated_data = {}
+	_registered_instances = []
 
 	def __init__(self, file_metadata=None, parent_view=None, handle=None):
+		self._must_free = True
 		if handle is not None:
 			self.handle = core.handle_of_type(handle, core.BNBinaryView)
 			if file_metadata is None:
@@ -697,6 +699,7 @@ class BinaryView(object):
 			self._cb = core.BNCustomBinaryView()
 			self._cb.context = 0
 			self._cb.init = self._cb.init.__class__(self._init)
+			self._cb.freeObject = self._cb.freeObject.__class__(self._free_object)
 			self._cb.read = self._cb.read.__class__(self._read)
 			self._cb.write = self._cb.write.__class__(self._write)
 			self._cb.insert = self._cb.insert.__class__(self._insert)
@@ -719,6 +722,8 @@ class BinaryView(object):
 			if parent_view is not None:
 				parent_view = parent_view.handle
 			self.handle = core.BNCreateCustomBinaryView(self.__class__.name, file_metadata.handle, parent_view, self._cb)
+			self.__class__._registered_instances.append(self)
+			self._must_free = False
 		self.notifications = {}
 		self.next_address = None  # Do NOT try to access view before init() is called, use placeholder
 
@@ -864,7 +869,8 @@ class BinaryView(object):
 	def __del__(self):
 		for i in self.notifications.values():
 			i._unregister()
-		core.BNFreeBinaryView(self.handle)
+		if self._must_free:
+			core.BNFreeBinaryView(self.handle)
 
 	def __iter__(self):
 		count = ctypes.c_ulonglong(0)
@@ -1310,6 +1316,12 @@ class BinaryView(object):
 			log.log_error(traceback.format_exc())
 			return False
 
+	def _free_object(self, ctxt):
+		try:
+			self.__class__._registered_instances.remove(self)
+		except:
+			log.log_error(traceback.format_exc())
+
 	def _read(self, ctxt, dest, offset, length):
 		try:
 			data = self.perform_read(offset, length)
diff --git a/python/downloadprovider.py b/python/downloadprovider.py
index eb9735f3..10f9dee1 100644
--- a/python/downloadprovider.py
+++ b/python/downloadprovider.py
@@ -39,6 +39,7 @@ from binaryninja import range
 
 
 class DownloadInstance(object):
+	_registered_instances = []
 	def __init__(self, provider, handle = None):
 		if handle is None:
 			self._cb = core.BNDownloadInstanceCallbacks()
@@ -46,15 +47,18 @@ class DownloadInstance(object):
 			self._cb.destroyInstance = self._cb.destroyInstance.__class__(self._destroy_instance)
 			self._cb.performRequest = self._cb.performRequest.__class__(self._perform_request)
 			self.handle = core.BNInitDownloadInstance(provider.handle, self._cb)
+			self.__class__._registered_instances.append(self)
 		else:
 			self.handle = core.handle_of_type(handle, core.BNDownloadInstance)
-		self._outputCallbacks = None
+		self._must_free = handle is not None
 
 	def __del__(self):
-		core.BNFreeDownloadInstance(self.handle)
+		if self._must_free:
+			core.BNFreeDownloadInstance(self.handle)
 
 	def _destroy_instance(self, ctxt):
 		try:
+			self.__class__._registered_instances.remove(self)
 			self.perform_destroy_instance()
 		except:
 			log.log_error(traceback.format_exc())