Fix inline return handling for nonstandard return registers

Avoid treating LLIL_JUMP(reg) as an inline return when the inlined callee
writes to that register before returning. Only suppress normal return-address
setup after proving the callee actually returns through an unmodified
caller-recorded return-address register

Also share the call-return rewrite path with direct LLIL_JUMP forms that have
call-like semantics, so nonstandard direct-call patterns get the same return,
jump, and tail-call handling as LLIL_CALL

Author: zznop

defaultarch.cpp

@@ -48,6 +48,37 @@ static bool IsReturnAddressRegisterExpr(LowLevelILInstruction expr, const set<ui
 }
 
 
+static void RemoveWrittenReturnAddressRegisters(LowLevelILInstruction instr, set<uint32_t>& returnAddressRegisters)
+{
+	switch (instr.operation)
+	{
+	case LLIL_SET_REG:
+		returnAddressRegisters.erase(instr.GetDestRegister<LLIL_SET_REG>());
+		break;
+	case LLIL_SET_REG_SPLIT:
+		returnAddressRegisters.erase(instr.GetHighRegister<LLIL_SET_REG_SPLIT>());
+		returnAddressRegisters.erase(instr.GetLowRegister<LLIL_SET_REG_SPLIT>());
+		break;
+	default:
+		break;
+	}
+}
+
+
+static bool IsReturnAddressRegisterJumpOrReturn(LowLevelILInstruction instr, const set<uint32_t>& returnAddressRegisters)
+{
+	switch (instr.operation)
+	{
+	case LLIL_JUMP:
+		return IsReturnAddressRegisterExpr(instr.GetDestExpr<LLIL_JUMP>(), returnAddressRegisters);
+	case LLIL_RET:
+		return IsReturnAddressRegisterExpr(instr.GetDestExpr<LLIL_RET>(), returnAddressRegisters);
+	default:
+		return false;
+	}
+}
+
+
 void Architecture::DefaultAnalyzeBasicBlocks(Function* function, BasicBlockAnalysisContext& context)
 {
 	auto data = function->GetView();
@@ -1033,11 +1064,34 @@ void FunctionLifterContext::CheckForInlinedCall(BasicBlock* block, size_t instrC
 				if (IsConstantPointer(instr.GetSourceExpr<LLIL_SET_REG>(), addr))
 					returnAddressRegisters.insert(instr.GetDestRegister<LLIL_SET_REG>());
 			}
+			bool hasCallSemantics = lastInstr.operation == LLIL_CALL
+				|| (lastInstr.operation == LLIL_JUMP && !returnAddressRegisters.empty());
 
-			if (lastInstr.operation == LLIL_CALL && returnAddressRegisters.empty())
+			// Copy the inlined code from the target function
+			Ref<Architecture> callArch = block->GetArchitecture();
+			auto blocks = PrepareToCopyForeignFunction(targetIL);
+			auto unresolvedIndirectBranches = targetFunc->GetUnresolvedIndirectBranches();
+			auto sourceLocation = inlineDuringAnalysis == InlineUsingCallAddress ? ILSourceLocation(lastInstr) : ILSourceLocation();
+			set<uint32_t> unmodifiedReturnAddressRegisters = returnAddressRegisters;
+			bool calleeReturnsThroughCallerReturnAddressRegister = false;
+			for (auto& block : blocks)
+			{
+				for (size_t instrIndex = block->GetStart(); instrIndex < block->GetEnd(); instrIndex++)
+					RemoveWrittenReturnAddressRegisters(targetIL->GetInstruction(instrIndex), unmodifiedReturnAddressRegisters);
+			}
+			for (auto& block : blocks)
+			{
+				for (size_t instrIndex = block->GetStart(); instrIndex < block->GetEnd(); instrIndex++)
+				{
+					if (IsReturnAddressRegisterJumpOrReturn(
+						targetIL->GetInstruction(instrIndex), unmodifiedReturnAddressRegisters))
+						calleeReturnsThroughCallerReturnAddressRegister = true;
+				}
+			}
+
+			if (hasCallSemantics && !calleeReturnsThroughCallerReturnAddressRegister)
 			{
 				// Set up return address according to the architecture
-				// TODO: Handle architectures that use a nonstandard way of calling functions
 				uint32_t linkReg = m_platform->GetArchitecture()->GetLinkRegister();
 				if (linkReg == BN_INVALID_REGISTER)
 				{
@@ -1056,20 +1110,14 @@ void FunctionLifterContext::CheckForInlinedCall(BasicBlock* block, size_t instrC
 
 					uint64_t addrToSet = addr;
 					if (block->GetArchitecture()->GetName() == "thumb2")
-						addrToSet |= 1; // XXX: hack moved here from lowlevelilfunction.cpp
+						addrToSet |= 1;
 
 					ExprId linkExpr = m_function->SetRegister(
 						regInfo.size, linkReg, m_function->ConstPointer(regInfo.size, addrToSet, lastInstr), 0, lastInstr);
 					m_function->SetExprAttributes(linkExpr, ILAllowDeadStoreElimination);
 					m_function->AddInstruction(linkExpr);
 				}
 			}
-
-			// Copy the inlined code from the target function
-			Ref<Architecture> callArch = block->GetArchitecture();
-			auto blocks = PrepareToCopyForeignFunction(targetIL);
-			auto unresolvedIndirectBranches = targetFunc->GetUnresolvedIndirectBranches();
-			auto sourceLocation = inlineDuringAnalysis == InlineUsingCallAddress ? ILSourceLocation(lastInstr) : ILSourceLocation();
 			for (auto& block : blocks)
 			{
 				m_function->PrepareToCopyBlock(block);
@@ -1078,7 +1126,7 @@ void FunctionLifterContext::CheckForInlinedCall(BasicBlock* block, size_t instrC
 					LowLevelILInstruction instr = targetIL->GetInstruction(instrIndex);
 					ArchAndAddr loc(block->GetArchitecture(), instr.address);
 
-					if (lastInstr.operation == LLIL_CALL && instr.operation == LLIL_RET)
+					if (hasCallSemantics && instr.operation == LLIL_RET)
 					{
 						// If the instruction is a return, emit the computation of the target
 						// location (it may affect the stack pointer) but go directly to the
@@ -1089,14 +1137,15 @@ void FunctionLifterContext::CheckForInlinedCall(BasicBlock* block, size_t instrC
 						m_function->AddInstruction(instr.GetDestExpr<LLIL_RET>().CopyTo(m_function, sourceLocation));
 						m_function->AddInstruction(m_function->Goto(end, sourceLocation));
 					}
-					else if (lastInstr.operation == LLIL_CALL && instr.operation == LLIL_JUMP
+					else if (hasCallSemantics && instr.operation == LLIL_JUMP
 						&& (block->GetArchitecture() == callArch)
 						&& (IsConstantPointer(instr.GetDestExpr<LLIL_JUMP>(), addr)
-							|| IsReturnAddressRegisterExpr(instr.GetDestExpr<LLIL_JUMP>(), returnAddressRegisters)))
+							|| IsReturnAddressRegisterExpr(instr.GetDestExpr<LLIL_JUMP>(),
+								unmodifiedReturnAddressRegisters)))
 					{
 						m_function->AddInstruction(m_function->Goto(end, sourceLocation));
 					}
-					else if (lastInstr.operation == LLIL_CALL && instr.operation == LLIL_JUMP
+					else if (hasCallSemantics && instr.operation == LLIL_JUMP
 						&& block->GetOutgoingEdges().empty() && (unresolvedIndirectBranches.count(loc) == 0))
 					{
 						// Jump without outgoing edges in the graph, and it is not marked as having