[SharedCache] Apply demangled names and types more broadly

Fixes the case of stub functions being applied with mangled names.

8 files changed, 50 insertions, 24 deletions

diff --git a/view/sharedcache/api/sharedcache.cpp b/view/sharedcache/api/sharedcache.cpp
index b92a0e64..922a7ff3 100644
--- a/view/sharedcache/api/sharedcache.cpp
+++ b/view/sharedcache/api/sharedcache.cpp
@@ -125,13 +125,19 @@ std::string SharedCacheAPI::GetRegionTypeAsString(const BNSharedCacheRegionType
 	}
 }
 
-Ref<Symbol> CacheSymbol::GetBNSymbol(BinaryView &view) const
+std::pair<std::string, Ref<Type>> CacheSymbol::DemangledName(BinaryView &view) const
 {
 	QualifiedName qname;
-	Ref<Type> outType;
+	Ref<Type> outType = nullptr;
 	std::string shortName = name;
 	if (DemangleGeneric(view.GetDefaultArchitecture(), name, outType, qname, &view, true))
 		shortName = qname.GetString();
+	return {shortName, outType};
+}
+
+Ref<Symbol> CacheSymbol::GetBNSymbol(BinaryView &view) const
+{
+	auto [shortName, _] = DemangledName(view);
 	return new Symbol(type, shortName, shortName, name, address, nullptr);
 }
 
diff --git a/view/sharedcache/api/sharedcacheapi.h b/view/sharedcache/api/sharedcacheapi.h
index db5eeefc..5a842cba 100644
--- a/view/sharedcache/api/sharedcacheapi.h
+++ b/view/sharedcache/api/sharedcacheapi.h
@@ -288,6 +288,7 @@ namespace SharedCacheAPI {
 		uint64_t address;
 		std::string name;
 
+		std::pair<std::string, BinaryNinja::Ref<BinaryNinja::Type>> DemangledName(BinaryNinja::BinaryView &view) const;
 		BinaryNinja::Ref<BinaryNinja::Symbol> GetBNSymbol(BinaryNinja::BinaryView& view) const;
 	};
 
diff --git a/view/sharedcache/core/MachOProcessor.cpp b/view/sharedcache/core/MachOProcessor.cpp
index 7f3da817..fdb56516 100644
--- a/view/sharedcache/core/MachOProcessor.cpp
+++ b/view/sharedcache/core/MachOProcessor.cpp
@@ -62,7 +62,10 @@ void SharedCacheMachOProcessor::ApplyHeader(SharedCacheMachOHeader& header)
 			// NOTE: This table is read relative to the link edit segment file base.
 			const auto symbols = header.ReadSymbolTable(*m_view, *m_vm);
 			for (const auto& sym : symbols)
-				ApplySymbol(m_view, typeLib, sym.ToBNSymbol(*m_view));
+			{
+				auto [symbol, symbolType] = sym.GetBNSymbolAndType(*m_view);
+				ApplySymbol(m_view, typeLib, symbol, symbolType);
+			}
 		}
 
 		// Apply symbols from export trie.
@@ -72,7 +75,10 @@ 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(*m_view));
+			{
+				auto [symbol, symbolType] = sym.GetBNSymbolAndType(*m_view);
+				ApplySymbol(m_view, typeLib, symbol, symbolType);
+			}
 		}
 		m_view->EndBulkModifySymbols();
 	}
diff --git a/view/sharedcache/core/SharedCache.cpp b/view/sharedcache/core/SharedCache.cpp
index 0b8653d7..2e902056 100644
--- a/view/sharedcache/core/SharedCache.cpp
+++ b/view/sharedcache/core/SharedCache.cpp
@@ -11,14 +11,21 @@ using namespace BinaryNinja;
 // The next id to use when calling Cache::AddEntry
 static CacheEntryId nextId = 1;
 
