[SharedCache] Fix some bugs

9 files changed, 47 insertions, 39 deletions

diff --git a/view/sharedcache/core/MachOProcessor.cpp b/view/sharedcache/core/MachOProcessor.cpp
index be2e9007..893aa8fe 100644
--- a/view/sharedcache/core/MachOProcessor.cpp
+++ b/view/sharedcache/core/MachOProcessor.cpp
@@ -35,11 +35,7 @@ void SharedCacheMachOProcessor::ApplyHeader(SharedCacheMachOHeader& header)
 			auto targetPlatform = m_view->GetDefaultPlatform();
 			auto functions = header.ReadFunctionTable(*m_vm);
 			for (const auto& func : functions)
-			{
-				// TODO: Check to make sure the func exists in a loaded region?
-				// TODO: ^ this check existed prior so we should add it back.
-				m_view->AddFunctionForAnalysis(targetPlatform, func);
-			}
+				m_view->AddFunctionForAnalysis(targetPlatform, func, true);
 		}
 
 		auto typeLib = m_view->GetTypeLibrary(header.installName);
@@ -52,10 +48,9 @@ void SharedCacheMachOProcessor::ApplyHeader(SharedCacheMachOHeader& header)
 			// Mach-O View symtab processing with
 			// a ton of stuff cut out so it can work
 			// NOTE: This table is read relative to the link edit segment file base.
-			// TODO: Remove this and use the m_symbols in the cache?
 			const auto symbols = header.ReadSymbolTable(*m_view, *m_vm);
 			for (const auto& sym : symbols)
-				ApplySymbol(m_view, typeLib, sym.ToBNSymbol());
+				ApplySymbol(m_view, typeLib, sym.ToBNSymbol(*m_view));
 		}
 
 		// Apply symbols from export trie.
@@ -65,7 +60,7 @@ void SharedCacheMachOProcessor::ApplyHeader(SharedCacheMachOHeader& header)
 			// TODO: Remove this and use the m_symbols in the cache?
 			const auto exportSymbols = header.ReadExportSymbolTrie(*m_vm);
 			for (const auto& sym : exportSymbols)
-				ApplySymbol(m_view, typeLib, sym.ToBNSymbol());
+				ApplySymbol(m_view, typeLib, sym.ToBNSymbol(*m_view));
 		}
 		m_view->EndBulkModifySymbols();
 	}
diff --git a/view/sharedcache/core/SharedCache.cpp b/view/sharedcache/core/SharedCache.cpp
index 56d2040d..d0c3d25b 100644
--- a/view/sharedcache/core/SharedCache.cpp
+++ b/view/sharedcache/core/SharedCache.cpp
@@ -11,11 +11,12 @@ using namespace BinaryNinja;
 // The next id to use when calling Cache::AddEntry
 static CacheEntryId nextId = 1;
 
-Ref<Symbol> CacheSymbol::ToBNSymbol() const
+Ref<Symbol> CacheSymbol::ToBNSymbol(BinaryView& view) const
 {
 	QualifiedName qname;
+	Ref<Type> outType;
 	std::string shortName = name;
-	if (DemangleLLVM(name, qname, true))
+	if (DemangleGeneric(view.GetDefaultArchitecture(), name, outType, qname, &view, true))
 		shortName = qname.GetString();
 	return new Symbol(type, shortName, shortName, name, address, nullptr);
 }
@@ -458,7 +459,7 @@ std::optional<CacheImage> SharedCache::GetImageContaining(const uint64_t address
 {
 	// TODO: What if we are using this on a shared region? Return a list of images?
 	auto region = GetRegionContaining(address);
-	if (region.has_value())
+	if (region.has_value() && region->imageStart.has_value())
 		return GetImageAt(*region->imageStart);
 	return std::nullopt;
 }
diff --git a/view/sharedcache/core/SharedCache.h b/view/sharedcache/core/SharedCache.h
index 54ec9dfd..674e00c3 100644
--- a/view/sharedcache/core/SharedCache.h
+++ b/view/sharedcache/core/SharedCache.h
@@ -32,7 +32,7 @@ struct CacheSymbol
 	CacheSymbol& operator=(CacheSymbol&& other) noexcept = default;
 
 	// NOTE: you should really only call this when adding the symbol to the view.
-	BinaryNinja::Ref<BinaryNinja::Symbol> ToBNSymbol() const;
+	BinaryNinja::Ref<BinaryNinja::Symbol> ToBNSymbol(BinaryNinja::BinaryView& view) const;
 };
 
 enum class CacheRegionType
