1 files changed, 198 insertions, 0 deletions
diff --git a/view/kernelcache/core/KernelCacheController.cpp b/view/kernelcache/core/KernelCacheController.cpp
new file mode 100644
index 00000000..ef2fbd04
--- /dev/null
+++ b/view/kernelcache/core/KernelCacheController.cpp
@@ -0,0 +1,198 @@
+#include "KernelCacheController.h"
+#include "MachOProcessor.h"
+
+using namespace BinaryNinja;
+using namespace BinaryNinja::KC;
+
+typedef uint64_t ViewId;
+
+std::shared_mutex GlobalControllersMutex;
+
+std::map<ViewId, KCRef<KernelCacheController>>& GlobalControllers()
+{
+	static std::map<ViewId, KCRef<KernelCacheController>> g_controllers = {};
+	return g_controllers;
+}
+
+
+ViewId GetViewIdFromFileMetadata(const FileMetadata& file)
+{
+	// 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 file.GetSessionId();
+}
+
+void DeleteController(const FileMetadata& file)
+{
+	const auto id = GetViewIdFromFileMetadata(file);
+	std::unique_lock<std::shared_mutex> lock(GlobalControllersMutex);
+	auto& controllers = GlobalControllers();
+	if (auto it = controllers.find(id); it != controllers.end())
+	{
+		auto controller = it->second;
+		// Someone is still holding the controller, lets warn about this.
+		// 2 is expected here because we have one held in `controllers` and one held by `controller`.
+		if (controller->m_refs > 2)
+			LogWarn("Deleting KernelCacheController for view %llx, but there are still %d references", id,
+				controller->m_refs.load());
+
+		controllers.erase(it);
+		LogDebug("Deleted KernelCacheController for view %s", file.GetFilename().c_str());
+	}
+}
+
+void RegisterKernelCacheControllerDestructor()
+{
+	BNObjectDestructionCallbacks callbacks = {};
+	callbacks.destructFileMetadata = [](void* ctx, BNFileMetadata* obj) -> void {
+		const auto file = FileMetadata(obj);
+		DeleteController(file);
+	};
+	BNRegisterObjectDestructionCallbacks(&callbacks);
+}
+
+KernelCacheController::KernelCacheController(KernelCache&& cache, Ref<Logger> logger) : m_cache(std::move(cache))
+{
+	INIT_KC_API_OBJECT();
+	m_logger = std::move(logger);
+	m_loadedImages = {};
+	m_regionFilter = std::regex(".*LINKEDIT.*");
+}
+
+KCRef<KernelCacheController> KernelCacheController::Initialize(BinaryView& view, KernelCache&& cache)
+{
+	auto id = GetViewIdFromFileMetadata(*view.GetFile());
+	std::unique_lock<std::shared_mutex> lock(GlobalControllersMutex);
+	auto logger = new Logger("KernelCache.Controller", view.GetFile()->GetSessionId());
+	KCRef<KernelCacheController> controller = new KernelCacheController(std::move(cache), logger);
+
+	// Pull the settings from the view.
+	if (Ref<Settings> settings = view.GetLoadSettings(KC_VIEW_NAME))
+	{
+		if (settings->Contains("loader.kc.regionFilter"))
+			controller->m_regionFilter = std::regex(settings->Get<std::string>("loader.kc.regionFilter", &view));
+	}
+
+	// TODO: Support old shared cache metadata
+	// TODO: Not strictly necessary as the user has already loaded the information into the database, this would just
+	// TODO: prevent incidental extra work from being done when loading a region or image.
+	// const uint64_t oldStateCount = view.GetUIntMetadata(OLD_METADATA_KEY_COUNT);
+
+	// Check the view auto metadata for shared cache information.
+	// This effectively restores the state of the opened database to when it was last saved.
+	// NOTE: We store on the parent view because hilariously, the metadata is not present until after view init.
+	if (auto loadedImageMetadata = view.GetParentView()->QueryMetadata("KernelCacheLoadedImages"))
+	{
+		auto loadedImageList = loadedImageMetadata->GetArray();
+		for (const auto & imageAddrMeta : loadedImageList)
+		{
+			auto imageAddr = imageAddrMeta->GetUnsignedInteger();
+			controller->m_loadedImages.insert(imageAddr);
+		}
+	}
+
+	GlobalControllers().insert({id, controller});
+	return controller;
+}
+
+KCRef<KernelCacheController> KernelCacheController::FromView(const BinaryView& view)
+{
+	auto id = GetViewIdFromFileMetadata(*view.GetFile());
+	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 KernelCacheController::ApplyImage(BinaryView& view, const CacheImage& image)
+{
+	// Load all regions of an image and mark the image as loaded.
+	// NOTE: The regions lock m_loadMutex themselves, so we do not hold it up here.
+	bool loadedRegion = false;
+
+	BNRelocationInfo reloc;
+	memset(&reloc, 0, sizeof(BNRelocationInfo));
+	reloc.type = StandardRelocationType;
+	reloc.size = 8;
+	reloc.nativeType = BINARYNINJA_MANUAL_RELOCATION;
+
+	// We check for a valid offset as we apply auto segments and re-call this function on view init
+	if (!view.IsValidOffset(image.headerVirtualAddress))
+	{
+		loadedRegion = true;
+		for (const auto& segment : image.header->segments)
+		{
+			view.AddAutoSegment(segment.vmaddr, segment.vmsize, segment.fileoff, segment.filesize, segment.flags);
+
+			auto relocations = m_cache.GetRelocations();
+
+			auto begin = std::lower_bound(relocations.begin(), relocations.end(), segment.vmaddr,
+				[](const std::pair<uint64_t, uint64_t>& reloc, uint64_t addr) {
+					return reloc.first < addr;
+				});
+
+			auto arch = view.GetDefaultArchitecture();
+			// Process relocations until the VM address is beyond our region
+			for (auto it = begin; it != relocations.end() && it->first < segment.vmaddr + segment.vmsize; ++it) {
+				reloc.address = it->first;
+				view.DefineRelocation(arch, reloc, it->second, reloc.address);
+			}
+		}
+	}
+
+	view.FinalizeNewSegments();
+
+	// The ApplyRegionAtAddress no longer holds the lock, we can take it now.
+	std::unique_lock<std::shared_mutex> lock(m_loadMutex);
+
+
+	// If there was no loaded regions than we just want to forgo loading the image.
+	// We also skip if we already loaded the image itself. We do this after loading regions
+	// as we regions have their own check.
+
+	// On view init, we exit here as we just need to re-add the segments and relocations.
+	if (!loadedRegion || m_loadedImages.find(image.headerVirtualAddress) != m_loadedImages.end())
+		return false;
+
+	if (image.headerVirtualAddress)
+	{
+		// Header information is applied to the view here, such as sections.
+		auto machoProcessor = KernelCacheMachOProcessor(&view);
+
+		// Adding a user section will mark all functions for updates unless we disable this.
+		// Because images are known separate compilation units, we have a real reason to make sure we don't mark all previously
+		// analyzed functions as updated.
+		auto prevDisabledState = view.GetFunctionAnalysisUpdateDisabled();
+		view.SetFunctionAnalysisUpdateDisabled(true);
+		machoProcessor.ApplyHeader(GetCache(), *image.header);
+		view.SetFunctionAnalysisUpdateDisabled(prevDisabledState);
+	}
+
+	m_loadedImages.insert(image.headerVirtualAddress);
+
+	m_logger->LogInfoF("Loaded image: '{}'", image.path);
+
+	// TODO: This needs to be done in a "database save" callback.
+	// NOTE: We store on the parent view because hilariously, the view metadata is not available in view init.
+	std::vector<uint64_t> loadedImages;
+	loadedImages.reserve(m_loadedImages.size());
+	// Store the loaded images in the metadata.
+	for (const uint64_t& addr : m_loadedImages)
+	{
+		loadedImages.push_back(addr);
+	}
+	view.GetParentView()->StoreMetadata("KernelCacheLoadedImages", new Metadata(loadedImages));
+
+	// TODO: Partial failure state (i.e. 2 regions loaded, one failed)
+	return true;
+}
+
+bool KernelCacheController::IsImageLoaded(const CacheImage& image)
+{
+	std::shared_lock<std::shared_mutex> lock(m_loadMutex);
+	return std::any_of(m_loadedImages.begin(), m_loadedImages.end(), [&](const auto& loadedImage) {
+		return loadedImage == image.headerVirtualAddress;
+	});
+}