-Ref<Symbol> CacheSymbol::ToBNSymbol(BinaryView& view) const
+std::pair<std::string, Ref<Type>> CacheSymbol::DemangledName(BinaryView &view) const
 {
 	QualifiedName qname;
 	Ref<Type> outType;
 	std::string shortName = name;
 	if (DemangleGeneric(view.GetDefaultArchitecture(), name, outType, qname, &view, true))
 		shortName = qname.GetString();
-	return new Symbol(type, shortName, shortName, name, address, nullptr);
+	return { shortName, outType };
+}
+
+std::pair<Ref<Symbol>, Ref<Type>> CacheSymbol::GetBNSymbolAndType(BinaryView& view) const
+{
+	auto [shortName, demangledType] = DemangledName(view);
+	auto symbol = new Symbol(type, shortName, shortName, name, address, nullptr);
+	return {symbol, demangledType};
 }
 
 std::vector<std::string> CacheImage::GetDependencies() const
diff --git a/view/sharedcache/core/SharedCache.h b/view/sharedcache/core/SharedCache.h
index 3211b90c..0bf04776 100644
--- a/view/sharedcache/core/SharedCache.h
+++ b/view/sharedcache/core/SharedCache.h
@@ -31,8 +31,10 @@ struct CacheSymbol
 
 	CacheSymbol& operator=(CacheSymbol&& other) noexcept = default;
 
+	std::pair<std::string, BinaryNinja::Ref<BinaryNinja::Type>> DemangledName(BinaryNinja::BinaryView& view) const;
+
 	// NOTE: you should really only call this when adding the symbol to the view.
-	BinaryNinja::Ref<BinaryNinja::Symbol> ToBNSymbol(BinaryNinja::BinaryView& view) const;
+	std::pair<BinaryNinja::Ref<BinaryNinja::Symbol>, BinaryNinja::Ref<BinaryNinja::Type>> GetBNSymbolAndType(BinaryNinja::BinaryView& view) const;
 };
 
 enum class CacheRegionType
@@ -151,11 +153,8 @@ public:
 	CacheEntry& operator=(CacheEntry&&) = default;
 
 	// Construct a cache entry from the file on disk.
-	// TODO: Seperate this out a bit more.
 	static std::optional<CacheEntry> FromFile(const std::string& filePath, const std::string& fileName, CacheEntryType type);
 
-	// TODO: From Project file?
-
 	WeakFileAccessor GetAccessor() const;
 
 	// Get the headers virtual address.
diff --git a/view/sharedcache/core/Utility.cpp b/view/sharedcache/core/Utility.cpp
index 396c3337..583770e6 100644
--- a/view/sharedcache/core/Utility.cpp
+++ b/view/sharedcache/core/Utility.cpp
@@ -65,9 +65,8 @@ uint64_t readValidULEB128(const uint8_t*& current, const uint8_t* end)
 	return value;
 }
 