diff --git a/view/sharedcache/core/SharedCacheController.cpp b/view/sharedcache/core/SharedCacheController.cpp
index 29c1192a..72d5d950 100644
--- a/view/sharedcache/core/SharedCacheController.cpp
+++ b/view/sharedcache/core/SharedCacheController.cpp
@@ -60,6 +60,9 @@ SharedCacheController::SharedCacheController(SharedCache cache, Ref<Logger> logg
 	m_logger = std::move(logger);
 	m_loadedRegions = {};
 	m_loadedImages = {};
+	m_processObjC = true;
+	m_processCFStrings = true;
+	m_regionFilter = std::regex(".*LINKEDIT.*");
 }
 
 DSCRef<SharedCacheController> SharedCacheController::Initialize(BinaryView& view, SharedCache cache)
@@ -111,16 +114,16 @@ bool SharedCacheController::ApplyRegionAtAddress(BinaryView& view, const uint64_
 
 bool SharedCacheController::ApplyRegion(BinaryView& view, const CacheRegion& region)
 {
-	// TODO: Check m_loadedRegions? This seems redundant...
+	std::unique_lock<std::shared_mutex> lock(m_loadMutex);
 	// 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))
+	// First check to make sure we haven't already loaded the region.
+	if (m_loadedRegions.find(region.start) != m_loadedRegions.end())
 		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);
+		m_logger->LogDebug("Skipping filtered region at %llx", region.start);
 		return false;
 	}
 
@@ -161,8 +164,9 @@ bool SharedCacheController::ApplyRegion(BinaryView& view, const CacheRegion& reg
 	return true;
 }
 
