Fix memory leak in functions calling BNCreateUserFunction

4 files changed, 20 insertions, 7 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 404651c7..d7ab80f5 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -4120,7 +4120,7 @@ namespace BinaryNinja {
 			\param platform Platform for the function to be loaded
 		    \param addr Virtual adddress of the function to be loaded
 		*/
-		void CreateUserFunction(Platform* platform, uint64_t start);
+		Ref<Function> CreateUserFunction(Platform* platform, uint64_t start);
 
 		/*! removes a user function from the list of functions
 
diff --git a/binaryview.cpp b/binaryview.cpp
index 6120e10f..db1974da 100644
--- a/binaryview.cpp
+++ b/binaryview.cpp
@@ -1795,9 +1795,12 @@ void BinaryView::RemoveAnalysisFunction(Function* func, bool updateRefs)
 }
 
 
-void BinaryView::CreateUserFunction(Platform* platform, uint64_t start)
+Ref<Function> BinaryView::CreateUserFunction(Platform* platform, uint64_t start)
 {
-	BNCreateUserFunction(m_object, platform->GetObject(), start);
+	BNFunction* func = BNCreateUserFunction(m_object, platform->GetObject(), start);
+	if (!func)
+		return nullptr;
+	return new Function(func);
 }
 
 
diff --git a/python/binaryview.py b/python/binaryview.py
index 1a17faa4..01d44471 100644
--- a/python/binaryview.py
+++ b/python/binaryview.py
@@ -3982,7 +3982,8 @@ class BinaryView:
 		"""
 		core.BNRemoveAnalysisFunction(self.handle, func.handle, update_refs)
 
-	def create_user_function(self, addr: int, plat: Optional['_platform.Platform'] = None) -> '_function.Function':
+	def create_user_function(self, addr: int, plat: Optional['_platform.Platform'] = None) \
+			-> Optional['_function.Function']:
 		"""
 		``create_user_function`` add a new *user* function of the given ``plat`` at the virtual address ``addr``
 
@@ -4000,7 +4001,10 @@ class BinaryView:
 			if self.platform is None:
 				raise Exception("Attempting to call create_user_function with no specified platform")
 			plat = self.platform
-		return _function.Function(self, core.BNCreateUserFunction(self.handle, plat.handle, addr))
+		func = core.BNCreateUserFunction(self.handle, plat.handle, addr)
+		if func is None:
+			return None
+		return _function.Function(self, func)
 
 	def remove_user_function(self, func: '_function.Function') -> None:
 		"""
diff --git a/rust/src/binaryview.rs b/rust/src/binaryview.rs
index 31a12105..d6916e2c 100644
--- a/rust/src/binaryview.rs
+++ b/rust/src/binaryview.rs
@@ -853,9 +853,15 @@ pub trait BinaryViewExt: BinaryViewBase {
         }
     }
 
-    fn create_user_function(&self, plat: &Platform, addr: u64) {
+    fn create_user_function(&self, plat: &Platform, addr: u64) -> Result<Ref<Function>> {
         unsafe {
-            BNCreateUserFunction(self.as_ref().handle, plat.handle, addr);
+            let func = BNCreateUserFunction(self.as_ref().handle, plat.handle, addr);
+
+            if func.is_null() {
+                return Err(());
+            }
+
+            Ok(Function::from_raw(func))
         }
     }