-void ApplySymbol(Ref<BinaryView> view, Ref<TypeLibrary> typeLib, Ref<Symbol> symbol)
+void ApplySymbol(Ref<BinaryView> view, Ref<TypeLibrary> typeLib, Ref<Symbol> symbol, Ref<Type> type)
 {
-	Ref<Function> func = nullptr;
 	auto symbolAddress = symbol->GetAddress();
 	auto symbolName = symbol->GetFullName();
 
@@ -78,24 +77,26 @@ void ApplySymbol(Ref<BinaryView> view, Ref<TypeLibrary> typeLib, Ref<Symbol> sym
 	// Define the symbol!
 	view->DefineAutoSymbol(symbol);
 
-	// Try and pull a type to apply at the symbol location.
-	Ref<Type> type = nullptr;
+	// Try and pull a type from a type library to apply at the symbol location.
+	// The type library type will take precedence over the passed in type.
+	Ref<Type> selectedType = type;
 	if (typeLib)
-		type = view->ImportTypeLibraryObject(typeLib, {symbolName});
+		selectedType = view->ImportTypeLibraryObject(typeLib, {symbolName});
 
+	Ref<Function> func = nullptr;
 	if (symbol->GetType() == FunctionSymbol)
 	{
 		Ref<Platform> targetPlatform = view->GetDefaultPlatform();
 		// Make sure to check for already added function from the function table.
 		// Unless we have retrieved a type here we don't need to make a new function.
 		func = view->GetAnalysisFunction(targetPlatform, symbolAddress);
-		if (!func || type)
-			func = view->AddFunctionForAnalysis(targetPlatform, symbolAddress, false, type);
+		if (!func || selectedType)
+			func = view->AddFunctionForAnalysis(targetPlatform, symbolAddress, false, selectedType);
 	}
 	else
 	{
 		// Other symbol types can just use this, they don't need to worry about linear sweep removing them.
-		view->DefineAutoSymbolAndVariableOrFunction(view->GetDefaultPlatform(), symbol, type);
+		view->DefineAutoSymbolAndVariableOrFunction(view->GetDefaultPlatform(), symbol, selectedType);
 	}
 
 	if (func)
diff --git a/view/sharedcache/core/Utility.h b/view/sharedcache/core/Utility.h
index e6f272e7..af242bb6 100644
--- a/view/sharedcache/core/Utility.h
+++ b/view/sharedcache/core/Utility.h
@@ -34,7 +34,7 @@ uint64_t readLEB128(const uint8_t*& current, const uint8_t* end);
 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);
+	BinaryNinja::Ref<BinaryNinja::Symbol> symbol, BinaryNinja::Ref<BinaryNinja::Type> type = nullptr);
 
 // Returns the "image name" for a given path.
 // /blah/foo/bar/libObjCThing.dylib -> libObjCThing.dylib
diff --git a/view/sharedcache/workflow/SharedCacheWorkflow.cpp b/view/sharedcache/workflow/SharedCacheWorkflow.cpp
index 5980d141..c0c90f87 100644
--- a/view/sharedcache/workflow/SharedCacheWorkflow.cpp
+++ b/view/sharedcache/workflow/SharedCacheWorkflow.cpp
@@ -87,22 +87,28 @@ void IdentifyStub(BinaryView& view, const SharedCacheController& controller, uin
 	if (!symbol.has_value())
 		return;
 
+	// Demangle if possible, the type pulled will be used, however type library will take precedence.
+	auto [demangledName, demangledType] = symbol->DemangledName(view);
+	auto rawName = STUB_PREFIX + symbol->name;
+	auto shortName = STUB_PREFIX + demangledName;
+
 	// Try and retrieve a type for the stub function using type libraries.
 	if (const auto targetFunc = view.GetAnalysisFunction(view.GetDefaultPlatform(), stubFuncAddr))
 	{
 		// NOTE: The type library name is expected to be the image name currently.
 		// Try and pull the type from the associated type library (if there is one)
-		Ref<Type> type = nullptr;
+		Ref<Type> selectedType = demangledType;
 		if (const auto image = controller.GetImageContaining(symbolAddr))
 			if (auto typeLib = TypeLibraryFromName(view, image->name))
-				type = view.ImportTypeLibraryObject(typeLib, {symbol->name});
+				selectedType = view.ImportTypeLibraryObject(typeLib, {symbol->name});
 
-		if (type)
-			targetFunc->SetAutoType(type);
+		if (selectedType)
+			targetFunc->SetAutoType(selectedType);
 	}
 
 	// Define the new symbol!
-	view.DefineAutoSymbol(new Symbol(symbol->type, STUB_PREFIX + symbol->name, stubFuncAddr));
+	auto bnSymbol = new Symbol(symbol->type, shortName, shortName, rawName, stubFuncAddr, nullptr);
+	view.DefineAutoSymbol(bnSymbol);
 }
 
 void AnalyzeStubFunction(Ref<Function> func, Ref<MediumLevelILFunction> mlil, SharedCacheController& controller, bool loadImage)