-bool SharedCacheController::IsRegionLoaded(const CacheRegion& region) const
+bool SharedCacheController::IsRegionLoaded(const CacheRegion& region)
 {
+	std::shared_lock<std::shared_mutex> lock(m_loadMutex);
 	return std::any_of(m_loadedRegions.begin(), m_loadedRegions.end(), [&](const auto& loadedRegion) {
 		return loadedRegion == region.start;
 	});
@@ -170,15 +174,19 @@ bool SharedCacheController::IsRegionLoaded(const CacheRegion& region) const
 
 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.
+	// NOTE: The regions lock m_loadMutex themselves, so we do not hold it up here.
 	bool loadedRegion = false;
 	for (const auto& regionStart : image.regionStarts)
 		if (ApplyRegionAtAddress(view, regionStart))
 			loadedRegion = true;
 
+	// 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.
-	if (!loadedRegion)
+	// We also skip if we already loaded the image itself. We do this after loading regions
+	// as we regions have their own check.
+	if (!loadedRegion || m_loadedImages.find(image.headerAddress) != m_loadedImages.end())
 		return false;
 
 	if (image.header)
@@ -187,17 +195,16 @@ bool SharedCacheController::ApplyImage(BinaryView& view, const CacheImage& image
 		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?
+		// Load objective-c information.
+		auto objcProcessor = DSCObjC::SharedCacheObjCProcessor(&view, false);
 		try
 		{
 			if (m_processObjC)
 				objcProcessor.ProcessObjCData(image.GetName());
-			if (m_processObjC)
-			objcProcessor.ProcessCFStrings(image.GetName());
+			if (m_processCFStrings)
+				objcProcessor.ProcessCFStrings(image.GetName());
 		}
 		catch (std::exception& e)
 		{
@@ -216,8 +223,9 @@ bool SharedCacheController::ApplyImage(BinaryView& view, const CacheImage& image
 	return true;
 }
 
-bool SharedCacheController::IsImageLoaded(const CacheImage& image) const
+bool SharedCacheController::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.headerAddress;
 	});
diff --git a/view/sharedcache/core/SharedCacheController.h b/view/sharedcache/core/SharedCacheController.h
index 354762eb..3b92f198 100644
--- a/view/sharedcache/core/SharedCacheController.h
+++ b/view/sharedcache/core/SharedCacheController.h
@@ -18,8 +18,11 @@ namespace BinaryNinja::DSC {
 	{
 		IMPLEMENT_DSC_API_OBJECT(BNSharedCacheController);
 		Ref<Logger> m_logger;
-
 		SharedCache m_cache;
+
+		// Locks on load attempts (region or image).
+		std::shared_mutex m_loadMutex;
+
 		// Store the open images.
 		// Things other than the cache here will be serialized.
 		std::unordered_set<uint64_t> m_loadedRegions;
@@ -49,13 +52,13 @@ namespace BinaryNinja::DSC {
 
 		bool ApplyRegion(BinaryView& view, const CacheRegion& region);
 
-		bool IsRegionLoaded(const CacheRegion& region) const;
+		bool IsRegionLoaded(const CacheRegion& region);
 
 		// Loads the relevant image info into the view. This does not update analysis so if you
 		// call this make sure at some point you update analysis and likely with linear sweep.
 		bool ApplyImage(BinaryView& view, const CacheImage& image);
 
-		bool IsImageLoaded(const CacheImage& image) const;
+		bool IsImageLoaded(const CacheImage& image);
 
 		// Get the metadata for saving the state of the shared cache.
 		Ref<Metadata> GetMetadata() const;
diff --git a/view/sharedcache/core/SharedCacheView.cpp b/view/sharedcache/core/SharedCacheView.cpp
index bf80c6eb..f27f7c14 100644
--- a/view/sharedcache/core/SharedCacheView.cpp
+++ b/view/sharedcache/core/SharedCacheView.cpp
@@ -82,7 +82,7 @@ Ref<Settings> SharedCacheViewType::GetLoadSettingsForData(BinaryView* data)
 	R"({
 		"title" : "Region Regex Filter",
 		"type" : "string",
-		"default" : "LINKEDIT",
+		"default" : ".*LINKEDIT.*",
 		"description" : "Regex filter for region names to skip loading, by default this filters out the link edit region which is not necessary to be applied to the view."
 		})");
 
@@ -784,19 +784,22 @@ bool SharedCacheView::Init()
 		// After this we should have all the mappings available as well.
 		CacheProcessor cacheProcessor = CacheProcessor(this);
 		auto startTime = std::chrono::high_resolution_clock::now();
-		if (!cacheProcessor.ProcessCache(sharedCache))
+		bool result = cacheProcessor.ProcessCache(sharedCache);
+		auto endTime = std::chrono::high_resolution_clock::now();
+		std::chrono::duration<double> elapsed = endTime - startTime;
+		logger->LogInfo("Processing %zu entries took %.3f seconds", sharedCache.GetEntries().size(), elapsed.count());
+
+		if (!result)
 		{
-			// Oh no, we failed to process the cache, this likely means the primary on-disk file was not able to be
-			// found.
 			// TODO: Prompt the user to select the primary cache file? We can still recover from here.
+			// Oh no, we failed to process the cache, this likely means the primary on-disk file was not able to be found.
 			logger->LogError("Failed to process cache, likely missing cache files.");
 			// NOTE: An interaction handler headlessly could select yes to this, however for the purposes of this we will just let them continue by defualt.
 			if (IsUIEnabled() && ShowMessageBox("Continue opening", "This shared cache file was unable to be processed, would you like to open without shared cache information?", YesNoButtonSet, QuestionIcon) != YesButton)
 				return false;
+			// Skip all the other shared cache stuff.
+			return true;
 		}
-		auto endTime = std::chrono::high_resolution_clock::now();
-		std::chrono::duration<double> elapsed = endTime - startTime;
-		logger->LogInfo("Processing %zu entries took %.3f seconds", sharedCache.GetEntries().size(), elapsed.count());
 	}
 
 	{
diff --git a/view/sharedcache/core/Utility.cpp b/view/sharedcache/core/Utility.cpp
index 91729578..c4891930 100644
--- a/view/sharedcache/core/Utility.cpp
+++ b/view/sharedcache/core/Utility.cpp
@@ -74,7 +74,7 @@ void ApplySymbol(Ref<BinaryView> view, Ref<TypeLibrary> typeLib, Ref<Symbol> sym
 	if (symbol->GetType() == FunctionSymbol)
 	{
 		Ref<Platform> targetPlatform = view->GetDefaultPlatform();
-		func = view->AddFunctionForAnalysis(targetPlatform, symbolAddress);
+		func = view->AddFunctionForAnalysis(targetPlatform, symbolAddress, true);
 	}
 
 	if (typeLib)
diff --git a/view/sharedcache/core/Utility.h b/view/sharedcache/core/Utility.h
index 86715557..e6f272e7 100644
--- a/view/sharedcache/core/Utility.h
+++ b/view/sharedcache/core/Utility.h
@@ -36,9 +36,6 @@ uint64_t readValidULEB128(const uint8_t*& current, const uint8_t* end);
 void ApplySymbol(BinaryNinja::Ref<BinaryNinja::BinaryView> view, BinaryNinja::Ref<BinaryNinja::TypeLibrary> typeLib,
 	BinaryNinja::Ref<BinaryNinja::Symbol> symbol);
 
-// Returns the on disk file for a given path, this is required to support project files.
-std::string ResolveFilePath(BinaryNinja::Ref<BinaryNinja::BinaryView> dscView, const std::string& path);
-
 // Returns the "image name" for a given path.
 // /blah/foo/bar/libObjCThing.dylib -> libObjCThing.dylib
 std::string BaseFileName(const std::string& path);
diff --git a/view/sharedcache/core/VirtualMemory.h b/view/sharedcache/core/VirtualMemory.h
index fb00ebaf..43a9acb4 100644
--- a/view/sharedcache/core/VirtualMemory.h
+++ b/view/sharedcache/core/VirtualMemory.h
@@ -55,6 +55,7 @@ public:
 
 	bool IsAddressMapped(uint64_t address);
 
+	// TODO: Bulk pointer writes here would alleviate a lot of the time spent in the slide info processor.
 	// Write a pointer at a given address. This pointer will be persisted
 	// for a given `VirtualMemoryRegion` region, unlike using the MappedFileAccessor directly.
 	// The persistence is provided through the WeakFileAccessor itself and thus is unique to the construction.