[WARP] Improved fetch dialog UX

- Respects views fetch batch size and allowed tags
- Saves and loads from the view settings instead of qt settings

7 files changed, 58 insertions, 83 deletions

diff --git a/plugins/warp/ui/matches.cpp b/plugins/warp/ui/matches.cpp
index 2345a7ea..b03d72c6 100644
--- a/plugins/warp/ui/matches.cpp
+++ b/plugins/warp/ui/matches.cpp
@@ -110,13 +110,6 @@ WarpCurrentFunctionWidget::WarpCurrentFunctionWidget()
 void WarpCurrentFunctionWidget::SetFetcher(std::shared_ptr<WarpFetcher> fetcher)
 {
     m_fetcher = fetcher;
-    // TODO: We need to remove the completion callback from the previously set fetcher.
-    m_fetcher->AddCompletionCallback([this]() {
-        // TODO: This is a little bit underspecified, we may end up updating more than strictly necessary.
-        // Once this function has been fetched, we need to update the matches for this widget.
-        UpdateMatches();
-        return KeepCallback;
-    });
 }
 
 void WarpCurrentFunctionWidget::SetCurrentFunction(FunctionRef current)
@@ -136,7 +129,8 @@ void WarpCurrentFunctionWidget::SetCurrentFunction(FunctionRef current)
         {
             BinaryNinja::WorkerPriorityEnqueue([this]() {
                 BinaryNinja::Ref bgTask = new BinaryNinja::BackgroundTask("Fetching WARP Functions...", true);
-                m_fetcher->FetchPendingFunctions();
+                const auto allowedTags = GetAllowedTagsFromView(m_current->GetView());
+                m_fetcher->FetchPendingFunctions(allowedTags);
                 bgTask->Finish();
             });
         }
diff --git a/plugins/warp/ui/plugin.cpp b/plugins/warp/ui/plugin.cpp
index d4860b91..b1cc9a5b 100644
--- a/plugins/warp/ui/plugin.cpp
+++ b/plugins/warp/ui/plugin.cpp
@@ -163,6 +163,9 @@ WarpSidebarWidget::WarpSidebarWidget(BinaryViewRef data) : SidebarWidget("WARP")
 	tabWidget->addTab(matchedFrame, "Matched Functions");
 	tabWidget->addTab(containerFrame, "Containers");
 
+	layout->addWidget(tabWidget);
+	this->setLayout(layout);
+
 	// Do a full update if analysis has been done, otherwise we may persist old data and not have new data.
 	m_analysisEvent = new AnalysisCompletionEvent(m_data, [this]() {
 		ExecuteOnMainThread([this]() {
@@ -170,15 +173,12 @@ WarpSidebarWidget::WarpSidebarWidget(BinaryViewRef data) : SidebarWidget("WARP")
 		});
 	});
 
-	std::shared_ptr<WarpFetcher> fetcher = WarpFetcher::Global();
+	const std::shared_ptr<WarpFetcher> fetcher = WarpFetcher::Global();
 	fetcher->AddCompletionCallback([this]() {
 		Update();
 		return KeepCallback;
 	});
 
-	layout->addWidget(tabWidget);
-	this->setLayout(layout);
-
 	// NOTE: This fetcher is shared with the fetch dialog that is constructed on initialization of this plugin.
 	m_currentFunctionWidget->SetFetcher(fetcher);
 }
diff --git a/plugins/warp/ui/shared/fetchdialog.cpp b/plugins/warp/ui/shared/fetchdialog.cpp
index 4b72e840..1d18aa92 100644
--- a/plugins/warp/ui/shared/fetchdialog.cpp
+++ b/plugins/warp/ui/shared/fetchdialog.cpp
@@ -7,6 +7,7 @@
 
 #include "action.h"
 #include "fetcher.h"
+#include "misc.h"
 
 using namespace BinaryNinja;
 
@@ -35,8 +36,6 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv,
     for (const auto &c: m_containers)
         m_containerCombo->addItem(QString::fromStdString(c->GetName()));
 
-    // TODO: Need to add tooltip to explain that a source must have atleast one of these tags to be considered.
-
     // Tags editor
     m_tagsList = new QListWidget(this);
     m_addTagBtn = new QPushButton(this);
