[WARP] Improve UX surrounding removal of matched functions

- Adds Python API to remove matched function
- Adds command + UI actions to remove matched function
- Adds command + UI actions to ignore function in subsequent matches

4 files changed, 88 insertions, 4 deletions

diff --git a/plugins/warp/src/lib.rs b/plugins/warp/src/lib.rs
index 50ad3889..9f40acef 100644
--- a/plugins/warp/src/lib.rs
+++ b/plugins/warp/src/lib.rs
@@ -58,6 +58,14 @@ fn get_warp_include_tag_type(view: &BinaryView) -> Ref<TagType> {
         .unwrap_or_else(|| view.create_tag_type(INCLUDE_TAG_NAME, INCLUDE_TAG_ICON))
 }
 
+const IGNORE_TAG_ICON: &str = "🧊";
+const IGNORE_TAG_NAME: &str = "WARP: Ignored Function";
+
+fn get_warp_ignore_tag_type(view: &BinaryView) -> Ref<TagType> {
+    view.tag_type_by_name(IGNORE_TAG_NAME)
+        .unwrap_or_else(|| view.create_tag_type(IGNORE_TAG_NAME, IGNORE_TAG_ICON))
+}
+
 pub fn core_signature_dir() -> PathBuf {
     // Get core signatures for the given platform
     let install_dir = binaryninja::install_directory();
diff --git a/plugins/warp/src/plugin.rs b/plugins/warp/src/plugin.rs
index 5746a9bc..c396ebca 100644
--- a/plugins/warp/src/plugin.rs
+++ b/plugins/warp/src/plugin.rs
@@ -216,6 +216,18 @@ pub extern "C" fn CorePluginInit() -> bool {
     );
 
     register_command_for_function(
+        "WARP\\Ignore Function",
+        "Add current function to the list of functions to ignore when matching",
+        function::IgnoreFunction {},
+    );
+
+    register_command_for_function(
+        "WARP\\Remove Matched Function",
+        "Remove the current match from the selected function, to prevent matches in future use 'Ignore Function'",
+        function::RemoveFunction {},
+    );
+
+    register_command_for_function(
         "WARP\\Copy GUID",
         "Copy the computed GUID for the function",
         function::CopyFunctionGUID {},
diff --git a/plugins/warp/src/plugin/function.rs b/plugins/warp/src/plugin/function.rs
index a4fd2c10..6202f9d7 100644
--- a/plugins/warp/src/plugin/function.rs
+++ b/plugins/warp/src/plugin/function.rs
@@ -1,9 +1,14 @@
-use crate::cache::{cached_function_guid, try_cached_function_guid};
-use crate::{get_warp_include_tag_type, INCLUDE_TAG_NAME};
+use crate::cache::{
+    cached_function_guid, insert_cached_function_match, try_cached_function_guid,
+    try_cached_function_match,
+};
+use crate::{
+    get_warp_ignore_tag_type, get_warp_include_tag_type, IGNORE_TAG_NAME, INCLUDE_TAG_NAME,
+};
 use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
 use binaryninja::command::{Command, FunctionCommand};
-use binaryninja::function::Function;
+use binaryninja::function::{Function, FunctionUpdateType};
 use binaryninja::rc::Guard;
 use rayon::iter::ParallelIterator;
 use std::thread;
@@ -43,6 +48,60 @@ impl FunctionCommand for IncludeFunction {
     }
 }
 
+pub struct IgnoreFunction;
+
+impl FunctionCommand for IgnoreFunction {
+    fn action(&self, view: &BinaryView, func: &Function) {
+        let sym_name = func.symbol().short_name();
+        let sym_name_str = sym_name.to_string_lossy();
+        let should_add_tag = func.function_tags(None, Some(IGNORE_TAG_NAME)).is_empty();
+        let ignore_tag_type = get_warp_ignore_tag_type(view);
+        match should_add_tag {
+            true => {
+                log::info!(
+                    "Ignoring function for matching '{}' at 0x{:x}",
+                    sym_name_str,
+                    func.start()
+                );
+                func.add_tag(&ignore_tag_type, "", None, false, None);
+            }
+            false => {
+                log::info!(
+                    "Including function for matching '{}' at 0x{:x}",
+                    sym_name_str,
+                    func.start()
+                );
+                func.remove_tags_of_type(&ignore_tag_type, None, false, None);
+            }
+        }
+    }
+
+    fn valid(&self, _view: &BinaryView, _func: &Function) -> bool {
+        true
+    }
+}
+
+pub struct RemoveFunction;
+
+impl FunctionCommand for RemoveFunction {
+    fn action(&self, _view: &BinaryView, func: &Function) {
+        let sym_name = func.symbol().short_name();
+        let sym_name_str = sym_name.to_string_lossy();
+        log::info!(
+            "Removing matched function '{}' at 0x{:x}",
+            sym_name_str,
+            func.start()
+        );
+        insert_cached_function_match(func, None);
+        func.reanalyze(FunctionUpdateType::UserFunctionUpdate);
+    }
+
+    fn valid(&self, _view: &BinaryView, func: &Function) -> bool {
+        // Only allow if the function actually has a match.
+        try_cached_function_match(func).is_some()
+    }
+}
+
 pub struct CopyFunctionGUID;
 
 impl FunctionCommand for CopyFunctionGUID {
diff --git a/plugins/warp/src/plugin/workflow.rs b/plugins/warp/src/plugin/workflow.rs
index a1b0e202..1fca8008 100644
--- a/plugins/warp/src/plugin/workflow.rs
+++ b/plugins/warp/src/plugin/workflow.rs
@@ -6,7 +6,7 @@ use crate::cache::{
 use crate::convert::{platform_to_target, to_bn_type};
 use crate::matcher::{Matcher, MatcherSettings};
 use crate::plugin::settings::PluginSettings;
-use crate::{get_warp_tag_type, relocatable_regions};
+use crate::{get_warp_ignore_tag_type, get_warp_tag_type, relocatable_regions, IGNORE_TAG_NAME};
 use binaryninja::architecture::RegisterId;
 use binaryninja::background_task::BackgroundTask;
 use binaryninja::binary_view::{BinaryView, BinaryViewExt};
@@ -62,11 +62,15 @@ struct FunctionSet {
 impl FunctionSet {
     fn from_view(view: &BinaryView) -> Option<Self> {
         let mut set = FunctionSet::default();
+        let is_ignored_func =
+            |f: &BNFunction| !f.function_tags(None, Some(IGNORE_TAG_NAME)).is_empty();
 
         // TODO: Par iter this? Using dashmap
         set.functions_by_target_and_guid = view
             .functions()
             .iter()
+            // Skip functions that have the ignored tag! Otherwise, we will match on them.
+            .filter(|f| !is_ignored_func(f))
             .filter_map(|f| {
                 let guid = try_cached_function_guid(&f)?;
                 let target = platform_to_target(&f.platform());
@@ -110,6 +114,7 @@ pub fn run_matcher(view: &BinaryView) {
     // TODO: Create the tag type so we dont have UB in the apply function workflow.
     let undo_id = view.file().begin_undo_actions(false);
     let _ = get_warp_tag_type(view);
+    let _ = get_warp_ignore_tag_type(view);
     view.file().forget_undo_actions(&undo_id);
 
     // Then we want to actually find matching functions.