api: indirect branch inlining support

7 files changed, 158 insertions, 95 deletions

diff --git a/architecture.cpp b/architecture.cpp
index bd476049..6db5b6b5 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -298,6 +298,15 @@ std::set<ArchAndAddr>& BasicBlockAnalysisContext::GetHaltedDisassemblyAddresses(
 }
 
 
+std::map<ArchAndAddr, ArchAndAddr>& BasicBlockAnalysisContext::GetInlinedUnresolvedIndirectBranches()
+{
+	if (!m_inlinedUnresolvedIndirectBranches)
+		m_inlinedUnresolvedIndirectBranches.emplace();
+
+	return *m_inlinedUnresolvedIndirectBranches;
+}
+
+
 void BasicBlockAnalysisContext::AddTempOutgoingReference(Function* targetFunc)
 {
 	BNAnalyzeBasicBlocksContextAddTempReference(m_context, targetFunc->m_object);
@@ -386,6 +395,29 @@ void BasicBlockAnalysisContext::Finalize()
 		delete[] haltedAddresses;
 	}
 
+	if (m_inlinedUnresolvedIndirectBranches)
+	{
+		auto& inlinedUnresolvedIndirectBranches = *m_inlinedUnresolvedIndirectBranches;
+
+		BNArchitectureAndAddress* locations = new BNArchitectureAndAddress[inlinedUnresolvedIndirectBranches.size() * 2];
+
+		size_t i = 0;
+		for (auto& pair : inlinedUnresolvedIndirectBranches)
+		{
+			locations[i].arch = pair.first.arch->GetObject();
+			locations[i].address = pair.first.address;
+
+			locations[i + 1].arch = pair.second.arch->GetObject();
+			locations[i + 1].address = pair.second.address;
+
+			i += 2;
+		}
+
+		BNAnalyzeBasicBlocksContextSetInlinedUnresolvedIndirectBranches(m_context, locations, inlinedUnresolvedIndirectBranches.size() * 2);
+
+		delete[] locations;
+	}
+
 	if (m_contextualReturns)
 	{
 		auto& contextualReturns = *m_contextualReturns;
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 0ca84526..2457bf6e 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -8137,6 +8137,7 @@ namespace BinaryNinja {
 		std::optional<std::map<uint64_t, std::set<ArchAndAddr>>> m_directCodeReferences;
 		std::optional<std::set<ArchAndAddr>> m_directNoReturnCalls;
 		std::optional<std::set<ArchAndAddr>> m_haltedDisassemblyAddresses;
+		std::optional<std::map<ArchAndAddr, ArchAndAddr>> m_inlinedUnresolvedIndirectBranches;
 
 	public:
 		BNBasicBlockAnalysisContext* m_context;
@@ -8160,6 +8161,7 @@ namespace BinaryNinja {
 		std::map<uint64_t, std::set<ArchAndAddr>>& GetDirectCodeReferences();
 		std::set<ArchAndAddr>& GetDirectNoReturnCalls();
 		std::set<ArchAndAddr>& GetHaltedDisassemblyAddresses();
+		std::map<ArchAndAddr, ArchAndAddr>& GetInlinedUnresolvedIndirectBranches();
 
 		void AddTempOutgoingReference(Function* targetFunc);
 
@@ -11495,7 +11497,7 @@ namespace BinaryNinja {
 
 		Ref<Function> GetCalleeForAnalysis(Ref<Platform> platform, uint64_t addr, bool exact);
 
-		std::vector<uint64_t> GetUnresolvedIndirectBranches();
+		std::vector<ArchAndAddr> GetUnresolvedIndirectBranches();
 		bool HasUnresolvedIndirectBranches();
 
 		void SetAutoCallTypeAdjustment(Architecture* arch, uint64_t addr, const Confidence<Ref<Type>>& adjust);
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 114eb191..253330c0 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -37,14 +37,14 @@
 // 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 116
+#define BN_CURRENT_CORE_ABI_VERSION 117
 
 // 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
 // will require rebuilding. The minimum version is increased when there are
 // incompatible changes that break binary compatibility, such as changes to
 // existing types or functions.
-#define BN_MINIMUM_CORE_ABI_VERSION 116
+#define BN_MINIMUM_CORE_ABI_VERSION 117
 
 #ifdef __GNUC__
 	#ifdef BINARYNINJACORE_LIBRARY
@@ -1904,6 +1904,9 @@ extern "C"
 
 		size_t haltedDisassemblyAddressesCount;
 		BNArchitectureAndAddress* haltedDisassemblyAddresses;
+
+		size_t inlinedUnresolvedIndirectBranchCount;
+		BNArchitectureAndAddress* inlinedUnresolvedIndirectBranches;
 	} BNBasicBlockAnalysisContext;
 
 	typedef struct BNCustomArchitecture
@@ -5141,7 +5144,7 @@ extern "C"
 	BINARYNINJACOREAPI BNFunction* BNGetCalleeForAnalysis(BNFunction* func, BNPlatform* platform,
 		uint64_t addr, bool exact);
 
-	BINARYNINJACOREAPI uint64_t* BNGetUnresolvedIndirectBranches(BNFunction* func, size_t* count);
+	BINARYNINJACOREAPI BNArchitectureAndAddress* BNGetUnresolvedIndirectBranches(BNFunction* func, size_t* count);
 	BINARYNINJACOREAPI bool BNHasUnresolvedIndirectBranches(BNFunction* func);
 
 	BINARYNINJACOREAPI void BNFunctionToggleRegion(BNFunction* func, uint64_t hash);
@@ -5214,6 +5217,7 @@ extern "C"
 	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetDirectNoReturnCalls(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, size_t count);
 	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetContextualFunctionReturns(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, bool* values, size_t count);
 	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetHaltedDisassemblyAddresses(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* sources, size_t count);
+	BINARYNINJACOREAPI void BNAnalyzeBasicBlocksContextSetInlinedUnresolvedIndirectBranches(BNBasicBlockAnalysisContext* abb, BNArchitectureAndAddress* locations, size_t count);
 
 	BINARYNINJACOREAPI BNAnalysisParameters BNGetParametersForAnalysis(BNBinaryView* view);
 	BINARYNINJACOREAPI void BNSetParametersForAnalysis(BNBinaryView* view, BNAnalysisParameters params);
diff --git a/defaultabb.cpp b/defaultabb.cpp
index 0847f0c8..1c4c38bb 100644
--- a/defaultabb.cpp
+++ b/defaultabb.cpp
@@ -75,6 +75,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 	auto& directRefs = context.GetDirectCodeReferences();
 	auto& directNoReturnCalls = context.GetDirectNoReturnCalls();
 	auto& haltedDisassemblyAddresses = context.GetHaltedDisassemblyAddresses();
+	auto& inlinedUnresolvedIndirectBranches = context.GetInlinedUnresolvedIndirectBranches();
 
 	bool hasInvalidInstructions = false;
 	set<ArchAndAddr> guidedSourceBlockTargets;
@@ -325,6 +326,81 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 				{
 					bool fastPath;
 
+					auto handleAsFallback = [&]() {
+						// Undefined type or target, check for targets from analysis and stop disassembling this block
+						endsBlock = true;
+
+						if (info.branchType[i] == IndirectBranch)
+						{
+							// Indirect calls need not end the block early.
+							Ref<LowLevelILFunction> ilFunc = new LowLevelILFunction(location.arch, nullptr);
+							location.arch->GetInstructionLowLevelIL(opcode, location.address, maxLen, *ilFunc);
+							for (size_t idx = 0; idx < ilFunc->GetInstructionCount(); idx++)
+							{
+								if ((*ilFunc)[idx].operation == LLIL_CALL)
+								{
+									endsBlock = false;
+									break;
+								}
+							}
+						}
+
+						indirectBranchIter = indirectBranches.find(location);
+						endIter = indirectBranches.end();
+						if (indirectBranchIter != endIter)
+						{
+							for (auto& branch : indirectBranchIter->second)
+							{
+								directRefs[branch.address].emplace(location);
+								Ref<Platform> targetPlatform = funcPlatform;
+								if (branch.arch != function->GetArchitecture())
+									targetPlatform = funcPlatform->GetRelatedPlatform(branch.arch);
+
+								// Normal analysis should not inline indirect targets that are function starts
+								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)
+								{
+									blocksToProcess.push(branch);
+									seenBlocks.insert(branch);
+								}
+							}
+						}
+						else if (info.branchType[i] == ExceptionBranch)
+						{
+							block->SetCanExit(false);
+						}
+						else if (info.branchType[i] == FunctionReturn && function->CanReturn().GetValue())
+						{
+							// Support for contextual function returns. This is mainly used for ARM/Thumb with 'blx lr'. It's most common for this to be treated
+							// as a function return, however it can also be a function call. For now this transform is described as follows:
+							// 1) Architecture lifts a call instruction as LLIL_CALL with a branch type of FunctionReturn
+							// 2) By default, contextualFunctionReturns is used to translate this to a LLIL_RET (conservative)
+							// 3) Downstream analysis uses dataflow to validate the return target
+							// 4) If the target is not the ReturnAddressValue, then we avoid the translation to a return and leave the instruction as a call
+							if (auto it = contextualFunctionReturns.find(location); it != contextualFunctionReturns.end())
+								endsBlock = it->second;
+							else
+							{
+								Ref<LowLevelILFunction> ilFunc = new LowLevelILFunction(location.arch, nullptr);
+								location.arch->GetInstructionLowLevelIL(opcode, location.address, maxLen, *ilFunc);
+								if (ilFunc->GetInstructionCount() && ((*ilFunc)[0].operation == LLIL_CALL))
+									contextualFunctionReturns[location] = true;
+							}
+						}
+						else
+						{
+							// If analysis did not find any valid branch targets, don't assume anything about global
+							// function state, such as __noreturn analysis, since we can't see the entire function->
+							block->SetUndeterminedOutgoingEdges(true);
+						}
+					};
+
 					switch (info.branchType[i])
 					{
 					case UnconditionalBranch:
@@ -375,7 +451,7 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 								calledFunctions.insert(otherFunc);
 								if (info.branchType[i] == UnconditionalBranch)
 								{
-									if (!otherFunc->CanReturn())
+									if (!otherFunc->CanReturn() && !otherFunc->IsInlinedDuringAnalysis().GetValue())
 									{
 										directNoReturnCalls.insert(location);
 										endsBlock = true;
@@ -465,98 +541,39 @@ void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnaly
 								break;
 							}
 
-							directRefs[target.address].emplace(location);
-							if (!func->CanReturn())
-							{
-								directNoReturnCalls.insert(location);
-								endsBlock = true;
-								block->SetCanExit(false);
-							}
 
 							// Add function as an early reference in case it gets updated before this
 							// function finishes analysis.
 							context.AddTempOutgoingReference(func);
 
 							calledFunctions.emplace(func);
-						}
-						break;
 
-					case SystemCall:
-						break;
-
-					default:
-						// Undefined type or target, check for targets from analysis and stop disassembling this block
-						endsBlock = true;
-
-						if (info.branchType[i] == IndirectBranch)
-						{
-							// Indirect calls need not end the block early.
-							Ref<LowLevelILFunction> ilFunc = new LowLevelILFunction(location.arch, nullptr);
-							location.arch->GetInstructionLowLevelIL(opcode, location.address, maxLen, *ilFunc);
-							for (size_t idx = 0; idx < ilFunc->GetInstructionCount(); idx++)
+							directRefs[target.address].emplace(location);
+							if (!func->CanReturn())
 							{
-								if ((*ilFunc)[idx].operation == LLIL_CALL)
+								if (func->IsInlinedDuringAnalysis().GetValue() && func->HasUnresolvedIndirectBranches())
 								{
-									endsBlock = false;
-									break;
+									auto unresolved = func->GetUnresolvedIndirectBranches();
+									if (unresolved.size() == 1)
+									{
+										inlinedUnresolvedIndirectBranches[location] = *unresolved.begin();
+										handleAsFallback();
+										break;
+									}
 								}
+
+								directNoReturnCalls.insert(location);
+								endsBlock = true;
+								block->SetCanExit(false);
 							}
 						}
+						break;
 
-						indirectBranchIter = indirectBranches.find(location);
-						endIter = indirectBranches.end();
-						if (indirectBranchIter != endIter)
-						{
-							for (auto& branch : indirectBranchIter->second)
-							{
-								directRefs[branch.address].emplace(location);
-								Ref<Platform> targetPlatform = funcPlatform;
-								if (branch.arch != function->GetArchitecture())
-									targetPlatform = funcPlatform->GetRelatedPlatform(branch.arch);
-
-								// Normal analysis should not inline indirect targets that are function starts
-								if (translateTailCalls && data->GetAnalysisFunction(targetPlatform, branch.address))
-									continue;
-
-								if (isGuidedSourceBlock)
-									guidedSourceBlockTargets.insert(branch);
+					case SystemCall:
+						break;
 
-								block->AddPendingOutgoingEdge(IndirectBranch, branch.address, branch.arch);
-								if (seenBlocks.count(branch) == 0)
-								{
-									blocksToProcess.push(branch);
-									seenBlocks.insert(branch);
-								}
-							}
-						}
-						else if (info.branchType[i] == ExceptionBranch)
-						{
-							block->SetCanExit(false);
-						}
-						else if (info.branchType[i] == FunctionReturn)
-						{
-							// Support for contextual function returns. This is mainly used for ARM/Thumb with 'blx lr'. It's most common for this to be treated
-							// as a function return, however it can also be a function call. For now this transform is described as follows:
-							// 1) Architecture lifts a call instruction as LLIL_CALL with a branch type of FunctionReturn
-							// 2) By default, contextualFunctionReturns is used to translate this to a LLIL_RET (conservative)
-							// 3) Downstream analysis uses dataflow to validate the return target
-							// 4) If the target is not the ReturnAddressValue, then we avoid the translation to a return and leave the instruction as a call
-							if (auto it = contextualFunctionReturns.find(location); it != contextualFunctionReturns.end())
-								endsBlock = it->second;
-							else
-							{
-								Ref<LowLevelILFunction> ilFunc = new LowLevelILFunction(location.arch, nullptr);
-								location.arch->GetInstructionLowLevelIL(opcode, location.address, maxLen, *ilFunc);
-								if (ilFunc->GetInstructionCount() && ((*ilFunc)[0].operation == LLIL_CALL))
-									contextualFunctionReturns[location] = true;
-							}
-						}
-						else
-						{
-							// If analysis did not find any valid branch targets, don't assume anything about global
-							// function state, such as __noreturn analysis, since we can't see the entire function->
-							block->SetUndeterminedOutgoingEdges(true);
-						}
+					default:
+						handleAsFallback();
 						break;
 					}
 				}
diff --git a/function.cpp b/function.cpp
index 9e7f4986..367dbe80 100644
--- a/function.cpp
+++ b/function.cpp
@@ -1852,13 +1852,15 @@ Ref<Function> Function::GetCalleeForAnalysis(Ref<Platform> platform, uint64_t ad
 }
 
 
-vector<uint64_t> Function::GetUnresolvedIndirectBranches()
+vector<ArchAndAddr> Function::GetUnresolvedIndirectBranches()
 {
 	size_t count;
-	uint64_t* addrs = BNGetUnresolvedIndirectBranches(m_object, &count);
-	vector<uint64_t> result;
-	result.insert(result.end(), addrs, &addrs[count]);
-	BNFreeAddressList(addrs);
+	BNArchitectureAndAddress* addresses = BNGetUnresolvedIndirectBranches(m_object, &count);
+	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});
+	BNFreeArchitectureAndAddressList(addresses);
 	return result;
 }
 