@@ -67,13 +66,13 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv,
     m_tagsList->setToolTip("A source must have atleast ONE of these tags to be considered");
 
     // Defaults from processor tags
-    for (const auto &t: m_fetchProcessor->GetTags())
+    for (const auto &t: GetAllowedTagsFromView(m_bv))
         AddListItem(m_tagsList, QString::fromStdString(t));
 
     // Batch size and matcher checkbox
     m_batchSize = new QSpinBox(this);
     m_batchSize->setRange(10, 1000);
-    m_batchSize->setValue(100);
+    m_batchSize->setValue(GetBatchSizeFromView(m_bv));
     m_batchSize->setToolTip("Number of functions to fetch in each batch");
 
     m_rerunMatcher = new QCheckBox("Re-run matcher after fetch", this);
@@ -84,9 +83,6 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv,
     m_clearProcessed->setChecked(false);
 
     form->addRow(new QLabel("Container: "), m_containerCombo);
-    // TODO: Need to plumb this through to the fetcher, and also likely have a blacklisted or whitelist mode for this dialog.
-    // TODO: Alos wan to prefill the list of sources from the view/global settings.
-    // form->addRow(new QLabel("Allowed Sources: "), srcWrapper);
     form->addRow(new QLabel("Allowed Tags: "), tagWrapper);
     form->addRow(new QLabel("Batch Size: "), m_batchSize);
     form->addRow(m_rerunMatcher);
@@ -94,7 +90,7 @@ WarpFetchDialog::WarpFetchDialog(BinaryViewRef bv,
 
     auto buttons = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel, this);
     connect(buttons, &QDialogButtonBox::accepted, this, &WarpFetchDialog::onAccept);
-    connect(buttons, &QDialogButtonBox::rejected, this, &QDialog::reject);
+    connect(buttons, &QDialogButtonBox::rejected, this, &WarpFetchDialog::onReject);
 
     auto root = new QVBoxLayout(this);
     root->addLayout(form);
@@ -149,12 +145,12 @@ void WarpFetchDialog::onAccept()
     if (idx > 0) // 0 == All Containers
         containerIndex = static_cast<size_t>(idx - 1);
 
-    auto tags = collectTags();
     const auto batch = static_cast<size_t>(m_batchSize->value());
     const bool rerun = m_rerunMatcher->isChecked();
 
-    // Persist tags to the shared processor for consistency across navigation
-    m_fetchProcessor->SetTags(tags);
+    const auto tags = collectTags();
+    // Persist tags to the view settings.
+    SetTagsToView(m_bv, tags);
 
     if (m_clearProcessed->isChecked())
         m_fetchProcessor->ClearProcessed();
@@ -165,8 +161,16 @@ void WarpFetchDialog::onAccept()
     accept();
 }
 
