[SharedCache] Add project support for split dsc

3 files changed, 77 insertions, 20 deletions

diff --git a/view/sharedcache/core/SharedCache.cpp b/view/sharedcache/core/SharedCache.cpp
index ec0db6e7..13a12fe2 100644
--- a/view/sharedcache/core/SharedCache.cpp
+++ b/view/sharedcache/core/SharedCache.cpp
@@ -166,7 +166,7 @@ uint64_t SharedCache::FastGetBackingCacheCount(BinaryNinja::Ref<BinaryNinja::Bin
 {
 	std::shared_ptr<MMappedFileAccessor> baseFile;
 	try {
-		baseFile = MMappedFileAccessor::Open(dscView->GetFile()->GetSessionId(), dscView->GetFile()->GetOriginalFilename())->lock();
+		baseFile = MMappedFileAccessor::Open(dscView, dscView->GetFile()->GetSessionId(), dscView->GetFile()->GetOriginalFilename())->lock();
 	}
 	catch (...){
 		LogError("SharedCache preload: Failed to open file");
@@ -188,7 +188,7 @@ uint64_t SharedCache::FastGetBackingCacheCount(BinaryNinja::Ref<BinaryNinja::Bin
 	{
 		if (header.cacheType != 2)
 		{
-			if (std::filesystem::exists(baseFile->Path() + ".01"))
+			if (std::filesystem::exists(ResolveFilePath(dscView, baseFile->Path() + ".01")))
 				cacheFormat = LargeCacheFormat;
 			else
 				cacheFormat = SplitCacheFormat;
@@ -229,7 +229,7 @@ void SharedCache::PerformInitialLoad()
 {
 	m_logger->LogInfo("Performing initial load of Shared Cache");
 	auto path = m_dscView->GetFile()->GetOriginalFilename();
-	auto baseFile = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), path)->lock();
+	auto baseFile = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), path)->lock();
 
     progressMutex.lock();
 	progressMap[m_dscView->GetFile()->GetSessionId()] = LoadProgressLoadingCaches;
@@ -259,7 +259,7 @@ void SharedCache::PerformInitialLoad()
 	{
 		if (primaryCacheHeader.cacheType != 2)
 		{
-			if (std::filesystem::exists(baseFile->Path() + ".01"))
+			if (std::filesystem::exists(ResolveFilePath(m_dscView, baseFile->Path() + ".01")))
 				m_cacheFormat = LargeCacheFormat;
 			else
 				m_cacheFormat = SplitCacheFormat;
@@ -394,7 +394,7 @@ void SharedCache::PerformInitialLoad()
 				subCachePath = mainFileName + entry.fileExtension;
 			else
 				subCachePath = mainFileName + "." + entry.fileExtension;
-			auto subCacheFile = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
+			auto subCacheFile = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
 
 			dyld_cache_header subCacheHeader {};
 			uint64_t headerSize = subCacheFile->ReadUInt32(16);
@@ -485,7 +485,7 @@ void SharedCache::PerformInitialLoad()
 		for (size_t i = 1; i <= subCacheCount; i++)
 		{
 			auto subCachePath = mainFileName + "." + std::to_string(i);
-			auto subCacheFile = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
+			auto subCacheFile = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
 
 			dyld_cache_header subCacheHeader {};
 			uint64_t headerSize = subCacheFile->ReadUInt32(16);
@@ -534,7 +534,7 @@ void SharedCache::PerformInitialLoad()
 		// Load .symbols subcache
 
 		auto subCachePath = mainFileName + ".symbols";
-		auto subCacheFile = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
+		auto subCacheFile = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
 
 		dyld_cache_header subCacheHeader {};
 		uint64_t headerSize = subCacheFile->ReadUInt32(16);
@@ -624,7 +624,7 @@ void SharedCache::PerformInitialLoad()
 			else
 				subCachePath = mainFileName + "." + entry.fileExtension;
 
-			auto subCacheFile = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
+			auto subCacheFile = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
 
 			dyld_cache_header subCacheHeader {};
 			uint64_t headerSize = subCacheFile->ReadUInt32(16);
@@ -688,7 +688,7 @@ void SharedCache::PerformInitialLoad()
 		try
 		{
 			auto subCachePath = mainFileName + ".symbols";
-			auto subCacheFile = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
+			auto subCacheFile = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), subCachePath)->lock();
 			dyld_cache_header subCacheHeader {};
 			uint64_t headerSize = subCacheFile->ReadUInt32(16);
 			if (subCacheFile->ReadUInt32(16) > sizeof(dyld_cache_header))
@@ -957,7 +957,7 @@ std::shared_ptr<VM> SharedCache::GetVMMap(bool mapPages)
 		{
 			for (const auto& mapping : cache.mappings)
 			{
-				vm->MapPages(m_dscView->GetFile()->GetSessionId(), mapping.second.first, mapping.first, mapping.second.second, cache.path,
+				vm->MapPages(m_dscView, m_dscView->GetFile()->GetSessionId(), mapping.second.first, mapping.first, mapping.second.second, cache.path,
 					[this, vm=vm](std::shared_ptr<MMappedFileAccessor> mmap){
 						ParseAndApplySlideInfoForFile(mmap);
 					});
@@ -2786,11 +2786,11 @@ std::vector<std::pair<std::string, Ref<Symbol>>> SharedCache::LoadAllSymbolsAndW
 		auto header = HeaderForAddress(img.headerLocation);
 		std::shared_ptr<MMappedFileAccessor> mapping;
 		try {
-			mapping = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
+			mapping = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
 		}
 		catch (...)
 		{
-			m_logger->LogWarn("Serious Error: Failed to open export trie for %s", header->installName.c_str());
+			m_logger->LogWarn("Serious Error: Failed to open export trie %s for %s", header->exportTriePath.c_str(), header->installName.c_str());
 			continue;
 		}
 		auto exportList = SharedCache::ParseExportTrie(mapping, *header);
@@ -2862,7 +2862,7 @@ void SharedCache::FindSymbolAtAddrAndApplyToAddr(uint64_t symbolLocation, uint64
 	{
 		std::shared_ptr<MMappedFileAccessor> mapping;
 		try {
-			mapping = MMappedFileAccessor::Open(m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
+			mapping = MMappedFileAccessor::Open(m_dscView, m_dscView->GetFile()->GetSessionId(), header->exportTriePath)->lock();
 		}
 		catch (...)
 		{
diff --git a/view/sharedcache/core/VM.cpp b/view/sharedcache/core/VM.cpp
index 7fe01471..32d6c938 100644
--- a/view/sharedcache/core/VM.cpp
+++ b/view/sharedcache/core/VM.cpp
@@ -38,6 +38,7 @@
 #include <memory>
 #include <cstring>
 #include <stdio.h>
+#include <filesystem>
 #include <binaryninjaapi.h>
 
 #ifdef _MSC_VER
@@ -60,6 +61,60 @@ void VMShutdown()
 	fileAccessors.clear();
 }
 
+
+std::string ResolveFilePath(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, const std::string& path)
+{
+	auto dscProjectFile = dscView->GetFile()->GetProjectFile();
+
+	// If we're not in a project, just return the path we were given
+	if (!dscProjectFile)
+	{
+		return path;
+	}
+
+	// TODO: do we need to support looking in subfolders?
+	// Replace project file path on disk with project file name for resolution
+	std::string projectFilePathOnDisk = dscProjectFile->GetPathOnDisk();
+	std::string cleanPath = path;
+	cleanPath.replace(cleanPath.find(projectFilePathOnDisk), projectFilePathOnDisk.size(), dscProjectFile->GetName());
+
+	size_t lastSlashPos = cleanPath.find_last_of("/\\");
+	std::string fileName;
+
+	if (lastSlashPos != std::string::npos) {
+		fileName = cleanPath.substr(lastSlashPos + 1);
+	} else {
+		fileName = cleanPath;
+	}
+
+	auto project = dscProjectFile->GetProject();
+	auto dscProjectFolder = dscProjectFile->GetFolder();
+	for (const auto& file : project->GetFiles())
+	{
+		auto fileFolder = file->GetFolder();
+		bool isSibling = false;
+		if (!dscProjectFolder && !fileFolder)
+		{
+			// Both top-level
+			isSibling = true;
+		}
+		else if (dscProjectFolder && fileFolder)
+		{
+			// Have same parent folder
+			isSibling = dscProjectFolder->GetId() == fileFolder->GetId();
+		}
+
+		if (isSibling && file->GetName() == fileName)
+		{
+			return file->GetPathOnDisk();
+		}
+	}
+
+	// If we couldn't find a sibling filename, just return the path we were given
+	return path;
+}
+
+
 void MMAP::Map()
 {
 	if (mapped)
@@ -143,14 +198,14 @@ void MMAP::Unmap()
 }
 
 
-std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> MMappedFileAccessor::Open(const uint64_t sessionID, const std::string &path, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine)
+std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> MMappedFileAccessor::Open(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, const uint64_t sessionID, const std::string &path, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine)
 {
 	std::scoped_lock<std::mutex> lock(fileAccessorsMutex);
 	if (fileAccessors.count(path) == 0)
 	{
 		auto fileAcccessor = std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>>(new SelfAllocatingWeakPtr<MMappedFileAccessor>(
 			// Allocator logic for the SelfAllocatingWeakPtr
-			[path=path, sessionID=sessionID](){
+			[path=path, sessionID=sessionID, dscView](){
 				std::unique_lock<std::mutex> _lock(fileAccessorDequeMutex);
 
 				// Iterate through held references and start removing them until we can get a file pointer
@@ -165,7 +220,7 @@ std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> MMappedFileAccessor:
 
 				mmapCount++;
 				_lock.unlock();
-				auto accessor = std::shared_ptr<MMappedFileAccessor>(new MMappedFileAccessor(path), [](MMappedFileAccessor* accessor){
+				auto accessor = std::shared_ptr<MMappedFileAccessor>(new MMappedFileAccessor(ResolveFilePath(dscView, path)), [](MMappedFileAccessor* accessor){
 					// worker thread or we can deadlock on exit here.
 					BinaryNinja::WorkerEnqueue([accessor](){
 						fileAccessorSemaphore.release();
@@ -428,7 +483,7 @@ VM::~VM()
 }
 
 
-void VM::MapPages(uint64_t sessionID, size_t vm_address, size_t fileoff, size_t size, std::string filePath, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine)
+void VM::MapPages(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, uint64_t sessionID, size_t vm_address, size_t fileoff, size_t size, std::string filePath, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine)
 {
 	// The mappings provided for shared caches will always be page aligned.
 	// We can use this to our advantage and gain considerable performance via page tables.
@@ -454,7 +509,7 @@ void VM::MapPages(uint64_t sessionID, size_t vm_address, size_t fileoff, size_t
 				throw MappingCollisionException();
 			}
 		}
-		m_map.insert_or_assign(page, PageMapping(filePath, MMappedFileAccessor::Open(sessionID, filePath, postAllocationRoutine), i + fileoff));
+		m_map.insert_or_assign(page, PageMapping(filePath, MMappedFileAccessor::Open(dscView, sessionID, filePath, postAllocationRoutine), i + fileoff));
 	}
 }
 
diff --git a/view/sharedcache/core/VM.h b/view/sharedcache/core/VM.h
index d85a766d..b000266b 100644
--- a/view/sharedcache/core/VM.h
+++ b/view/sharedcache/core/VM.h
@@ -9,6 +9,8 @@
 
 void VMShutdown();
 
+std::string ResolveFilePath(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, const std::string& path);
+
 class counting_semaphore {
 public:
 	explicit counting_semaphore(int count = 0) : count_(count) {}
@@ -121,7 +123,7 @@ public:
 	MMappedFileAccessor(const std::string &path);
 	~MMappedFileAccessor();
 
-	static std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> Open(const uint64_t sessionID, const std::string &path, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine = nullptr);
+	static std::shared_ptr<SelfAllocatingWeakPtr<MMappedFileAccessor>> Open(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, const uint64_t sessionID, const std::string &path, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine = nullptr);
 
 	static void CloseAll(const uint64_t sessionID);
 
@@ -226,7 +228,7 @@ public:
 
     ~VM();
 
-    void MapPages(uint64_t sessionID, size_t vm_address, size_t fileoff, size_t size, std::string filePath, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine);
+    void MapPages(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, uint64_t sessionID, size_t vm_address, size_t fileoff, size_t size, std::string filePath, std::function<void(std::shared_ptr<MMappedFileAccessor>)> postAllocationRoutine);
 
     bool AddressIsMapped(uint64_t address);