1 files changed, 124 insertions, 105 deletions

diff --git a/view/sharedcache/ui/SharedCacheUINotifications.cpp b/view/sharedcache/ui/SharedCacheUINotifications.cpp
index 9d5893f9..a7b190ff 100644
--- a/view/sharedcache/ui/SharedCacheUINotifications.cpp
+++ b/view/sharedcache/ui/SharedCacheUINotifications.cpp
@@ -3,14 +3,15 @@
 //
 
 #include "SharedCacheUINotifications.h"
-#include <QLayout>
 #include <sharedcacheapi.h>
 #include "ui/sidebar.h"
 #include "ui/linearview.h"
 #include "ui/viewframe.h"
 #include "dscpicker.h"
 #include "progresstask.h"
-#include "SharedCacheBDNotifications.h"
+
+using namespace BinaryNinja;
+using namespace SharedCacheAPI;
 
 UINotifications* UINotifications::m_instance = nullptr;
 
@@ -20,123 +21,141 @@ void UINotifications::init()
 	UIContext::registerNotification(m_instance);
 }
 
-
 void UINotifications::OnViewChange(UIContext* context, ViewFrame* frame, const QString& type)
 {
 	if (!frame)
 		return;
 
-	// FIXME there is a bv func for this
-	static std::function<bool(Ref<BinaryView>, uint64_t)> isAddrMapped = [](Ref<BinaryView> view, uint64_t addr) {
-		if (view && view->GetTypeName() == VIEW_NAME)
+	auto view = frame->getCurrentBinaryView();
+	if (!view || view->GetTypeName() != VIEW_NAME)
+		return;
+
+	auto viewInt = frame->getCurrentViewInterface();
+	if (!viewInt)
+		return;
+
+	auto ah = viewInt->actionHandler();
+	// Check to see if we have already bound these actions.
+	if (ah->isBoundAction("Load Image by Name"))
+		return;
+
+	static auto loadRegionAtAddr = [](BinaryView& view, uint64_t addr) {
+		auto controller = SharedCacheController::GetController(view);
+		if (!controller)
+			return;
+		if (auto foundRegion = controller->GetRegionContaining(addr))
 		{
-			for (const auto& seg : view->GetSegments())
-			{
-				if (seg->GetStart() <= addr && seg->GetEnd() > addr)
-					return true;
-			}
+			// If we did not load the region, then we don't need to run analysis.
+			if (!controller->ApplyRegion(view, *foundRegion))
+				return;
+			view.AddAnalysisOption("linearsweep");
+			view.UpdateAnalysis();
 		}
-		return false;
 	};
 
-	auto view = frame->getCurrentBinaryView();
-	if (view && view->GetTypeName() == VIEW_NAME)
-	{
-		if (auto viewInt = frame->getCurrentViewInterface())
+	static auto loadImageAtAddr = [](BinaryView& view, uint64_t addr) {
+		auto controller = SharedCacheController::GetController(view);
+		if (!controller)
+			return;
+		if (auto foundImage = controller->GetImageContaining(addr))
 		{
-			auto ah = viewInt->actionHandler();
-			if (!ah->isBoundAction("Load Image by Name"))
-			{
-				ah->bindAction("Load Image by Name", UIAction([view = view](const UIActionContext& ctx) {
-					DisplayDSCPicker(ctx.context, view);
-				}));
-				ah->bindAction("Load Section by Address", UIAction([view = view](const UIActionContext& ctx) {
-					Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(ctx.binaryView);
-					uint64_t addr = 0;
-					bool gotAddr = GetAddressInput(addr, "Address", "Address");
-					if (gotAddr)
-					{
-						BackgroundThread::create(ctx.context->mainWindow())->thenBackground(
-						[cache=cache, addr=addr]() {
-							cache->LoadSectionAtAddress(addr);
-						})->start();
-					}
-				}));
-				ah->bindAction("Load ADDRHERE",
-					UIAction(
-						[](const UIActionContext& ctx) {
-							Ref<BinaryView> view = ctx.binaryView;
-							Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(ctx.binaryView);
-							uint64_t addr = ctx.token.token.value;
-							if (addr)
-							{
-								BackgroundThread::create(ctx.context->mainWindow())->thenBackground(
-								[cache=cache, addr=addr]() {
-									cache->LoadSectionAtAddress(addr);
-								})->start();
-							}
-						},
-						[](const UIActionContext& ctx) {
-							Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(ctx.binaryView);
-							uint64_t addr = ctx.token.token.value;
-							if (isAddrMapped(ctx.binaryView, addr))
-								return false;
-							return addr && cache->GetNameForAddress(addr) != "";	 // bool
-						}));
-				ah->bindAction("Load IMGHERE",
-					UIAction(
-						[](const UIActionContext& ctx) {
-							Ref<BinaryView> view = ctx.binaryView;
-							Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(view);
-							uint64_t addr = ctx.token.token.value;
-							if (addr)
-							{
-								BackgroundThread::create(ctx.context->mainWindow())->thenBackground(
-								[cache=cache, addr=addr]() {
-									cache->LoadImageContainingAddress(addr);
-								})->start();
-							}
-						},
-						[](const UIActionContext& ctx) {
-							Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(ctx.binaryView);
-							uint64_t addr = ctx.token.token.value;
-							if (isAddrMapped(ctx.binaryView, addr))
-								return false;
-							return addr && cache->GetImageNameForAddress(addr) != "";  // bool
-						}));
-				ah->setActionDisplayName("Load ADDRHERE", [](const UIActionContext& ctx) {
-					Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(ctx.binaryView);
-					uint64_t addr = ctx.token.token.value;
-					if (addr)
-						return QString("Load ") + cache->GetNameForAddress(addr).c_str();
-					return QString("Error");
-				});
-				ah->setActionDisplayName("Load IMGHERE", [](const UIActionContext& ctx) {
-					Ref<SharedCacheAPI::SharedCache> cache = new SharedCacheAPI::SharedCache(ctx.binaryView);
-					uint64_t addr = ctx.token.token.value;
-					if (addr)
-						return QString("Load ") + cache->GetImageNameForAddress(addr).c_str();
-					return QString("Error");
-				});
-				if (auto linearView = qobject_cast<LinearView*>(viewInt->widget()))
-				{
-					linearView->contextMenu().addAction("Load ADDRHERE", VIEW_NAME);
-					linearView->contextMenu().addAction("Load IMGHERE", VIEW_NAME);
-					linearView->contextMenu().addAction("Load Image by Name", "DSCView2");
-					linearView->contextMenu().addAction("Load Section by Address", "DSCView2");
-					linearView->contextMenu().setGroupOrdering(VIEW_NAME, 0);
-					linearView->contextMenu().setGroupOrdering("DSCView2", 1);
-				}
-			}
+			// If we did not load the image, then we don't need to run analysis.
+			if (!controller->ApplyImage(view, *foundImage))
+				return;
+			view.AddAnalysisOption("linearsweep");
+			view.UpdateAnalysis();
 		}
+	};
+
+	auto loadImageNameAction = [](const UIActionContext& ctx) {
+		DisplayDSCPicker(ctx.context, ctx.binaryView);
+	};
+
+	auto loadSectionAddrAction = [](const UIActionContext& ctx) {
+		uint64_t addr = 0;
+		if (GetAddressInput(addr, "Address", "Address"))
+		{
+			BackgroundThread::create(ctx.context->mainWindow())
+				->thenBackground([ctx, addr](){ loadRegionAtAddr(*ctx.binaryView, addr); })
+				->start();
+		}
+	};
+
+	auto loadRegionTokenAction = [](const UIActionContext& ctx) {
+		BackgroundThread::create(ctx.context->mainWindow())
+			->thenBackground([ctx](){ loadRegionAtAddr(*ctx.binaryView, ctx.token.token.value); })
+			->start();
+	};
+
+	auto loadImageTokenAction = [](const UIActionContext& ctx) {
+		BackgroundThread::create(ctx.context->mainWindow())
+			->thenBackground([ctx](){ loadImageAtAddr(*ctx.binaryView, ctx.token.token.value); })
+			->start();
+	};
+
+	auto isValidUnloadedRegionAction = [](const UIActionContext& ctx) {
+		uint64_t addr = ctx.token.token.value;
+		// Check if the region is already loaded in the view.
+		if (!ctx.binaryView->GetSectionsAt(addr).empty())
+			return false;
+		auto controller = SharedCacheController::GetController(*ctx.binaryView);
+		if (!controller)
+			return false;
+		return controller->GetRegionContaining(addr).has_value();
+	};
+
+	auto isValidUnloadedImageAction = [](const UIActionContext& ctx) {
+		uint64_t addr = ctx.token.token.value;
+		// Check if the image is already loaded in the view.
+		if (!ctx.binaryView->GetSectionsAt(addr).empty())
+			return false;
+		auto controller = SharedCacheController::GetController(*ctx.binaryView);
+		if (!controller)
+			return false;
+		return controller->GetImageContaining(addr).has_value();
+	};
+
+	ah->bindAction("Load Image by Name", UIAction(loadImageNameAction));
+	ah->bindAction("Load Section by Address", UIAction(loadSectionAddrAction));
+
+	ah->bindAction("Load ADDRHERE", UIAction(loadRegionTokenAction, isValidUnloadedRegionAction));
+	ah->bindAction("Load IMGHERE", UIAction(loadImageTokenAction, isValidUnloadedImageAction));
+
+	ah->setActionDisplayName("Load ADDRHERE", [](const UIActionContext& ctx) {
+		auto controller = SharedCacheController::GetController(*ctx.binaryView);
+		if (!controller)
+			return QString("NO CONTROLLER");
+		uint64_t addr = ctx.token.token.value;
+		auto region = controller->GetRegionContaining(addr);
+		if (!region)
+			return QString("NO REGION");
+		return QString("Load ") + region->name.c_str();
+	});
+
+	ah->setActionDisplayName("Load IMGHERE", [](const UIActionContext& ctx) {
+		auto controller = SharedCacheController::GetController(*ctx.binaryView);
+		if (!controller)
+			return QString("NO CONTROLLER");
+		uint64_t addr = ctx.token.token.value;
+		auto image = controller->GetImageContaining(addr);
+		if (!image)
+			return QString("NO IMAGE");
+		return QString("Load ") + image->name.c_str();
+	});
+
+	// Finally add the actions to the context menu.
+	if (auto linearView = qobject_cast<LinearView*>(viewInt->widget()))
+	{
+		linearView->contextMenu().addAction("Load ADDRHERE", VIEW_NAME);
+		linearView->contextMenu().addAction("Load IMGHERE", VIEW_NAME);
+		linearView->contextMenu().addAction("Load Image by Name", "DSCView2");
+		linearView->contextMenu().addAction("Load Section by Address", "DSCView2");
+		linearView->contextMenu().setGroupOrdering(VIEW_NAME, 0);
+		linearView->contextMenu().setGroupOrdering("DSCView2", 1);
 	}
 }
+
 void UINotifications::OnAfterOpenFile(UIContext* context, FileContext* file, ViewFrame* frame)
 {
-	if (frame->getCurrentBinaryView())
-	{
-		auto listener = new SharedCacheBDNotifications(frame->getCurrentBinaryView());
-		frame->getCurrentBinaryView()->RegisterNotification(listener);
-	}
 	UIContextNotification::OnAfterOpenFile(context, file, frame);
 }