From f6659517c08799cf52fddd09424eaea867ff3ed7 Mon Sep 17 00:00:00 2001
From: Rusty Wagner <rusty@vector35.com>
Date: Fri, 10 May 2019 16:15:48 -0400
Subject: Fix memory leak and use after free bugs in BinaryView and FlowGraph
 objects

---
 python/binaryview.py | 17 ++++++++++-------
 python/flowgraph.py  | 16 ++++++++++++++++
 2 files changed, 26 insertions(+), 7 deletions(-)

(limited to 'python')

diff --git a/python/binaryview.py b/python/binaryview.py
index 1baeab2b..66c30d06 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -735,7 +735,6 @@ class BinaryView(object):
 	_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:
@@ -755,7 +754,8 @@ 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.externalRefTaken = self._cb.externalRefTaken.__class__(self._external_ref_taken)
+			self._cb.externalRefReleased = self._cb.externalRefReleased.__class__(self._external_ref_released)
 			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)
@@ -778,8 +778,6 @@ 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
 
@@ -925,8 +923,7 @@ class BinaryView(object):
 	def __del__(self):
 		for i in self.notifications.values():
 			i._unregister()
-		if self._must_free:
-			core.BNFreeBinaryView(self.handle)
+		core.BNFreeBinaryView(self.handle)
 
 	def __iter__(self):
 		count = ctypes.c_ulonglong(0)
@@ -1388,7 +1385,13 @@ class BinaryView(object):
 			log.log_error(traceback.format_exc())
 			return False
 
-	def _free_object(self, ctxt):
+	def _external_ref_taken(self, ctxt):
+		try:
+			self.__class__._registered_instances.append(self)
+		except:
+			log.log_error(traceback.format_exc())
+
+	def _external_ref_released(self, ctxt):
 		try:
 			self.__class__._registered_instances.remove(self)
 		except:
diff --git a/python/flowgraph.py b/python/flowgraph.py
index 58add8b0..189b075c 100644
--- a/python/flowgraph.py
+++ b/python/flowgraph.py
@@ -338,6 +338,8 @@ class FlowGraph(object):
 	from incoming edges, or graphs that have disjoint subgraphs will not render correctly. This will be fixed \
 	in a future version.
 	"""
+	_registered_instances = []
+
 	def __init__(self, handle = None):
 		if handle is None:
 			self._ext_cb = core.BNCustomFlowGraph()
@@ -346,6 +348,8 @@ class FlowGraph(object):
 			self._ext_cb.populateNodes = self._ext_cb.populateNodes.__class__(self._populate_nodes)
 			self._ext_cb.completeLayout = self._ext_cb.completeLayout.__class__(self._complete_layout)
 			self._ext_cb.update = self._ext_cb.update.__class__(self._update)
+			self._ext_cb.externalRefTaken = self._ext_cb.externalRefTaken.__class__(self._external_ref_taken)
+			self._ext_cb.externalRefReleased = self._ext_cb.externalRefReleased.__class__(self._external_ref_released)
 			handle = core.BNCreateCustomFlowGraph(self._ext_cb)
 		self.handle = handle
 
@@ -390,6 +394,18 @@ class FlowGraph(object):
 			log.log_error(traceback.format_exc())
 			return None
 
+	def _external_ref_taken(self, ctxt):
+		try:
+			self.__class__._registered_instances.append(self)
+		except:
+			log.log_error(traceback.format_exc())
+
+	def _external_ref_released(self, ctxt):
+		try:
+			self.__class__._registered_instances.remove(self)
+		except:
+			log.log_error(traceback.format_exc())
+
 	def finish_prepare_for_layout(self):
 		"""
 		``finish_prepare_for_layout`` signals that preparations for rendering a graph are complete.
-- 
cgit v1.3.1