Initial support for guided disassembly.

5 files changed, 137 insertions, 2 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 8ca9ad33..41a93d08 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -11435,6 +11435,12 @@ namespace BinaryNinja {
 		void SetUserIndirectBranches(
 		    Architecture* sourceArch, uint64_t source, const std::vector<ArchAndAddr>& branches);
 
+		// Guided Disassembly Support
+		void SetGuidedSourceBlocks(const std::vector<ArchAndAddr>& addresses);
+		void AddGuidedSourceBlocks(const std::vector<ArchAndAddr>& addresses);
+		void RemoveGuidedSourceBlocks(const std::vector<ArchAndAddr>& addresses);
+		std::vector<ArchAndAddr> GetGuidedSourceBlocks();
+
 		std::vector<IndirectBranchInfo> GetIndirectBranches();
 		std::vector<IndirectBranchInfo> GetIndirectBranchesAt(Architecture* arch, uint64_t addr);
 
diff --git a/binaryninjacore.h b/binaryninjacore.h
index cdfbfcee..46d1cdef 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,7 +37,7 @@
 // Current ABI version for linking to the core. This is incremented any time
 // there are changes to the API that affect linking, including new functions,
 // new types, or modifications to existing functions or types.
-#define BN_CURRENT_CORE_ABI_VERSION 113
+#define BN_CURRENT_CORE_ABI_VERSION 114
 
 // Minimum ABI version that is supported for loading of plugins. Plugins that
 // are linked to an ABI version less than this will not be able to load and
@@ -5112,6 +5112,12 @@ extern "C"
 	BINARYNINJACOREAPI void BNSetUserIndirectBranches(BNFunction* func, BNArchitecture* sourceArch, uint64_t source,
 	    BNArchitectureAndAddress* branches, size_t count);
 
+	BINARYNINJACOREAPI void BNSetGuidedSourceBlocks(BNFunction* func, BNArchitectureAndAddress* addresses, size_t count);
+	BINARYNINJACOREAPI void BNAddGuidedSourceBlocks(BNFunction* func, BNArchitectureAndAddress* addresses, size_t count);
+	BINARYNINJACOREAPI void BNRemoveGuidedSourceBlocks(BNFunction* func, BNArchitectureAndAddress* addresses, size_t count);
+	BINARYNINJACOREAPI BNArchitectureAndAddress* BNGetGuidedSourceBlocks(BNFunction* func, size_t* count);
+	BINARYNINJACOREAPI void BNFreeArchitectureAndAddressList(BNArchitectureAndAddress* addresses);
+
 	BINARYNINJACOREAPI BNIndirectBranchInfo* BNGetIndirectBranches(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI BNIndirectBranchInfo* BNGetIndirectBranchesAt(
 	    BNFunction* func, BNArchitecture* arch, uint64_t addr, size_t* count);
diff --git a/defaultabb.cpp b/defaultabb.cpp
index df5d0c5c..477e5722 100644
--- a/defaultabb.cpp
+++ b/defaultabb.cpp
@@ -76,6 +76,13 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 	auto& directNoReturnCalls = context.GetDirectNoReturnCalls();
 	auto& haltedDisassemblyAddresses = context.GetHaltedDisassemblyAddresses();
 
+	bool hasInvalidInstructions = false;
+	set<ArchAndAddr> guidedSourceBlockTargets;
+	auto guidedSourceBlocks = function->GetGuidedSourceBlocks();
+	set<ArchAndAddr> guidedSourceBlocksSet;
+	for (const auto& block : guidedSourceBlocks)
+		guidedSourceBlocksSet.insert(block);
+
 	BNStringReference strRef;
 	auto targetExceedsByteLimit = [](const BNStringReference& strRef) {
 			size_t byteLimit = 8;
@@ -164,6 +171,8 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 		ArchAndAddr instructionGroupStart = location;
 		blocksToProcess.pop();
 
+		bool isGuidedSourceBlock = guidedSourceBlocksSet.count(location) ? true : false;
+
 		// Create a new basic block
 		Ref<BasicBlock> block = context.CreateBasicBlock(location.arch, location.address);
 
@@ -385,6 +394,8 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 							}
 							else
 							{
+								if (isGuidedSourceBlock)
+									guidedSourceBlockTargets.insert(target);
 
 								block->AddPendingOutgoingEdge(info.branchType[i], target.address, target.arch);
 								// Add the block to the list of blocks to process if it is not already processed
@@ -507,6 +518,9 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 								if (translateTailCalls && data->GetAnalysisFunction(targetPlatform, branch.address))
 									continue;
 
+								if (isGuidedSourceBlock)
+									guidedSourceBlockTargets.insert(branch);
+
 								block->AddPendingOutgoingEdge(IndirectBranch, branch.address, branch.arch);
 								if (seenBlocks.count(branch) == 0)
 								{
@@ -644,13 +658,21 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 			break;
 
 		if (haltOnInvalidInstructions && block->HasInvalidInstructions())
+			hasInvalidInstructions = true;
+
+		if (hasInvalidInstructions)
 		{
+			queue<ArchAndAddr> guidedBlocksToProcess;
 			while (!blocksToProcess.empty())
 			{
 				auto i = blocksToProcess.front();
 				blocksToProcess.pop();
-				haltedDisassemblyAddresses.emplace(i);
+				if (guidedSourceBlockTargets.count(i))
+					guidedBlocksToProcess.emplace(i);
+				else
+					haltedDisassemblyAddresses.emplace(i);
 			}
+			blocksToProcess = guidedBlocksToProcess;
 		}
 	}
 
diff --git a/function.cpp b/function.cpp
index 103fec95..2f04af9a 100644
--- a/function.cpp
+++ b/function.cpp
@@ -1741,6 +1741,62 @@ void Function::SetUserIndirectBranches(
 }
 
 
+void Function::SetGuidedSourceBlocks(const std::vector<ArchAndAddr>& addresses)
+{
+	BNArchitectureAndAddress* addressList = new BNArchitectureAndAddress[addresses.size()];
+	for (size_t i = 0; i < addresses.size(); i++)
+	{
+		addressList[i].arch = addresses[i].arch->GetObject();
+		addressList[i].address = addresses[i].address;
+	}
+	BNSetGuidedSourceBlocks(m_object, addressList, addresses.size());
+	delete[] addressList;
+}
+
+
+void Function::AddGuidedSourceBlocks(const std::vector<ArchAndAddr>& addresses)
+{
+	BNArchitectureAndAddress* addressList = new BNArchitectureAndAddress[addresses.size()];
+	for (size_t i = 0; i < addresses.size(); i++)
+	{
+		addressList[i].arch = addresses[i].arch->GetObject();
+		addressList[i].address = addresses[i].address;
+	}
+	BNAddGuidedSourceBlocks(m_object, addressList, addresses.size());
+	delete[] addressList;
+}
+
+
+void Function::RemoveGuidedSourceBlocks(const std::vector<ArchAndAddr>& addresses)
+{
+	BNArchitectureAndAddress* addressList = new BNArchitectureAndAddress[addresses.size()];
+	for (size_t i = 0; i < addresses.size(); i++)
+	{
+		addressList[i].arch = addresses[i].arch->GetObject();
+		addressList[i].address = addresses[i].address;
+	}
+	BNRemoveGuidedSourceBlocks(m_object, addressList, addresses.size());
+	delete[] addressList;
+}
+
+
+std::vector<ArchAndAddr> Function::GetGuidedSourceBlocks()
+{
+	size_t count;
+	BNArchitectureAndAddress* addresses = BNGetGuidedSourceBlocks(m_object, &count);
+
+	std::vector<ArchAndAddr> result;
+	result.reserve(count);
+	for (size_t i = 0; i < count; i++)
+	{
+		result.push_back({new CoreArchitecture(addresses[i].arch), addresses[i].address});
+	}
+
+	delete[] addresses;
+	return result;
+}
+
+
 vector<IndirectBranchInfo> Function::GetIndirectBranches()
 {
 	size_t count;
diff --git a/python/function.py b/python/function.py
index 2750e40a..45863f21 100644
--- a/python/function.py
+++ b/python/function.py
@@ -2295,6 +2295,51 @@ class Function:
 			branch_list[i].address = branches[i][1]
 		core.BNSetUserIndirectBranches(self.handle, source_arch.handle, source, branch_list, len(branches))
 
+	def set_guided_source_blocks(
+	    self, addresses: List[Tuple['architecture.Architecture', int]]
+	) -> None:
+		address_list = (core.BNArchitectureAndAddress * len(addresses))()
+		for i in range(len(addresses)):
+			address_list[i].arch = addresses[i][0].handle
+			address_list[i].address = addresses[i][1]
+		core.BNSetGuidedSourceBlocks(self.handle, address_list, len(addresses))
+
+	def add_guided_source_blocks(
+	    self, addresses: List[Tuple['architecture.Architecture', int]]
+	) -> None:
+		address_list = (core.BNArchitectureAndAddress * len(addresses))()
+		for i in range(len(addresses)):
+			address_list[i].arch = addresses[i][0].handle
+			address_list[i].address = addresses[i][1]
+		core.BNAddGuidedSourceBlocks(self.handle, address_list, len(addresses))
+
+	def remove_guided_source_blocks(
+	    self, addresses: List[Tuple['architecture.Architecture', int]]
+	) -> None:
+		address_list = (core.BNArchitectureAndAddress * len(addresses))()
+		for i in range(len(addresses)):
+			address_list[i].arch = addresses[i][0].handle
+			address_list[i].address = addresses[i][1]
+		core.BNRemoveGuidedSourceBlocks(self.handle, address_list, len(addresses))
+
+	def get_guided_source_blocks(
+	    self
+	) -> List[Tuple['architecture.Architecture', int]]:
+		count = ctypes.c_ulonglong()
+		addresses = core.BNGetGuidedSourceBlocks(self.handle, count)
+		try:
+			assert addresses is not None, "core.BNGetGuidedSourceBlocks returned None"
+			result = []
+			for i in range(count.value):
+				result.append((
+					architecture.CoreArchitecture._from_cache(addresses[i].arch),
+					addresses[i].address
+				))
+			return result
+		finally:
+			if addresses is not None:
+				core.BNFreeArchitectureAndAddressList(addresses)
+
 	def get_indirect_branches_at(
 	    self, addr: int, arch: Optional['architecture.Architecture'] = None
 	) -> List['variable.IndirectBranchInfo']: