2 files changed, 147 insertions, 36 deletions

diff --git a/rust/src/binary_view.rs b/rust/src/binary_view.rs
index a5d3ea6b..822c421a 100644
--- a/rust/src/binary_view.rs
+++ b/rust/src/binary_view.rs
@@ -947,29 +947,36 @@ pub trait BinaryViewExt: BinaryViewBase {
         MemoryMap::new(self.as_ref().to_owned())
     }
 
-    fn add_auto_function(&self, plat: &Platform, addr: u64) -> Option<Ref<Function>> {
-        unsafe {
-            let handle = BNAddFunctionForAnalysis(
-                self.as_ref().handle,
-                plat.handle,
-                addr,
-                false,
-                std::ptr::null_mut(),
-            );
-
-            if handle.is_null() {
-                return None;
-            }
+    /// Add an auto function at the given `address` with the views default platform.
+    ///
+    /// Use [`BinaryViewExt::add_auto_function_with_platform`] if you wish to specify a platform.
+    ///
+    /// NOTE: The default platform **must** be set for this view!
+    fn add_auto_function(&self, address: u64) -> Option<Ref<Function>> {
+        let platform = self.default_platform()?;
+        self.add_auto_function_with_platform(address, &platform)
+    }
 
-            Some(Function::ref_from_raw(handle))
-        }
+    /// Add an auto function at the given `address` with the `platform`.
+    ///
+    /// Use [`BinaryViewExt::add_auto_function_ext`] if you wish to specify a function type.
+    ///
+    /// NOTE: If the view's default platform is not set, this will set it to `platform`.
+    fn add_auto_function_with_platform(
+        &self,
+        address: u64,
+        platform: &Platform,
+    ) -> Option<Ref<Function>> {
+        self.add_auto_function_ext(address, platform, None)
     }
 
-    fn add_function_with_type(
+    /// Add an auto function at the given `address` with the `platform` and function type.
+    ///
+    /// NOTE: If the view's default platform is not set, this will set it to `platform`.
+    fn add_auto_function_ext(
         &self,
-        plat: &Platform,
-        addr: u64,
-        auto_discovered: bool,
+        address: u64,
+        platform: &Platform,
         func_type: Option<&Type>,
     ) -> Option<Ref<Function>> {
         unsafe {
@@ -980,9 +987,9 @@ pub trait BinaryViewExt: BinaryViewBase {
 
             let handle = BNAddFunctionForAnalysis(
                 self.as_ref().handle,
-                plat.handle,
-                addr,
-                auto_discovered,
+                platform.handle,
+                address,
+                true,
                 func_type,
             );
 
@@ -994,28 +1001,75 @@ pub trait BinaryViewExt: BinaryViewBase {
         }
     }
 
-    fn add_entry_point(&self, plat: &Platform, addr: u64) {
+    /// Remove an auto function from the view.
+    ///
+    /// Pass `true` for `update_refs` to update all references of the function.
+    ///
+    /// NOTE: Unlike [`BinaryViewExt::remove_user_function`], this will NOT prohibit the function from
+    /// being re-added in the future, use [`BinaryViewExt::remove_user_function`] to blacklist the
+    /// function from being automatically created.
+    fn remove_auto_function(&self, func: &Function, update_refs: bool) {
         unsafe {
-            BNAddEntryPointForAnalysis(self.as_ref().handle, plat.handle, addr);
+            BNRemoveAnalysisFunction(self.as_ref().handle, func.handle, update_refs);
         }
     }
 
-    fn create_user_function(&self, plat: &Platform, addr: u64) -> Result<Ref<Function>> {
-        unsafe {
-            let func = BNCreateUserFunction(self.as_ref().handle, plat.handle, addr);
+    /// Add a user function at the given `address` with the views default platform.
+    ///
+    /// Use [`BinaryViewExt::add_user_function_with_platform`] if you wish to specify a platform.
+    ///
+    /// NOTE: The default platform **must** be set for this view!
+    fn add_user_function(&self, addr: u64) -> Option<Ref<Function>> {
+        let platform = self.default_platform()?;
+        self.add_user_function_with_platform(addr, &platform)
+    }
 
+    /// Add an auto function at the given `address` with the `platform`.
+    ///
+    /// NOTE: If the view's default platform is not set, this will set it to `platform`.
+    fn add_user_function_with_platform(
+        &self,
+        addr: u64,
+        platform: &Platform,
+    ) -> Option<Ref<Function>> {
+        unsafe {
+            let func = BNCreateUserFunction(self.as_ref().handle, platform.handle, addr);
             if func.is_null() {
-                return Err(());
+                return None;
             }
-
-            Ok(Function::ref_from_raw(func))
+            Some(Function::ref_from_raw(func))
         }
     }
 
+    /// Removes the function from the view and blacklists it from being created automatically.
+    ///
+    /// NOTE: If you call [`BinaryViewExt::add_user_function`], it will override the blacklist.
+    fn remove_user_function(&self, func: &Function) {
+        unsafe { BNRemoveUserFunction(self.as_ref().handle, func.handle) }
+    }
+
     fn has_functions(&self) -> bool {
         unsafe { BNHasFunctions(self.as_ref().handle) }
     }
 
+    /// Add an entry point at the given `address` with the view's default platform.
+    ///
+    /// NOTE: The default platform **must** be set for this view!
+    fn add_entry_point(&self, addr: u64) {
+        if let Some(platform) = self.default_platform() {
+            self.add_entry_point_with_platform(addr, &platform);
+        }
+    }
+
+    /// Add an entry point at the given `address` with the `platform`.
+    ///
+    /// NOTE: If the view's default platform is not set, this will set it to `platform`.
+    fn add_entry_point_with_platform(&self, addr: u64, platform: &Platform) {
+        unsafe {
+            BNAddEntryPointForAnalysis(self.as_ref().handle, platform.handle, addr);
+        }
+    }
+
     fn entry_point_function(&self) -> Option<Ref<Function>> {
         unsafe {
             let raw_func_ptr = BNGetAnalysisEntryPoint(self.as_ref().handle);
diff --git a/rust/tests/function.rs b/rust/tests/function.rs
index 455ba752..ac15fd91 100644
--- a/rust/tests/function.rs
+++ b/rust/tests/function.rs
@@ -1,7 +1,7 @@
-use binaryninja::binary_view::BinaryViewExt;
+use binaryninja::binary_view::{BinaryView, BinaryViewExt};
+use binaryninja::file_metadata::FileMetadata;
 use binaryninja::headless::Session;
-use binaryninja::metadata::Metadata;
-use binaryninja::rc::Ref;
+use binaryninja::platform::Platform;
 use std::path::PathBuf;
 
 #[test]
@@ -15,7 +15,7 @@ fn store_and_query_function_metadata() {
 
     // Store key/value pairs to user and auto metadata
     func.store_metadata("one", "one", false);
-    func.store_metadata("two", 2 as u64, true);
+    func.store_metadata("two", 2u64, true);
     func.store_metadata("three", "three", true);
     func.remove_metadata("three");
 
@@ -35,8 +35,9 @@ fn store_and_query_function_metadata() {
             .unwrap(),
         2
     );
-    assert!(
-        func.query_metadata("three") == None,
+    assert_eq!(
+        func.query_metadata("three"),
+        None,
         "Query for key \"three\" returned a value"
     );
 
@@ -67,3 +68,59 @@ fn store_and_query_function_metadata() {
     );
     assert_eq!(auto_metadata.get("one"), Ok(None));
 }
+
+#[test]
+fn add_function() {
+    let _session = Session::new().expect("Failed to initialize session");
+    let out_dir = env!("OUT_DIR").parse::<PathBuf>().unwrap();
+    let view = binaryninja::load(out_dir.join("atox.obj")).expect("Failed to create view");
+    let mut func = view
+        .entry_point_function()
+        .expect("Failed to get entry point function");
+
+    // Remove the function then as an auto function then add as auto function.
+    // This tests to make sure that the function is not blacklisted from being added back.
+    view.remove_auto_function(&func, false);
+
+    assert_eq!(
+        view.functions_at(func.start()).len(),
+        0,
+        "Function was not removed"
+    );
+    func = view
+        .add_auto_function(func.start())
+        .expect("Failed to add function");
+    assert_eq!(
+        view.functions_at(func.start()).len(),
+        1,
+        "Function was not added back"
+    );
+
+    // Use the user version of remove to blacklist the function, auto function should be prohibited.
+    view.remove_user_function(&func);
+    assert_eq!(
+        view.add_auto_function(func.start()),
+        None,
+        "Function was not blacklisted"
+    );
+
+    // Adding back as a user should override the blacklist.
+    func = view
+        .add_user_function(func.start())
+        .expect("Failed to add function as user");
+    assert_eq!(
+        view.functions_at(func.start()).len(),
+        1,
+        "Function was not added back"
+    );
+
+    // Make sure you cannot add a function without a default platform.
+    let code = &[0xa1, 0xfa, 0xf8, 0xf0, 0x99, 0x83, 0xc0, 0x37, 0xc3];
+    let view = BinaryView::from_data(&FileMetadata::new(), code).expect("Failed to create view");
+    assert!(view.add_user_function(0).is_none());
+    assert!(view.add_auto_function(0).is_none());
+
+    // Now set it to verify we can add it.
+    let platform = Platform::by_name("x86").expect("Failed to get platform");
+    assert!(view.add_user_function_with_platform(0, &platform).is_some());
+}