[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

11 files changed, 224 insertions, 16 deletions

diff --git a/plugins/warp/api/python/warp.py b/plugins/warp/api/python/warp.py
index 29537895..866d179b 100644
--- a/plugins/warp/api/python/warp.py
+++ b/plugins/warp/api/python/warp.py
@@ -195,6 +195,10 @@ class WarpFunction:
     def apply(self, function: Function):
         warpcore.BNWARPFunctionApply(self.handle, function.handle)
 
+    @staticmethod
+    def remove_matched(function: Function):
+        warpcore.BNWARPFunctionApply(None, function.handle)
+
 
 class WarpContainerSearchQuery:
     def __init__(self, query: str, offset: Optional[int] = None, limit: Optional[int] = None, source: Optional[Source] = None, source_tags: Optional[List[str]] = None):
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.
diff --git a/plugins/warp/ui/matched.cpp b/plugins/warp/ui/matched.cpp
index aff6779f..7112e74e 100644
--- a/plugins/warp/ui/matched.cpp
+++ b/plugins/warp/ui/matched.cpp
@@ -27,17 +27,18 @@ WarpMatchedWidget::WarpMatchedWidget(BinaryViewRef current)
 	m_tableWidget->setContentsMargins(0, 0, 0, 0);
 	m_splitter->addWidget(m_tableWidget);
 
-	// Toggle the applying workflow, this workflow sets all the data for the function based on the matched function
-	// data.
+	// Removes the match for the function, this is irreversible currently, and the user must run the matcher again.
+	// TODO: We previously were trying to instead toggle the application of the match, but because the symbols are applied
+	// TODO: when applying the match metadata we would persist that regardless.
 	m_tableWidget->RegisterContextMenuAction(
-		"Toggle Application", [this](WarpFunctionItem*, std::optional<uint64_t> address) {
+		"Remove Match", [this](WarpFunctionItem*, std::optional<uint64_t> address) {
 			if (!address.has_value())
 				return;
 			for (const auto& func : m_current->GetAnalysisFunctionsForAddress(*address))
 			{
-				const bool previous = BinaryNinja::Settings::Instance()->Get<bool>(WARP_APPLY_ACTIVITY, func);
-				BinaryNinja::Settings::Instance()->Set(WARP_APPLY_ACTIVITY, !previous, func);
-				func->Reanalyze();
+				WarpRemoveMatchDialog dlg(this, func);
+				if (dlg.execute())
+					Update();
 			}
 		});
 
@@ -74,10 +75,9 @@ void WarpMatchedWidget::Update()
 	for (const auto& analysisFunction : m_current->GetAnalysisFunctionList())
 	{
 		if (const auto& matchedFunction = Warp::Function::GetMatched(*analysisFunction))
-		{
-			uint64_t startAddress = analysisFunction->GetStart();
-			m_tableWidget->InsertFunction(startAddress, new WarpFunctionItem(matchedFunction, analysisFunction));
-		}
+			m_tableWidget->InsertFunction(analysisFunction->GetStart(), new WarpFunctionItem(matchedFunction, analysisFunction));
+		else
+			m_tableWidget->RemoveFunction(analysisFunction->GetStart());
 	}
 	m_tableWidget->GetTableView()->setModel(m_tableWidget->GetProxyModel());
 	m_tableWidget->GetProxyModel()->setSourceModel(m_tableWidget->GetModel());
diff --git a/plugins/warp/ui/matches.cpp b/plugins/warp/ui/matches.cpp
index 5ec3da1a..4b8f0495 100644
--- a/plugins/warp/ui/matches.cpp
+++ b/plugins/warp/ui/matches.cpp
@@ -56,6 +56,24 @@ WarpCurrentFunctionWidget::WarpCurrentFunctionWidget()
 		// So it shows visually as selected.
 		m_tableWidget->GetModel()->SetMatchedFunction(selectedFunction);
 	});
+	// If the selected function is the current match, let the user remove the match.
+	m_tableWidget->RegisterContextMenuAction("Remove Match",
+		[this](WarpFunctionItem*, std::optional<uint64_t>) {
+			WarpRemoveMatchDialog dlg(this, m_current);
+			if (dlg.execute())
+				m_tableWidget->GetModel()->SetMatchedFunction(nullptr);
+		},
+		[this](WarpFunctionItem* item, std::optional<uint64_t>) {
+		if (item == nullptr)
+			return false;
+		Warp::Ref<Warp::Function> selectedFunction = item->GetFunction();
+		if (!selectedFunction)
+			return false;
+		Warp::Ref<Warp::Function> matchedFunction = m_tableWidget->GetModel()->GetMatchedFunction();
+		if (!matchedFunction)
+			return false;
+		return BNWARPFunctionsEqual(selectedFunction->m_object, matchedFunction->m_object);
+	});
 	m_tableWidget->RegisterContextMenuAction(
 		"Search for Source", [this](WarpFunctionItem* item, std::optional<uint64_t>) {
 			// Apply the source as the filter.
@@ -79,7 +97,6 @@ WarpCurrentFunctionWidget::WarpCurrentFunctionWidget()
 		m_infoWidget->UpdateInfo();
 	});
 
