[SharedCache] Refactor Shared Cache

In absence of a better name, this commit refactors the shared cache code.

1 files changed, 261 insertions, 0 deletions

diff --git a/view/sharedcache/core/SharedCacheController.cpp b/view/sharedcache/core/SharedCacheController.cpp
new file mode 100644
index 00000000..29c1192a
--- /dev/null
+++ b/view/sharedcache/core/SharedCacheController.cpp
@@ -0,0 +1,261 @@
+#include "SharedCacheController.h"
+
+#include <utility>
+
+#include "MachOProcessor.h"
+#include "ObjC.h"
+#include "SlideInfo.h"
+
+using namespace BinaryNinja;
+using namespace BinaryNinja::DSC;
+
+// Unique ID for a given Binary View.
+typedef uint64_t ViewId;
+
+std::shared_mutex GlobalControllersMutex;
+
+std::map<ViewId, DSCRef<SharedCacheController>>& GlobalControllers()
+{
+	// To make initialization order consistent we place the static in a function.
+	static std::map<ViewId, DSCRef<SharedCacheController>> g_dscViews = {};
+	return g_dscViews;
+}
+
+ViewId GetViewIdFromView(BinaryView& view)
+{
+	// Currently the view id is just the views session id.
+	// NOTE: If we want more than one shared cache controller per view we would need to make this more unique.
+	return view.GetFile()->GetSessionId();
+}
+
+void DeleteController(BinaryView& view)
+{
+	const auto id = GetViewIdFromView(view);
+	std::unique_lock<std::shared_mutex> lock(GlobalControllersMutex);
+	auto& dscViews = GlobalControllers();
+	if (auto it = dscViews.find(id); it != dscViews.end())
+	{
+		// Someone is still holding the controller, lets warn about this.
+		if (it->second->m_refs > 1)
+			LogWarn("Deleting SharedCacheController for view %llx, but there are still %d references", id,
+				it->second->m_refs.load());
+		dscViews.erase(it);
+		LogDebug("Deleted SharedCacheController for view %s", view.GetFile()->GetFilename().c_str());
+	}
+}
+
+void RegisterSharedCacheControllerDestructor()
+{
+	BNObjectDestructionCallbacks callbacks = {};
+	callbacks.destructBinaryView = [](void* ctx, BNBinaryView* obj) -> void {
+		auto view = BinaryView(obj);
+		DeleteController(view);
+	};
+	BNRegisterObjectDestructionCallbacks(&callbacks);
+}
+
+SharedCacheController::SharedCacheController(SharedCache cache, Ref<Logger> logger) : m_cache(std::move(cache))
+{
+	INIT_DSC_API_OBJECT();
+	m_logger = std::move(logger);
+	m_loadedRegions = {};
+	m_loadedImages = {};
+}
+
+DSCRef<SharedCacheController> SharedCacheController::Initialize(BinaryView& view, SharedCache cache)
+{
+	auto id = GetViewIdFromView(view);
+	std::unique_lock<std::shared_mutex> lock(GlobalControllersMutex);
+	auto logger = new Logger("SharedCacheController", view.GetFile()->GetSessionId());
+	DSCRef<SharedCacheController> dscView = new SharedCacheController(std::move(cache), logger);
+
+	// Pull the settings from the view.
+	if (Ref<Settings> settings = view.GetLoadSettings(VIEW_NAME))
+	{
+		if (settings->Contains("loader.dsc.processObjC"))
+			dscView->m_processObjC = settings->Get<bool>("loader.dsc.processObjC", &view);
+		if (settings->Contains("loader.dsc.processCFStrings"))
+			dscView->m_processCFStrings = settings->Get<bool>("loader.dsc.processCFStrings", &view);
+		if (settings->Contains("loader.dsc.regionFilter"))
+			dscView->m_regionFilter = std::regex(settings->Get<std::string>("loader.dsc.regionFilter", &view));
+	}
+
+	// Check the view auto metadata for shared cache information.
+	// This effectively restores the state of the opened database to when it was last saved.
+	auto metadata = view.GetAutoMetadata()->GetKeyValueStore();
+	if (metadata.find(METADATA_KEY) != metadata.end())
+		dscView->LoadMetadata(*metadata[METADATA_KEY]);
+
+	GlobalControllers().insert({id, dscView});
+	return dscView;
+}
+
+DSCRef<SharedCacheController> SharedCacheController::FromView(BinaryView& view)
+{
+	auto id = GetViewIdFromView(view);
+	std::shared_lock<std::shared_mutex> lock(GlobalControllersMutex);
+	auto& dscViews = GlobalControllers();
+	auto dscView = dscViews.find(id);
+	if (dscView == dscViews.end())
+		return nullptr;
+	return dscView->second;
+}
+
+bool SharedCacheController::ApplyRegionAtAddress(BinaryView& view, const uint64_t address)
+{
+	auto region = m_cache.GetRegionAt(address);
+	if (!region)
+		return false;
+	return ApplyRegion(view, *region);
+}
+
+bool SharedCacheController::ApplyRegion(BinaryView& view, const CacheRegion& region)
+{
+	// TODO: Check m_loadedRegions? This seems redundant...
+	// Loads the given region into the BinaryView and marks it as loaded.
+	// First check to make sure there isn't already a segment at the address in the view.
+	if (view.GetSegmentAt(region.start))
+		return false;
+
+	// Skip filtered regions, this defaults to just LINKEDIT regions.
+	if (std::regex_match(region.name, m_regionFilter))
+	{
+		LogDebug("Skipping filtered region at %llx", region.start);
+		return false;
+	}
+
+	auto vm = m_cache.GetVirtualMemory();
+	auto reader = VirtualMemoryReader(vm);
+	DataBuffer buffer = {};
+	try
+	{
+		buffer = reader.ReadBuffer(region.start, region.size);
+	}
+	catch (std::exception& e)
+	{
+		// This happens if we have not mapped in all the relevant entries.
+		m_logger->LogError("Failed to read region: %s", e.what());
+		return false;
+	}
+
+	// Unique memory region name so that we don't cause collisions.
+	// TODO: Better name? I dont really think so...
+	const auto memoryRegionName = fmt::format("{}_0x{:x}", region.name, region.start);
+
+	// NOTE: Adding a data memory region will store the entire contents of the region in the BNDB.
+	// TODO: We can use the AddRemoteMemoryRegion if we want to reload on view init.
+	// TODO: ^ The above is only useful if we assume that all files will be available across database loads.
+	// TODO: we might allow a user to select non-persisted memory regions as an option.
+	view.GetMemoryMap()->AddDataMemoryRegion(memoryRegionName, region.start, buffer, region.flags);
+	// TODO: We might want to make this auto if we decide to "reload" all loaded region in view init.
+	// If we are not associated with an image we can create a section here to set the semantics.
+	// This is important for stub regions, as they will deref non image data that we want to retrieve the value of.
+	if (region.type != CacheRegionType::Image)
+		view.AddUserSection(memoryRegionName, region.start, region.size, region.SectionSemanticsForRegion());
+
+	m_loadedRegions.insert(region.start);
+
+	// TODO: This needs to be done in a "database save" callback.
+	view.StoreMetadata(METADATA_KEY, GetMetadata());
+
+	return true;
+}
+
+bool SharedCacheController::IsRegionLoaded(const CacheRegion& region) const
+{
+	return std::any_of(m_loadedRegions.begin(), m_loadedRegions.end(), [&](const auto& loadedRegion) {
+		return loadedRegion == region.start;
+	});
+}
+
+bool SharedCacheController::ApplyImage(BinaryView& view, const CacheImage& image)
+{
+	// TODO: Check m_loadedImages? This seems redundant...
+	// Load all regions of an image and mark the image as loaded.
+	bool loadedRegion = false;
+	for (const auto& regionStart : image.regionStarts)
+		if (ApplyRegionAtAddress(view, regionStart))
+			loadedRegion = true;
+
+	// If there was no loaded regions than we just want to forgo loading the image.
+	if (!loadedRegion)
+		return false;
+
+	if (image.header)
+	{
+		// Header information is applied to the view here, such as sections.
+		auto machoProcessor = SharedCacheMachOProcessor(&view, m_cache.GetVirtualMemory());
+		machoProcessor.ApplyHeader(*image.header);
+
+		// TODO: Make this optional.
+		// Load objective-c information.
+		auto objcProcessor = DSCObjC::SharedCacheObjCProcessor(&view, false);
+		// TODO: Passing in an image name here is weird considering this is shared with the MACHO view.
+		// TODO: We should abstract out the "image" into an objc image type that represents what is required, which ig is the name?
+		try
+		{
+			if (m_processObjC)
+				objcProcessor.ProcessObjCData(image.GetName());
+			if (m_processObjC)
+			objcProcessor.ProcessCFStrings(image.GetName());
+		}
+		catch (std::exception& e)
+		{
+			// Let the user know there was an error in processing the objc stuff but let the image load
+			// regardless, as its non-critical.
+			m_logger->LogError("Failed to process ObjC information: %s", e.what());
+		}
+	}
+
+	m_loadedImages.insert(image.headerAddress);
+
+	// TODO: This needs to be done in a "database save" callback.
+	view.StoreMetadata(METADATA_KEY, GetMetadata());
+
+	// TODO: Partial failure state (i.e. 2 regions loaded, one failed)
+	return true;
+}
+
+bool SharedCacheController::IsImageLoaded(const CacheImage& image) const
+{
+	return std::any_of(m_loadedImages.begin(), m_loadedImages.end(), [&](const auto& loadedImage) {
+		return loadedImage == image.headerAddress;
+	});
+}
+
+Ref<Metadata> SharedCacheController::GetMetadata() const
+{
+	std::map<std::string, Ref<Metadata>> controllerMeta;
+
+	std::vector<uint64_t> loadedImages;
+	std::vector<uint64_t> loadedRegions;
+	loadedImages.reserve(m_loadedImages.size());
+	loadedRegions.reserve(m_loadedRegions.size());
+	for (const auto& loadedImage : m_loadedImages)
+		loadedImages.push_back(loadedImage);
+	for (const auto& loadedRegion : m_loadedRegions)
+		loadedRegions.push_back(loadedRegion);
+
+	controllerMeta["loadedImages"] = new Metadata(loadedImages);
+	controllerMeta["loadedRegions"] = new Metadata(loadedRegions);
+
+	return new Metadata(controllerMeta);
+}
+
+void SharedCacheController::LoadMetadata(const Metadata& metadata)
+{
+	auto controllerMeta = metadata.GetKeyValueStore();
+	if (controllerMeta.find("loadedImages") != controllerMeta.end())
+	{
+		const auto loadedImages = controllerMeta["loadedImages"]->GetUnsignedIntegerList();
+		for (const auto& image : loadedImages)
+			m_loadedImages.insert(image);
+	}
+
+	if (controllerMeta.find("loadedRegions") != controllerMeta.end())
+	{
+		const auto loadedRegions = controllerMeta["loadedRegions"]->GetUnsignedIntegerList();
+		for (const auto& region : loadedRegions)
+			m_loadedImages.insert(region);
+	}
+}