Add unit test to detect memory leaks

5 files changed, 78 insertions, 1 deletions

diff --git a/binaryninjaapi.cpp b/binaryninjaapi.cpp
index f9b5f47b..66ddbffd 100644
--- a/binaryninjaapi.cpp
+++ b/binaryninjaapi.cpp
@@ -355,3 +355,16 @@ string BinaryNinja::GetUniqueIdentifierString()
 	BNFreeString(str);
 	return result;
 }
+
+
+map<string, uint64_t> BinaryNinja::GetMemoryUsageInfo()
+{
+	size_t count;
+	BNMemoryUsageInfo* info = BNGetMemoryUsageInfo(&count);
+
+	map<string, uint64_t> result;
+	for (size_t i = 0; i < count; i++)
+		result[info[i].name] = info[i].value;
+	BNFreeMemoryUsageInfo(info, count);
+	return result;
+}
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index b6d71165..2df09818 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -697,6 +697,8 @@ namespace BinaryNinja
 
 	std::string GetUniqueIdentifierString();
 
+	std::map<std::string, uint64_t> GetMemoryUsageInfo();
+
 	class DataBuffer
 	{
 		BNDataBuffer* m_buffer;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 383e6fb2..25bf1e46 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1913,6 +1913,12 @@ extern "C"
 		double seconds;
 	};
 
+	struct BNMemoryUsageInfo
+	{
+		char* name;
+		uint64_t value;
+	};
+
 	enum BNMetadataType
 	{
 		InvalidDataType,
@@ -3923,6 +3929,8 @@ extern "C"
 
 	BINARYNINJACOREAPI void* BNRegisterObjectRefDebugTrace(const char* typeName);
 	BINARYNINJACOREAPI void BNUnregisterObjectRefDebugTrace(const char* typeName, void* trace);
+	BINARYNINJACOREAPI BNMemoryUsageInfo* BNGetMemoryUsageInfo(size_t* count);
+	BINARYNINJACOREAPI void BNFreeMemoryUsageInfo(BNMemoryUsageInfo* info, size_t count);
 
 #ifdef __cplusplus
 }
diff --git a/python/__init__.py b/python/__init__.py
index 711a178a..c659aa1c 100644
--- a/python/__init__.py
+++ b/python/__init__.py
@@ -304,4 +304,14 @@ def core_ui_enabled():
 
 def core_set_license(licenseData):
 	'''Set the Binary Ninja license data, an alternative to storing a ``license.dat`` file'''
-	core.BNSetLicense(licenseData)
\ No newline at end of file
+	core.BNSetLicense(licenseData)
+
+
+def get_memory_usage_info():
+	count = ctypes.c_ulonglong()
+	info = core.BNGetMemoryUsageInfo(count)
+	result = {}
+	for i in range(0, count.value):
+		result[info[i].name] = info[i].value
+	core.BNFreeMemoryUsageInfo(info, count.value)
+	return result
diff --git a/suite/testcommon.py b/suite/testcommon.py
index 70f815a0..02ed13c4 100644
--- a/suite/testcommon.py
+++ b/suite/testcommon.py
@@ -862,3 +862,47 @@ class VerifyBuilder(Builder):
             return [str(functions == bndb_functions and comments == bndb_comments)]
         finally:
             self.delete_package("helloworld")
+
+    def test_memory_leaks(self):
+        """Detected memory leaks during analysis"""
+        # This test will attempt to detect object leaks during headless analysis
+        file_name = self.unpackage_file("helloworld")
+        try:
+            # Open the binary once and let any persistent structures be created (typically types)
+            bv = binja.BinaryViewType['ELF'].open(file_name)
+            bv.update_analysis_and_wait()
+            # Hold on to a graph reference while tearing down the binary view. This will keep a reference
+            # in the core. If we directly free the view, the teardown will happen in a worker thread and
+            # we will not be able to get a reliable object count. By keeping a reference in a different
+            # object in the core, the teardown will occur immediately upon freeing the other object.
+            graph = bv.functions[0].create_graph()
+            bv.file.close()
+            del bv
+            import gc
+            gc.collect()
+            del graph
+            gc.collect()
+
+            initial_object_counts = binja.get_memory_usage_info()
+
+            # Analyze the binary again
+            bv = binja.BinaryViewType['ELF'].open(file_name)
+            bv.update_analysis_and_wait()
+            graph = bv.functions[0].create_graph()
+            bv.file.close()
+            del bv
+            gc.collect()
+            del graph
+            gc.collect()
+
+            # Capture final object count
+            final_object_counts = binja.get_memory_usage_info()
+
+            # Check for leaks
+            ok = True
+            for i in initial_object_counts.keys():
+                if final_object_counts[i] > initial_object_counts[i]:
+                    ok = False
+            return ok
+        finally:
+            self.delete_package("helloworld")