+void WarpFetchDialog::onReject()
+{
+    const auto tags = collectTags();
+    // Persist tags to the view settings.
+    SetTagsToView(m_bv, tags);
+    reject();
+}
+
 void WarpFetchDialog::runBatchedFetch(const std::optional<size_t> &containerIndex,
-                                      const std::vector<Warp::SourceTag> &tags,
+                                      const std::vector<Warp::SourceTag> &allowedTags,
                                       size_t batchSize,
                                       bool rerunMatcher)
 {
@@ -186,7 +190,7 @@ void WarpFetchDialog::runBatchedFetch(const std::optional<size_t> &containerInde
     auto bv = m_bv;
 
     // TODO: Too many captures in this thing lol.
-    WorkerInteractiveEnqueue([fetcher, bv, funcs = std::move(funcs), batchSize, rerunMatcher, task]() mutable {
+    WorkerInteractiveEnqueue([fetcher, bv, funcs = std::move(funcs), batchSize, rerunMatcher, task, allowedTags]() mutable {
         size_t processed = 0;
         size_t batchIndex = 0;
 
@@ -198,7 +202,7 @@ void WarpFetchDialog::runBatchedFetch(const std::optional<size_t> &containerInde
             for (size_t i = 0; i < thisBatchCount; ++i)
                 fetcher->AddPendingFunction(funcs[processed + i]);
 
-            fetcher->FetchPendingFunctions();
+            fetcher->FetchPendingFunctions(allowedTags);
 
             ++batchIndex;
             processed += thisBatchCount;
@@ -219,8 +223,6 @@ void RegisterWarpFetchFunctionsCommand()
 {
     // Register a UI action and bind it globally. Add it to the Tools menu.
     const QString actionName = "WARP\\Fetch";
-
-    // TODO: Because we register this in every widget this will happen, this is bad behavior!
     if (!UIAction::isActionRegistered(actionName))
         UIAction::registerAction(actionName);
 
diff --git a/plugins/warp/ui/shared/fetchdialog.h b/plugins/warp/ui/shared/fetchdialog.h
index 37bfa0bb..a99eef19 100644
--- a/plugins/warp/ui/shared/fetchdialog.h
+++ b/plugins/warp/ui/shared/fetchdialog.h
@@ -45,13 +45,15 @@ private slots:
 
     void onAccept();
 
+    void onReject();
+
 private:
     void populateContainers();
 
     std::vector<Warp::SourceTag> collectTags() const;
 
     void runBatchedFetch(const std::optional<size_t> &containerIndex,
-                         const std::vector<Warp::SourceTag> &tags,
+                         const std::vector<Warp::SourceTag> &allowedTags,
                          size_t batchSize,
                          bool rerunMatcher);
 };
diff --git a/plugins/warp/ui/shared/fetcher.cpp b/plugins/warp/ui/shared/fetcher.cpp
index 767932e3..299050a5 100644
--- a/plugins/warp/ui/shared/fetcher.cpp
+++ b/plugins/warp/ui/shared/fetcher.cpp
@@ -1,26 +1,8 @@
 #include "fetcher.h"
 
-#include <QSettings>
-
 WarpFetcher::WarpFetcher()
 {
     m_logger = new BinaryNinja::Logger("WARP Fetcher");
-    QSettings qtSettings;
-    const QString key = "warp/allowedTags";
-
-    QStringList tags = qtSettings.value(key).toStringList();
-    if (tags.isEmpty()) {
-        tags = QStringList{ "official", "trusted" };
-        qtSettings.setValue(key, tags);
-        qtSettings.sync();
-    }
-
-    std::vector<Warp::SourceTag> initialTags;
-    initialTags.reserve(tags.size());
-    for (const auto& t : tags)
-        initialTags.emplace_back(t.trimmed().toStdString());
-
-    SetTags(initialTags);
 }
 
 void WarpFetcher::AddPendingFunction(const FunctionRef &func)
@@ -47,7 +29,7 @@ void WarpFetcher::ExecuteCompletionCallback()
         std::lock_guard<std::mutex> lock(m_requestMutex);
         m_completionCallbacks.erase(
             std::ranges::remove_if(m_completionCallbacks,
-                                   [](const auto &cb) { return cb() != RemoveCallback; }).begin(),
+                                   [](const auto &cb) { return cb() == RemoveCallback; }).begin(),
             m_completionCallbacks.end());
     });
 }
@@ -58,7 +40,7 @@ std::shared_ptr<WarpFetcher> WarpFetcher::Global()
     return global;
 }
 
-void WarpFetcher::FetchPendingFunctions()
+void WarpFetcher::FetchPendingFunctions(const std::vector<Warp::SourceTag>& allowedTags)
 {
     m_requestInProgress = true;
     const auto requests = FlushPendingFunctions();
@@ -82,13 +64,12 @@ void WarpFetcher::FetchPendingFunctions()
         platformMappedGuids[platform].push_back(guid.value());
     }
 
-    const auto tags = GetTags();
     for (const auto &[platform, guids] : platformMappedGuids)
     {
         m_logger->LogDebugF("Fetching {} functions for platform {}", guids.size(), platform->GetName());
         auto target = Warp::Target::FromPlatform(*platform);
         for (const auto &container: Warp::Container::All())
-            container->FetchFunctions(*target, guids, tags);
+            container->FetchFunctions(*target, guids, allowedTags);
 
         std::lock_guard<std::mutex> lock(m_requestMutex);
         for (const auto &guid: guids)
diff --git a/plugins/warp/ui/shared/fetcher.h b/plugins/warp/ui/shared/fetcher.h
index a7195ac4..917ac912 100644
--- a/plugins/warp/ui/shared/fetcher.h
+++ b/plugins/warp/ui/shared/fetcher.h
@@ -5,7 +5,6 @@
 #include <unordered_set>
 #include <vector>
 #include <functional>
-#include <QSettings>
 
 #include "warp.h"
 #include "binaryninjaapi.h"
@@ -24,7 +23,6 @@ class WarpFetcher
 
     std::mutex m_requestMutex;
     std::vector<FunctionRef> m_pendingRequests;
-    // TODO: Easy way to clear this if user wants to refetch.
     std::unordered_set<Warp::FunctionGUID> m_processedGuids;
 
     // List of callbacks to call when done fetching data, assume that others are using this as well.
@@ -38,35 +36,6 @@ public:
 
     std::atomic<bool> m_requestInProgress = false;
 
-    // Set the allowed source tags, sources with none of these tags will not be fetched from.
-    void SetTags(const std::vector<Warp::SourceTag> &tags)
-    {
-        std::lock_guard<std::mutex> lock(m_requestMutex);
-        // TODO: This is kinda a hack, the fetcher instance should not sync through qt settings!
-        QStringList qtTags = {};
-        for (const auto& t : tags)
-            qtTags.append(QString::fromStdString(t));
-        QSettings().setValue("warp/allowedTags", qtTags);
-    }
-
-    std::vector<Warp::SourceTag> GetTags() const
-    {
-        std::lock_guard<std::mutex> lock(const_cast<std::mutex &>(m_requestMutex));
-        // TODO: This is kinda a hack, the fetcher instance should not sync through qt settings!
-        QSettings qtSettings;
-        QStringList tags = qtSettings.value("warp/allowedTags").toStringList();
-        if (tags.isEmpty()) {
-            // The default tags to allow.
-            tags = QStringList{ "official", "trusted" };
-            qtSettings.setValue("warp/allowedTags", tags);
-            qtSettings.sync();
-        }
-        std::vector<Warp::SourceTag> initialTags = {};
-        for (const auto& t : tags)
-            initialTags.emplace_back(t.trimmed().toStdString());
-        return initialTags;
-    }
-
     void AddCompletionCallback(std::function<WarpFetchCompletionStatus()> cb)
     {
         std::lock_guard<std::mutex> lock(m_requestMutex);
@@ -75,7 +44,7 @@ public:
 
     void AddPendingFunction(const FunctionRef &func);
 
-    void FetchPendingFunctions();
+    void FetchPendingFunctions(const std::vector<Warp::SourceTag>& allowedTags);
 
     void ClearProcessed();
 private:
diff --git a/plugins/warp/ui/shared/misc.h b/plugins/warp/ui/shared/misc.h
index 35f640f5..57d71070 100644
--- a/plugins/warp/ui/shared/misc.h
+++ b/plugins/warp/ui/shared/misc.h
@@ -112,3 +112,30 @@ struct ParsedQuery
         return it.value();
     }
 };
+
+constexpr const char* ALLOWED_TAGS_SETTING = "warp.fetcher.allowedSourceTags";
+constexpr const char* BATCH_SIZE_SETTING = "warp.fetcher.fetchBatchSize";
+
+inline std::vector<Warp::SourceTag> GetAllowedTagsFromView(const BinaryViewRef& view)
+{
+    auto settings = BinaryNinja::Settings::Instance();
+    if (!settings->Contains(ALLOWED_TAGS_SETTING))
+        return {};
+    return settings->Get<std::vector<std::string>>(ALLOWED_TAGS_SETTING, view);
+}
+
+inline void SetTagsToView(const BinaryViewRef& view, const std::vector<Warp::SourceTag>& tags)
+{
+    auto settings = BinaryNinja::Settings::Instance();
+    if (!settings->Contains(ALLOWED_TAGS_SETTING))
+        return;
+    settings->Set(ALLOWED_TAGS_SETTING, tags, view);
+}
+
+inline int GetBatchSizeFromView(const BinaryViewRef& view)
+{
+    auto settings = BinaryNinja::Settings::Instance();
+    if (!settings->Contains(BATCH_SIZE_SETTING))
+        return 100;
+    return settings->Get<uint64_t>(BATCH_SIZE_SETTING, view);
+}
\ No newline at end of file