diff --git a/python/function.py b/python/function.py
index f290bebd..37bd3396 100644
--- a/python/function.py
+++ b/python/function.py
@@ -1277,18 +1277,22 @@ class Function:
 		return result
 
 	@property
-	def unresolved_indirect_branches(self) -> List[int]:
+	def unresolved_indirect_branches(self) -> List[Tuple['architecture.Architecture', int]]:
 		"""List of unresolved indirect branches (read-only)"""
 		count = ctypes.c_ulonglong()
-		addrs = core.BNGetUnresolvedIndirectBranches(self.handle, count)
-		assert addrs is not None, "core.BNGetUnresolvedIndirectBranches returned None"
+		addresses = core.BNGetUnresolvedIndirectBranches(self.handle, count)
 		try:
+			assert addresses is not None, "core.BNGetUnresolvedIndirectBranches returned None"
 			result = []
-			for i in range(0, count.value):
-				result.append(addrs[i])
+			for i in range(count.value):
+				result.append((
+					architecture.CoreArchitecture._from_cache(addresses[i].arch),
+					addresses[i].address
+				))
 			return result
 		finally:
-			core.BNFreeAddressList(addrs)
+			if addresses is not None:
+				core.BNFreeArchitectureAndAddressList(addresses)
 
 	@property
 	def has_unresolved_indirect_branches(self) -> bool:
diff --git a/rust/src/function.rs b/rust/src/function.rs
index c61e4291..de6d2d41 100644
--- a/rust/src/function.rs
+++ b/rust/src/function.rs
@@ -2372,11 +2372,13 @@ impl Function {
     }
 
     /// List of address of unresolved indirect branches
-    pub fn unresolved_indirect_branches(&self) -> Array<UnresolvedIndirectBranches> {
+    /*
+    pub fn unresolved_indirect_branches(&self) -> Array<Arch> {
         let mut count = 0;
         let result = unsafe { BNGetUnresolvedIndirectBranches(self.handle, &mut count) };
         unsafe { Array::new(result, count, ()) }
     }
+    */
 
     /// Returns a string representing the provenance. This portion of the API
     /// is under development. Currently the provenance information is