Rework Export Trie parser to avoid recursion, improve error checking

3 files changed, 241 insertions, 80 deletions

diff --git a/view/macho/machoview.cpp b/view/macho/machoview.cpp
index 2a04b7ed..a550efea 100644
--- a/view/macho/machoview.cpp
+++ b/view/macho/machoview.cpp
@@ -2498,58 +2498,171 @@ bool MachoView::GetSectionPermissions(MachOHeader& header, uint64_t address, uin
 	return false;
 }
 
-void MachoView::ParseExportTrie(BinaryReader& reader, linkedit_data_command exportTrie)
+
+bool MachoView::AddExportTerminalSymbol(
+	const std::string& symbolName, uint64_t symbolFlags, uint64_t imageOffset)
 {
-	try {
-		uint32_t endGuard = exportTrie.datasize;
-		DataBuffer buffer = GetParentView()->ReadBuffer(m_universalImageOffset + exportTrie.dataoff, exportTrie.datasize);
+	if (symbolFlags & EXPORT_SYMBOL_FLAGS_REEXPORT)
+	{
+		m_logger->LogTrace("Export symbol is a re-export, not supported: %s", symbolName.c_str());
+		return false;
+	}
 
-		ReadExportNode(GetStart(), buffer, "", 0, endGuard);
+	uint64_t symbolAddress = GetStart() + imageOffset;
+	if (symbolName.empty() || symbolAddress == 0)
+	{
+		m_logger->LogTrace("Export symbol is malformed: %s", symbolName.c_str());
+		return false;
 	}
-	catch (ReadException&)
+
+	// Tries to get the symbol type based off the section containing it.
+	auto sectionSymbolType = [&]() -> BNSymbolType {
+		uint32_t sectionFlags = 0;
+		for (const auto& section : m_allSections)
+		{
+			if (symbolAddress >= section.addr && symbolAddress < section.addr + section.size)
+			{
+				// Take the flags from the first containing section.
+				sectionFlags = section.flags;
+				break;
+			}
+		}
+
+		// TODO: Is this enough to determine a function symbol?
+		// TODO: Might be the cause of https://github.com/Vector35/binaryninja-api/issues/6526
+		// Check the sections flags to see if we actually have a function symbol instead.
+		if (sectionFlags & S_ATTR_PURE_INSTRUCTIONS || sectionFlags & S_ATTR_SOME_INSTRUCTIONS)
+			return FunctionSymbol;
+
+		// FIXME: See above, no it is not. Fallback on old logic here to avoid breaking export symbols in __text on regular Mach-Os.
+		auto symbolType = GetAnalysisFunctionsForAddress(GetStart() + imageOffset).size() ? FunctionSymbol : DataSymbol;
+		return symbolType;
+	};
+
+	switch (symbolFlags & EXPORT_SYMBOL_FLAGS_KIND_MASK)
 	{
-		m_logger->LogError("Error while parsing Export Trie");
+	case EXPORT_SYMBOL_FLAGS_KIND_REGULAR:
+	case EXPORT_SYMBOL_FLAGS_KIND_THREAD_LOCAL:
+		m_logger->LogTrace("Export symbol is a regular or thread local symbol: %d %s", sectionSymbolType(), symbolName.c_str());
+		DefineMachoSymbol(sectionSymbolType(), symbolName, symbolAddress, GlobalBinding, false);
+		break;
+	case EXPORT_SYMBOL_FLAGS_KIND_ABSOLUTE:
+		m_logger->LogTrace("Export symbol is an absolute symbol: %s", symbolName.c_str());
+		DefineMachoSymbol(DataSymbol, symbolName, symbolAddress, GlobalBinding, false);
+		break;
+	default:
+		m_logger->LogWarn("Unhandled export symbol kind: %llx", symbolFlags & EXPORT_SYMBOL_FLAGS_KIND_MASK);
+		return false;
 	}
+
+	m_logger->LogTrace("Successfully added export symbol: %s", symbolName.c_str());
+
+	return true;
 }
 
-void MachoView::ReadExportNode(uint64_t viewStart, DataBuffer& buffer, const std::string& currentText, size_t cursor, uint32_t endGuard)
+void MachoView::ParseExportTrie(BinaryReader& reader, linkedit_data_command exportTrie)
 {
-	if (cursor > endGuard)
-		throw ReadException();
+	try {
+		uint32_t endGuard = exportTrie.datasize;
+		DataBuffer buffer = GetParentView()
+								->ReadBuffer(m_universalImageOffset + exportTrie.dataoff, exportTrie.datasize);
+
+		struct Node
+		{
+			uint64_t cursor;
+			std::string text;
+		};
+		std::vector<Node> stack;
+		stack.reserve(64);
+		stack.push_back({ /* cursor */ 0, /* text */ "" });
 
-	uint64_t terminalSize = readValidULEB128(buffer, cursor);
-	uint64_t childOffset = cursor + terminalSize;
-	if (terminalSize != 0) {
-		uint64_t imageOffset = 0;
-		uint64_t flags = readValidULEB128(buffer, cursor);
-		if (!(flags & EXPORT_SYMBOL_FLAGS_REEXPORT))
+		while (!stack.empty())
 		{
-			imageOffset = readValidULEB128(buffer, cursor);
-			auto symbolType = GetAnalysisFunctionsForAddress(viewStart + imageOffset).size() ? FunctionSymbol : DataSymbol;
-			DefineMachoSymbol(symbolType, currentText, imageOffset + viewStart, GlobalBinding, true);
+			m_logger->LogTrace("Export Trie: Processing node %s with cursor %llu", stack.back().text.c_str(), stack.back().cursor);
+			Node node = std::move(stack.back());
+			stack.pop_back();
+
+			uint64_t cursor = node.cursor;
+			const std::string currentText = std::move(node.text);
+
+			if (cursor > endGuard)
+			{
+				m_logger->LogError("Export Trie: Cursor left trie during initial bounds check");
+				throw ReadException();
+			}
+
+			size_t localCursor = cursor;
+			uint64_t terminalSize = readValidULEB128(buffer, localCursor);
+			uint64_t childOffset = localCursor + terminalSize;
+
+			// If there's terminal data, define the symbol
+			if (terminalSize != 0)
+			{
+				uint64_t flags = readValidULEB128(buffer, localCursor);
+				uint64_t imageOffset = readValidULEB128(buffer, localCursor);
+				m_logger->LogTrace("Export Trie: Found terminal node %s with flags %llx and image offset %llx", currentText.c_str(), flags, imageOffset);
+
+				AddExportTerminalSymbol(currentText, flags, imageOffset);
+			}
+
+			localCursor = childOffset;
+			if (localCursor > endGuard)
+			{
+				m_logger->LogError("Export Trie: Cursor left trie while moving to child offset");
+				throw ReadException();
+			}
+
+			uint8_t childCount = buffer[localCursor++];
+			if (localCursor > endGuard)
+			{
+				m_logger->LogError("Export Trie: Cursor left trie while reading child count");
+				throw ReadException();
+			}
+
+			std::vector<Node> children;
+			children.reserve(childCount);
+			for (uint8_t i = 0; i < childCount; ++i)
+			{
+				if (localCursor > endGuard)
+				{
+					m_logger->LogError("Export Trie: Cursor left trie while reading child count");
+					throw ReadException();
+				}
+
+				std::string childText;
+				while (localCursor <= endGuard && buffer[localCursor] != 0) {
+					childText.push_back(buffer[localCursor++]);
+				}
+				localCursor++;  // skip the `\0`
+				if (localCursor > endGuard)
+				{
+					m_logger->LogError("Export Trie: Cursor left trie while reading child text");
+					throw ReadException();
+				}
+
+				uint64_t nextOffset = readValidULEB128(buffer, localCursor);
+				if (nextOffset == 0)
+				{
+					m_logger->LogError("Export Trie: Child offset is zero");
+					throw ReadException();
+				}
+
+				children.push_back({ nextOffset, currentText + childText });
+			}
+
+			// Push in reverse so that the first child is processed next
+			for (auto it = children.rbegin(); it != children.rend(); ++it)
+			{
+				stack.push_back(*it);
+			}
 		}
 	}
-	cursor = childOffset;
-	uint8_t childCount = buffer[cursor];
-	cursor++;
-	if (cursor > endGuard)
-		throw ReadException();
-	for (uint8_t i = 0; i < childCount; ++i)
+	catch (ReadException&)
 	{
-		std::string childText;
-		while (buffer[cursor] != 0 & cursor <= endGuard)
-			childText.push_back(buffer[cursor++]);
-		cursor++;
-		if (cursor > endGuard)
-			throw ReadException();
-		auto next = readValidULEB128(buffer, cursor);
-		if (next == 0)
-			throw ReadException();
-		ReadExportNode(viewStart, buffer, currentText + childText, next, endGuard);
+		m_logger->LogError("Export trie is malformed. Could not load Exported symbol names.");
 	}
 }
 
-
 void MachoView::ParseRebaseTable(BinaryReader& reader, MachOHeader& header, uint32_t tableOffset, uint32_t tableSize)
 {
 	if (tableSize == 0 || tableOffset == 0)
diff --git a/view/macho/machoview.h b/view/macho/machoview.h
index 787ae36b..3081ff13 100644
--- a/view/macho/machoview.h
+++ b/view/macho/machoview.h
@@ -1483,6 +1483,7 @@ namespace BinaryNinja
 		void ParseFunctionStarts(Platform* platform, uint64_t textBase, function_starts_command functionStarts);
 		bool ParseRelocationEntry(const relocation_info& info, uint64_t start, BNRelocationInfo& result);
 
+		bool AddExportTerminalSymbol(const std::string& symbolName, uint64_t symbolFlags, uint64_t imageOffset);
 		void ParseExportTrie(BinaryReader& reader, linkedit_data_command exportTrie);
 		void ReadExportNode(uint64_t viewStart, DataBuffer& buffer, const std::string& currentText,
 			size_t cursor, uint32_t endGuard);
diff --git a/view/sharedcache/core/MachO.cpp b/view/sharedcache/core/MachO.cpp
index 49953911..969591a0 100644
--- a/view/sharedcache/core/MachO.cpp
+++ b/view/sharedcache/core/MachO.cpp
@@ -606,59 +606,106 @@ bool SharedCacheMachOHeader::AddExportTerminalSymbol(
 	return true;
 }
 
-// TODO: This is like 90% of the runtime.
-bool SharedCacheMachOHeader::ProcessLinkEditTrie(std::vector<CacheSymbol>& symbols, const std::string& currentText,
-	const uint8_t* begin, const uint8_t* current, const uint8_t* end) const
+std::vector<CacheSymbol> SharedCacheMachOHeader::ReadExportSymbolTrie(VirtualMemory& vm) const
 {
-	if (current >= end)
-		return false;
+	// nothing to do if there’s no export‐trie
+	if (exportTrie.datasize == 0 || exportTrie.dataoff == 0)
+		return {};
+	std::vector<CacheSymbol> symbols = {};
+	try {
+		auto [begin, end] = vm.ReadSpan(GetLinkEditFileBase() + exportTrie.dataoff, exportTrie.datasize);
+		const uint8_t *cursor = begin;
 
-	uint64_t terminalSize = readValidULEB128(current, end);
-	const uint8_t* child = current + terminalSize;
+		struct Node
+		{
+			const uint8_t* cursor;
+			std::string text;
+		};
+		std::vector<Node> stack;
+		stack.reserve(64);
+		stack.push_back({ /* cursor */ begin, /* text */ "" });
 
-	// The terminal is an export symbol.
-	if (terminalSize != 0)
-		AddExportTerminalSymbol(symbols, currentText, current, end);
+		while (!stack.empty())
+		{
+			Node node = std::move(stack.back());
+			stack.pop_back();
 
-	// TODO: Make this look better
-	current = child;
-	uint8_t childCount = *current++;
-	std::string childText = currentText;
-	for (uint8_t i = 0; i < childCount; ++i)
-	{
-		if (current >= end)
-			return false;
-		const auto it = std::find(current, end, 0);
-		childText.append(current, it);
-		current = it + 1;
-		if (current >= end)
-			return false;
-		const auto next = readValidULEB128(current, end);
-		if (next == 0)
-			return false;
-		if (!ProcessLinkEditTrie(symbols, childText, begin, begin + next, end))
-			return false;
-		childText.resize(currentText.size());
-	}
+			cursor = node.cursor;
+			const std::string currentText = std::move(node.text);
 
-	return true;
-}
+			if (cursor > end)
+			{
+				LogError("Export Trie: Cursor left trie during initial bounds check");
+				throw ReadException();
+			}
 
-std::vector<CacheSymbol> SharedCacheMachOHeader::ReadExportSymbolTrie(VirtualMemory& vm) const
-{
-	if (exportTrie.datasize == 0)
-		return {};
+			uint64_t terminalSize = readValidULEB128(cursor, end);
+			const uint8_t* childCursor = cursor + terminalSize;
 
-	uint64_t exportTrieAddress = GetLinkEditFileBase() + exportTrie.dataoff;
-	std::vector<CacheSymbol> symbols = {};
-	try
-	{
-		auto [begin, end] = vm.ReadSpan(exportTrieAddress, exportTrie.datasize);
-		ProcessLinkEditTrie(symbols, "", begin, begin, end);
+			// If there's terminal data, define the symbol
+			if (terminalSize != 0)
+			{
+				AddExportTerminalSymbol(symbols, currentText, cursor, end);
+			}
+
+			cursor = childCursor;
+			if (cursor > end)
+			{
+				LogError("Export Trie: Cursor left trie while moving to child offset");
+				throw ReadException();
+			}
+
+			uint8_t childCount = *cursor;
+			cursor++;
+			if (cursor > end)
+			{
+				LogError("Export Trie: Cursor left trie while reading child count");
+				throw ReadException();
+			}
+
+			std::vector<Node> children;
+			children.reserve(childCount);
+			for (uint8_t i = 0; i < childCount; ++i)
+			{
+				if (cursor > end)
+				{
+					LogError("Export Trie: Cursor left trie while reading children");
+					throw ReadException();
+				}
+
+				std::string childText;
+				while (cursor <= end && *cursor != 0) {
+					childText.push_back(*cursor);
+					cursor++;
+				}
+				cursor++;  // skip the `\0`
+				if (cursor > end)
+				{
+					LogError("Export Trie: Cursor left trie while reading child text");
+					throw ReadException();
+				}
+
+				uint64_t nextOffset = readValidULEB128(cursor, end);
+				if (nextOffset == 0)
+				{
+					LogError("Export Trie: Child offset is zero");
+					throw ReadException();
+				}
+
+				children.push_back({ begin + nextOffset, currentText + childText });
+			}
+
+			// Push in reverse so that the first child is processed next
+			for (auto it = children.rbegin(); it != children.rend(); ++it)
+			{
+				stack.push_back(*it);
+			}
+		}
 	}
-	catch (std::exception& e)
+	catch (ReadException&)
 	{
-		BNLogError("Failed to read Export Trie: %s", e.what());
+		LogError("Export trie is malformed. Could not load Exported symbol names.");
 	}
+
 	return symbols;
 }