Skip to content

Commit 40f7d40

Browse files
emesareemesare
committed
[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
1 parent 49558b8 commit 40f7d40

File tree

12 files changed

+232
-16
lines changed

12 files changed

+232
-16
lines changed

‎plugins/warp/api/python/warp.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,10 @@ def get_matched(function: Function) -> Optional['WarpFunction']:
195195
def apply(self, function: Function):
196196
warpcore.BNWARPFunctionApply(self.handle, function.handle)
197197

198+
@staticmethod
199+
def remove_matched(function: Function):
200+
warpcore.BNWARPFunctionApply(None, function.handle)
201+
198202

199203
class WarpContainerSearchQuery:
200204
def __init__(self, query: str, offset: Optional[int] = None, limit: Optional[int] = None, source: Optional[Source] = None, source_tags: Optional[List[str]] = None):

‎plugins/warp/src/lib.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,14 @@ fn get_warp_include_tag_type(view: &BinaryView) -> Ref<TagType> {
5858
.unwrap_or_else(|| view.create_tag_type(INCLUDE_TAG_NAME, INCLUDE_TAG_ICON))
5959
}
6060

61+
const IGNORE_TAG_ICON: &str = "🧊";
62+
const IGNORE_TAG_NAME: &str = "WARP: Ignored Function";
63+
64+
fn get_warp_ignore_tag_type(view: &BinaryView) -> Ref<TagType> {
65+
view.tag_type_by_name(IGNORE_TAG_NAME)
66+
.unwrap_or_else(|| view.create_tag_type(IGNORE_TAG_NAME, IGNORE_TAG_ICON))
67+
}
68+
6169
pub fn core_signature_dir() -> PathBuf {
6270
// Get core signatures for the given platform
6371
let install_dir = binaryninja::install_directory();

‎plugins/warp/src/plugin.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,18 @@ pub extern "C" fn CorePluginInit() -> bool {
215215
function::IncludeFunction {},
216216
);
217217

218+
register_command_for_function(
219+
"WARP\\Ignore Function",
220+
"Add current function to the list of functions to ignore when matching",
221+
function::IgnoreFunction {},
222+
);
223+
224+
register_command_for_function(
225+
"WARP\\Remove Matched Function",
226+
"Remove the current match from the selected function, to prevent matches in future use 'Ignore Function'",
227+
function::RemoveFunction {},
228+
);
229+
218230
register_command_for_function(
219231
"WARP\\Copy GUID",
220232
"Copy the computed GUID for the function",

‎plugins/warp/src/plugin/function.rs‎

Lines changed: 62 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
1-
use crate::cache::{cached_function_guid, try_cached_function_guid};
2-
use crate::{get_warp_include_tag_type, INCLUDE_TAG_NAME};
1+
use crate::cache::{
2+
cached_function_guid, insert_cached_function_match, try_cached_function_guid,
3+
try_cached_function_match,
4+
};
5+
use crate::{
6+
get_warp_ignore_tag_type, get_warp_include_tag_type, IGNORE_TAG_NAME, INCLUDE_TAG_NAME,
7+
};
38
use binaryninja::background_task::BackgroundTask;
49
use binaryninja::binary_view::{BinaryView, BinaryViewExt};
510
use binaryninja::command::{Command, FunctionCommand};
6-
use binaryninja::function::Function;
11+
use binaryninja::function::{Function, FunctionUpdateType};
712
use binaryninja::rc::Guard;
813
use rayon::iter::ParallelIterator;
914
use std::thread;
@@ -43,6 +48,60 @@ impl FunctionCommand for IncludeFunction {
4348
}
4449
}
4550

51+
pub struct IgnoreFunction;
52+
53+
impl FunctionCommand for IgnoreFunction {
54+
fn action(&self, view: &BinaryView, func: &Function) {
55+
let sym_name = func.symbol().short_name();
56+
let sym_name_str = sym_name.to_string_lossy();
57+
let should_add_tag = func.function_tags(None, Some(IGNORE_TAG_NAME)).is_empty();
58+
let ignore_tag_type = get_warp_ignore_tag_type(view);
59+
match should_add_tag {
60+
true => {
61+
log::info!(
62+
"Ignoring function for matching '{}' at 0x{:x}",
63+
sym_name_str,
64+
func.start()
65+
);
66+
func.add_tag(&ignore_tag_type, "", None, false, None);
67+
}
68+
false => {
69+
log::info!(
70+
"Including function for matching '{}' at 0x{:x}",
71+
sym_name_str,
72+
func.start()
73+
);
74+
func.remove_tags_of_type(&ignore_tag_type, None, false, None);
75+
}
76+
}
77+
}
78+
79+
fn valid(&self, _view: &BinaryView, _func: &Function) -> bool {
80+
true
81+
}
82+
}
83+
84+
pub struct RemoveFunction;
85+
86+
impl FunctionCommand for RemoveFunction {
87+
fn action(&self, _view: &BinaryView, func: &Function) {
88+
let sym_name = func.symbol().short_name();
89+
let sym_name_str = sym_name.to_string_lossy();
90+
log::info!(
91+
"Removing matched function '{}' at 0x{:x}",
92+
sym_name_str,
93+
func.start()
94+
);
95+
insert_cached_function_match(func, None);
96+
func.reanalyze(FunctionUpdateType::UserFunctionUpdate);
97+
}
98+
99+
fn valid(&self, _view: &BinaryView, func: &Function) -> bool {
100+
// Only allow if the function actually has a match.
101+
try_cached_function_match(func).is_some()
102+
}
103+
}
104+
46105
pub struct CopyFunctionGUID;
47106

48107
impl FunctionCommand for CopyFunctionGUID {

‎plugins/warp/src/plugin/workflow.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crate::cache::{
66
use crate::convert::{platform_to_target, to_bn_type};
77
use crate::matcher::{Matcher, MatcherSettings};
88
use crate::plugin::settings::PluginSettings;
9-
use crate::{get_warp_tag_type, relocatable_regions};
9+
use crate::{get_warp_ignore_tag_type, get_warp_tag_type, relocatable_regions, IGNORE_TAG_NAME};
1010
use binaryninja::architecture::RegisterId;
1111
use binaryninja::background_task::BackgroundTask;
1212
use binaryninja::binary_view::{BinaryView, BinaryViewExt};
@@ -62,11 +62,15 @@ struct FunctionSet {
6262
impl FunctionSet {
6363
fn from_view(view: &BinaryView) -> Option<Self> {
6464
let mut set = FunctionSet::default();
65+
let is_ignored_func =
66+
|f: &BNFunction| !f.function_tags(None, Some(IGNORE_TAG_NAME)).is_empty();
6567

6668
// TODO: Par iter this? Using dashmap
6769
set.functions_by_target_and_guid = view
6870
.functions()
6971
.iter()
72+
// Skip functions that have the ignored tag! Otherwise, we will match on them.
73+
.filter(|f| !is_ignored_func(f))
7074
.filter_map(|f| {
7175
let guid = try_cached_function_guid(&f)?;
7276
let target = platform_to_target(&f.platform());
@@ -110,6 +114,7 @@ pub fn run_matcher(view: &BinaryView) {
110114
// TODO: Create the tag type so we dont have UB in the apply function workflow.
111115
let undo_id = view.file().begin_undo_actions(false);
112116
let _ = get_warp_tag_type(view);
117+
let _ = get_warp_ignore_tag_type(view);
113118
view.file().forget_undo_actions(&undo_id);
114119

115120
// Then we want to actually find matching functions.

‎plugins/warp/ui/matched.cpp‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,17 +27,18 @@ WarpMatchedWidget::WarpMatchedWidget(BinaryViewRef current)
2727
m_tableWidget->setContentsMargins(0, 0, 0, 0);
2828
m_splitter->addWidget(m_tableWidget);
2929

30-
// Toggle the applying workflow, this workflow sets all the data for the function based on the matched function
31-
// data.
30+
// Removes the match for the function, this is irreversible currently, and the user must run the matcher again.
31+
// TODO: We previously were trying to instead toggle the application of the match, but because the symbols are applied
32+
// TODO: when applying the match metadata we would persist that regardless.
3233
m_tableWidget->RegisterContextMenuAction(
33-
"Toggle Application", [this](WarpFunctionItem*, std::optional<uint64_t> address) {
34+
"Remove Match", [this](WarpFunctionItem*, std::optional<uint64_t> address) {
3435
if (!address.has_value())
3536
return;
3637
for (const auto& func : m_current->GetAnalysisFunctionsForAddress(*address))
3738
{
38-
const bool previous = BinaryNinja::Settings::Instance()->Get<bool>(WARP_APPLY_ACTIVITY, func);
39-
BinaryNinja::Settings::Instance()->Set(WARP_APPLY_ACTIVITY, !previous, func);
40-
func->Reanalyze();
39+
WarpRemoveMatchDialog dlg(this, func);
40+
if (dlg.execute())
41+
Update();
4142
}
4243
});
4344

@@ -74,10 +75,9 @@ void WarpMatchedWidget::Update()
7475
for (const auto& analysisFunction : m_current->GetAnalysisFunctionList())
7576
{
7677
if (const auto& matchedFunction = Warp::Function::GetMatched(*analysisFunction))
77-
{
78-
uint64_t startAddress = analysisFunction->GetStart();
79-
m_tableWidget->InsertFunction(startAddress, new WarpFunctionItem(matchedFunction, analysisFunction));
80-
}
78+
m_tableWidget->InsertFunction(analysisFunction->GetStart(), new WarpFunctionItem(matchedFunction, analysisFunction));
79+
else
80+
m_tableWidget->RemoveFunction(analysisFunction->GetStart());
8181
}
8282
m_tableWidget->GetTableView()->setModel(m_tableWidget->GetProxyModel());
8383
m_tableWidget->GetProxyModel()->setSourceModel(m_tableWidget->GetModel());

‎plugins/warp/ui/matches.cpp‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,24 @@ WarpCurrentFunctionWidget::WarpCurrentFunctionWidget()
5656
// So it shows visually as selected.
5757
m_tableWidget->GetModel()->SetMatchedFunction(selectedFunction);
5858
});
59+
// If the selected function is the current match, let the user remove the match.
60+
m_tableWidget->RegisterContextMenuAction("Remove Match",
61+
[this](WarpFunctionItem*, std::optional<uint64_t>) {
62+
WarpRemoveMatchDialog dlg(this, m_current);
63+
if (dlg.execute())
64+
m_tableWidget->GetModel()->SetMatchedFunction(nullptr);
65+
},
66+
[this](WarpFunctionItem* item, std::optional<uint64_t>) {
67+
if (item == nullptr)
68+
return false;
69+
Warp::Ref<Warp::Function> selectedFunction = item->GetFunction();
70+
if (!selectedFunction)
71+
return false;
72+
Warp::Ref<Warp::Function> matchedFunction = m_tableWidget->GetModel()->GetMatchedFunction();
73+
if (!matchedFunction)
74+
return false;
75+
return BNWARPFunctionsEqual(selectedFunction->m_object, matchedFunction->m_object);
76+
});
5977
m_tableWidget->RegisterContextMenuAction(
6078
"Search for Source", [this](WarpFunctionItem* item, std::optional<uint64_t>) {
6179
// Apply the source as the filter.
@@ -79,7 +97,6 @@ WarpCurrentFunctionWidget::WarpCurrentFunctionWidget()
7997
m_infoWidget->UpdateInfo();
8098
});
8199

82-
83100
connect(m_tableWidget->GetTableView(), &QTableView::doubleClicked, this, [=, this](const QModelIndex& index) {
84101
if (m_current == nullptr)
85102
return;

‎plugins/warp/ui/shared/function.cpp‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,15 @@ void WarpFunctionItemModel::InsertFunction(uint64_t address, WarpFunctionItem* i
7070
m_insertableFunctionRows[address] = rowCount() - 1;
7171
}
7272

73+
void WarpFunctionItemModel::RemoveFunction(uint64_t address)
74+
{
75+
const auto iter = m_insertableFunctionRows.find(address);
76+
if (iter == m_insertableFunctionRows.end())
77+
return;
78+
removeRow(iter->second);
79+
m_insertableFunctionRows.erase(iter);
80+
}
81+
7382
WarpFunctionItem* WarpFunctionItemModel::GetItem(const QModelIndex& index) const
7483
{
7584
if (!index.isValid())
@@ -173,7 +182,7 @@ WarpFunctionTableWidget::WarpFunctionTableWidget(QWidget* parent) : QWidget(pare
173182
m_table->setSelectionBehavior(QAbstractItemView::SelectRows);
174183
m_table->setSelectionMode(QAbstractItemView::SingleSelection);
175184
m_table->setEditTriggers(QAbstractItemView::NoEditTriggers);
176-
m_table->setFocusPolicy(Qt::NoFocus);
185+
m_table->setFocusPolicy(Qt::FocusPolicy::StrongFocus);
177186
m_table->setShowGrid(false);
178187
m_table->setAlternatingRowColors(false);
179188
m_table->setSortingEnabled(true);
@@ -219,6 +228,15 @@ WarpFunctionTableWidget::WarpFunctionTableWidget(QWidget* parent) : QWidget(pare
219228
if (!item || !item->GetFunction())
220229
return;
221230

231+
for (QAction* action : m_contextMenu->actions())
232+
{
233+
bool enabled = true;
234+
auto iter = m_contextMenuIsValid.find(action->text());
235+
if (iter != m_contextMenuIsValid.end())
236+
enabled = iter->second(item, m_model->GetAddress(sourceIndex));
237+
action->setEnabled(enabled);
238+
}
239+
222240
// Execute the menu and get the selected action
223241
const QAction* selectedAction = m_contextMenu->exec(m_table->viewport()->mapToGlobal(pos));
224242
if (!selectedAction)
@@ -238,6 +256,16 @@ void WarpFunctionTableWidget::RegisterContextMenuAction(
238256
m_contextMenuActions[name] = callback;
239257
}
240258

259+
void WarpFunctionTableWidget::RegisterContextMenuAction(
260+
const QString& name,
261+
const std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>& callback,
262+
const std::function<bool(WarpFunctionItem*, std::optional<uint64_t>)>& isValid)
263+
{
264+
// Reuse existing registration then add optional validator
265+
RegisterContextMenuAction(name, callback);
266+
m_contextMenuIsValid[name] = isValid;
267+
}
268+
241269
void WarpFunctionTableWidget::SetFunctions(QVector<WarpFunctionItem*> functions)
242270
{
243271
// Clear matches as they are no longer valid.
@@ -266,6 +294,11 @@ void WarpFunctionTableWidget::InsertFunction(uint64_t address, WarpFunctionItem*
266294
m_model->InsertFunction(address, function);
267295
}
268296

297+
void WarpFunctionTableWidget::RemoveFunction(uint64_t address)
298+
{
299+
m_model->RemoveFunction(address);
300+
}
301+
269302
void WarpFunctionTableWidget::setFilter(const std::string& filter)
270303
{
271304
m_proxyModel->setFilterFixedString(QString::fromStdString(filter));

‎plugins/warp/ui/shared/function.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ class WarpFunctionItemModel : public QStandardItemModel
4949

5050
void InsertFunction(uint64_t address, WarpFunctionItem* item);
5151

52+
void RemoveFunction(uint64_t address);
53+
5254
WarpFunctionItem* GetItem(const QModelIndex& index) const;
5355

5456
std::optional<uint64_t> GetAddress(const QModelIndex& index) const;
@@ -68,6 +70,8 @@ class WarpFunctionItemModel : public QStandardItemModel
6870
emit dataChanged(topLeft, bottomRight);
6971
}
7072
}
73+
74+
Warp::Ref<Warp::Function> GetMatchedFunction() const { return m_matchedFunction; }
7175
};
7276

7377
class WarpFunctionFilterModel : public QSortFilterProxyModel
@@ -95,6 +99,7 @@ class WarpFunctionTableWidget : public QWidget, public FilterTarget
9599
FilteredView* m_filterView;
96100
QMenu* m_contextMenu;
97101
std::map<QString, std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>> m_contextMenuActions;
102+
std::map<QString, std::function<bool(WarpFunctionItem*, std::optional<uint64_t>)>> m_contextMenuIsValid;
98103

99104
public:
100105
explicit WarpFunctionTableWidget(QWidget* parent = nullptr);
@@ -107,10 +112,17 @@ class WarpFunctionTableWidget : public QWidget, public FilterTarget
107112
void RegisterContextMenuAction(
108113
const QString& name, const std::function<void(WarpFunctionItem*, std::optional<uint64_t>)>& callback);
109114

115+
void RegisterContextMenuAction(
116+
const QString &name,
117+
const std::function<void(WarpFunctionItem *, std::optional<uint64_t>)> &callback,
118+
const std::function<bool(WarpFunctionItem *, std::optional<uint64_t>)> &isValid);
119+
110120
void SetFunctions(QVector<WarpFunctionItem*> functions);
111121

112122
void InsertFunction(uint64_t address, WarpFunctionItem* function);
113123

124+
void RemoveFunction(uint64_t address);
125+
114126
void setFilter(const std::string&) override;
115127

116128
void scrollToFirstItem() override {}

‎plugins/warp/ui/shared/misc.cpp‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "misc.h"
22

3+
#include <QDialogButtonBox>
34
#include <QGridLayout>
45
#include <QHeaderView>
56

@@ -140,3 +141,44 @@ ParsedQuery::ParsedQuery(const QString& rawQuery)
140141
// Normalize whitespace
141142
query = query.simplified();
142143
}
144+
145+
WarpRemoveMatchDialog::WarpRemoveMatchDialog(QWidget *parent, FunctionRef func) : QDialog(parent), m_func(func)
146+
{
147+
setWindowTitle("Remove Matching Function");
148+
setModal(true);
149+
150+
auto* vbox = new QVBoxLayout(this);
151+
auto* text = new QLabel("Remove the match for this function? You can also mark it as ignored to prevent future automatic matches.");
152+
text->setWordWrap(true);
153+
vbox->addWidget(text);
154+
155+
m_ignoreCheck = new QCheckBox("Tag function as ignored");
156+
m_ignoreCheck->setChecked(true);
157+
vbox->addWidget(m_ignoreCheck);
158+
159+
auto* buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
160+
connect(buttons, &QDialogButtonBox::accepted, this, &QDialog::accept);
161+
connect(buttons, &QDialogButtonBox::rejected, this, &QDialog::reject);
162+
vbox->addWidget(buttons);
163+
}
164+
165+
bool WarpRemoveMatchDialog::execute()
166+
{
167+
if (!m_func)
168+
return false;
169+
if (exec() != QDialog::Accepted)
170+
return false;
171+
Warp::Function::RemoveMatch(*m_func);
172+
if (m_ignoreCheck->isChecked())
173+
{
174+
// TODO: For now we just assume the tag type to exist (the matcher activity will create it)
175+
const TagTypeRef tagType = m_func->GetView()->GetTagTypeByName("WARP: Ignored Function");
176+
if (!tagType)
177+
return false;
178+
const TagRef tag = new BinaryNinja::Tag(tagType, "");
179+
if (tagType)
180+
m_func->AddUserFunctionTag(tag);
181+
}
182+
m_func->Reanalyze();
183+
return true;
184+
}

0 commit comments

Comments
 (0)