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

---
 binaryninjaapi.h | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

(limited to 'binaryninjaapi.h')

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 2df09818..daf3db4f 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -95,15 +95,22 @@ namespace BinaryNinja
 		{
 #ifdef WIN32
 			if (InterlockedDecrement((LONG*)&m_refs) == 0)
-				delete this;
+			{
+				if (!m_registeredRef)
+					delete this;
+			}
 #else
 			if (__sync_fetch_and_add(&m_refs, -1) == 1)
-				delete this;
+			{
+				if (!m_registeredRef)
+					delete this;
+			}
 #endif
 		}
 
 	public:
 		int m_refs;
+		bool m_registeredRef = false;
 		T* m_object;
 		CoreRefCountObject(): m_refs(0), m_object(nullptr) {}
 		virtual ~CoreRefCountObject() {}
@@ -133,13 +140,15 @@ namespace BinaryNinja
 
 		void AddRefForRegistration()
 		{
-			AddRefInternal();
+			m_registeredRef = true;
 		}
 
 		void ReleaseForRegistration()
 		{
 			m_object = nullptr;
-			ReleaseInternal();
+			m_registeredRef = false;
+			if (m_refs == 0)
+				delete this;
 		}
 	};
 
@@ -2806,6 +2815,7 @@ namespace BinaryNinja
 		static void PopulateNodesCallback(void* ctxt);
 		static void CompleteLayoutCallback(void* ctxt);
 		static BNFlowGraph* UpdateCallback(void* ctxt);
+		static void FreeObjectCallback(void* ctxt);
 
 	protected:
 		FlowGraph(BNFlowGraph* graph);
-- 
cgit v1.3.1