-
 	connect(m_tableWidget->GetTableView(), &QTableView::doubleClicked, this, [=, this](const QModelIndex& index) {
 		if (m_current == nullptr)
 			return;
diff --git a/plugins/warp/ui/shared/function.cpp b/plugins/warp/ui/shared/function.cpp
index e177c285..06d4cbcf 100644
--- a/plugins/warp/ui/shared/function.cpp
+++ b/plugins/warp/ui/shared/function.cpp
@@ -70,6 +70,15 @@ void WarpFunctionItemModel::InsertFunction(uint64_t address, WarpFunctionItem* i
 	m_insertableFunctionRows[address] = rowCount() - 1;
 }
 
+void WarpFunctionItemModel::RemoveFunction(uint64_t address)
+{
+	const auto iter = m_insertableFunctionRows.find(address);
+	if (iter == m_insertableFunctionRows.end())
+		return;
+	removeRow(iter->second);
+	m_insertableFunctionRows.erase(iter);
+}
+
 WarpFunctionItem* WarpFunctionItemModel::GetItem(const QModelIndex& index) const
 {
 	if (!index.isValid())
@@ -173,7 +182,7 @@ WarpFunctionTableWidget::WarpFunctionTableWidget(QWidget* parent) : QWidget(pare
 	m_table->setSelectionBehavior(QAbstractItemView::SelectRows);
 	m_table->setSelectionMode(QAbstractItemView::SingleSelection);
 	m_table->setEditTriggers(QAbstractItemView::NoEditTriggers);
-	m_table->setFocusPolicy(Qt::NoFocus);
+	m_table->setFocusPolicy(Qt::FocusPolicy::StrongFocus);
 	m_table->setShowGrid(false);
 	m_table->setAlternatingRowColors(false);
 	m_table->setSortingEnabled(true);
@@ -219,6 +228,15 @@ WarpFunctionTableWidget::WarpFunctionTableWidget(QWidget* parent) : QWidget(pare
 		if (!item || !item->GetFunction())
 			return;
 
+		for (QAction* action : m_contextMenu->actions())
+		{
+			bool enabled = true;
+			auto iter = m_contextMenuIsValid.find(action->text());
+			if (iter != m_contextMenuIsValid.end())
+				enabled = iter->second(item, m_model->GetAddress(sourceIndex));
+			action->setEnabled(enabled);
+		}
+
 		// Execute the menu and get the selected action
 		const QAction* selectedAction = m_contextMenu->exec(m_table->viewport()->mapToGlobal(pos));
 		if (!selectedAction)
@@ -238,6 +256,16 @@ void WarpFunctionTableWidget::RegisterContextMenuAction(
 	m_contextMenuActions[name] = callback;
 }
 
+void WarpFunctionTableWidget::RegisterContextMenuAction(
+	const QString& name,
+	const std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>& callback,
+	const std::function<bool(WarpFunctionItem*, std::optional<uint64_t>)>& isValid)
+{
+	// Reuse existing registration then add optional validator
+	RegisterContextMenuAction(name, callback);
+	m_contextMenuIsValid[name] = isValid;
+}
+
 void WarpFunctionTableWidget::SetFunctions(QVector<WarpFunctionItem*> functions)
 {
 	// Clear matches as they are no longer valid.
@@ -266,6 +294,11 @@ void WarpFunctionTableWidget::InsertFunction(uint64_t address, WarpFunctionItem*
 	m_model->InsertFunction(address, function);
 }
 
+void WarpFunctionTableWidget::RemoveFunction(uint64_t address)
+{
+	m_model->RemoveFunction(address);
+}
+
 void WarpFunctionTableWidget::setFilter(const std::string& filter)
 {
 	m_proxyModel->setFilterFixedString(QString::fromStdString(filter));
diff --git a/plugins/warp/ui/shared/function.h b/plugins/warp/ui/shared/function.h
index 5cfaff75..6281944d 100644
--- a/plugins/warp/ui/shared/function.h
+++ b/plugins/warp/ui/shared/function.h
@@ -49,6 +49,8 @@ public:
 
 	void InsertFunction(uint64_t address, WarpFunctionItem* item);
 
+	void RemoveFunction(uint64_t address);
+
 	WarpFunctionItem* GetItem(const QModelIndex& index) const;
 
 	std::optional<uint64_t> GetAddress(const QModelIndex& index) const;
@@ -68,6 +70,8 @@ public:
 			emit dataChanged(topLeft, bottomRight);
 		}
 	}
+
+	Warp::Ref<Warp::Function> GetMatchedFunction() const { return m_matchedFunction; }
 };
 
 class WarpFunctionFilterModel : public QSortFilterProxyModel
@@ -95,6 +99,7 @@ class WarpFunctionTableWidget : public QWidget, public FilterTarget
 	FilteredView* m_filterView;
 	QMenu* m_contextMenu;
 	std::map<QString, std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>> m_contextMenuActions;
+	std::map<QString, std::function<bool(WarpFunctionItem*, std::optional<uint64_t>)>> m_contextMenuIsValid;
 
 public:
 	explicit WarpFunctionTableWidget(QWidget* parent = nullptr);
@@ -107,10 +112,17 @@ public:
 	void RegisterContextMenuAction(
 		const QString& name, const std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>& callback);
 
+	void RegisterContextMenuAction(
+		const QString &name,
+		const std::function<void(WarpFunctionItem *, std::optional<uint64_t>)> &callback,
+		const std::function<bool(WarpFunctionItem *, std::optional<uint64_t>)> &isValid);
+
 	void SetFunctions(QVector<WarpFunctionItem*> functions);
 
 	void InsertFunction(uint64_t address, WarpFunctionItem* function);
 
+	void RemoveFunction(uint64_t address);
+
 	void setFilter(const std::string&) override;
 
 	void scrollToFirstItem() override {}
diff --git a/plugins/warp/ui/shared/misc.cpp b/plugins/warp/ui/shared/misc.cpp
index f84237c9..6de314b1 100644
--- a/plugins/warp/ui/shared/misc.cpp
+++ b/plugins/warp/ui/shared/misc.cpp
@@ -1,5 +1,6 @@
 #include "misc.h"
 
+#include <QDialogButtonBox>
 #include <QGridLayout>
 #include <QHeaderView>
 
@@ -140,3 +141,44 @@ ParsedQuery::ParsedQuery(const QString& rawQuery)
 	// Normalize whitespace
 	query = query.simplified();
 }
+
+WarpRemoveMatchDialog::WarpRemoveMatchDialog(QWidget *parent, FunctionRef func) : QDialog(parent), m_func(func)
+{
+	setWindowTitle("Remove Matching Function");
+	setModal(true);
+
+	auto* vbox = new QVBoxLayout(this);
+	auto* text = new QLabel("Remove the match for this function? You can also mark it as ignored to prevent future automatic matches.");
+	text->setWordWrap(true);
+	vbox->addWidget(text);
+
+	m_ignoreCheck = new QCheckBox("Tag function as ignored");
+	m_ignoreCheck->setChecked(true);
+	vbox->addWidget(m_ignoreCheck);
+
+	auto* buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
+	connect(buttons, &QDialogButtonBox::accepted, this, &QDialog::accept);
+	connect(buttons, &QDialogButtonBox::rejected, this, &QDialog::reject);
+	vbox->addWidget(buttons);
+}
+
+bool WarpRemoveMatchDialog::execute()
+{
+	if (!m_func)
+		return false;
+	if (exec() != QDialog::Accepted)
+		return false;
+	Warp::Function::RemoveMatch(*m_func);
+	if (m_ignoreCheck->isChecked())
+	{
+		// TODO: For now we just assume the tag type to exist (the matcher activity will create it)
+		const TagTypeRef tagType = m_func->GetView()->GetTagTypeByName("WARP: Ignored Function");
+		if (!tagType)
+			return false;
+		const TagRef tag = new BinaryNinja::Tag(tagType, "");
+		if (tagType)
+			m_func->AddUserFunctionTag(tag);
+	}
+	m_func->Reanalyze();
+	return true;
+}
diff --git a/plugins/warp/ui/shared/misc.h b/plugins/warp/ui/shared/misc.h
index 9d187594..97adcab0 100644
--- a/plugins/warp/ui/shared/misc.h
+++ b/plugins/warp/ui/shared/misc.h
@@ -1,4 +1,6 @@
 #pragma once
+#include <QCheckBox>
+#include <QDialog>
 #include <qmetatype.h>
 #include <QSortFilterProxyModel>
 #include <qstandarditemmodel.h>
@@ -103,6 +105,20 @@ struct ParsedQuery
 	}
 };
 
+// TODO: Consolidate with `WARP\\Remove Matched Function` plugin command?
+class WarpRemoveMatchDialog : public QDialog
+{
+	Q_OBJECT
+public:
+	explicit WarpRemoveMatchDialog(QWidget* parent, FunctionRef func);
+
+	bool execute();
+
+private:
+	FunctionRef m_func;
+	QCheckBox* m_ignoreCheck{nullptr};
+};
+
 constexpr const char* ALLOWED_TAGS_SETTING = "warp.fetcher.allowedSourceTags";
 constexpr const char* BATCH_SIZE_SETTING = "warp.fetcher.fetchBatchSize";