Allow calling conventions to specify a list of registers that are required to be considered for heuristic calling convention detection

Author: D0ntPanic

arch/x86/arch_x86.cpp

@@ -3814,6 +3814,11 @@ class X86ThiscallCallingConvention: public X86BaseCallingConvention
 		return vector<uint32_t>{ XED_REG_ECX };
 	}
 
+	virtual vector<uint32_t> GetRequiredArgumentRegisters() override
+	{
+		return vector<uint32_t>{ XED_REG_ECX };
+	}
+
 	virtual bool IsStackAdjustedOnReturn() override
 	{
 		return true;

binaryninjaapi.h

@@ -17488,6 +17488,8 @@ namespace BinaryNinja {
 		static uint32_t* GetCalleeSavedRegistersCallback(void* ctxt, size_t* count);
 		static uint32_t* GetIntegerArgumentRegistersCallback(void* ctxt, size_t* count);
 		static uint32_t* GetFloatArgumentRegistersCallback(void* ctxt, size_t* count);
+		static uint32_t* GetRequiredArgumentRegistersCallback(void* ctxt, size_t* count);
+		static uint32_t* GetRequiredClobberedRegistersCallback(void* ctxt, size_t* count);
 		static void FreeRegisterListCallback(void* ctxt, uint32_t* regs, size_t len);
 
 		static bool AreArgumentRegistersSharedIndexCallback(void* ctxt);
@@ -17520,6 +17522,21 @@ namespace BinaryNinja {
 
 		virtual std::vector<uint32_t> GetIntegerArgumentRegisters();
 		virtual std::vector<uint32_t> GetFloatArgumentRegisters();
+
+		/*! Gets the set of registers that must be arguments for heuristic calling convention
+			detection to consider this calling convention as a valid option.
+
+			\return The set of registers that must be arguments
+		*/
+		virtual std::vector<uint32_t> GetRequiredArgumentRegisters();
+
+		/*! Gets the set of registers that must be clobbered for heuristic calling convention
+			detection to consider this calling convention as a valid option.
+
+			\return The set of registers that must be clobbered
+		*/
+		virtual std::vector<uint32_t> GetRequiredClobberedRegisters();
+
 		virtual bool AreArgumentRegistersSharedIndex();
 		virtual bool AreArgumentRegistersUsedForVarArgs();
 		virtual bool IsStackReservedForArgumentRegisters();
@@ -17552,6 +17569,8 @@ namespace BinaryNinja {
 
 		virtual std::vector<uint32_t> GetIntegerArgumentRegisters() override;
 		virtual std::vector<uint32_t> GetFloatArgumentRegisters() override;
+		virtual std::vector<uint32_t> GetRequiredArgumentRegisters() override;
+		virtual std::vector<uint32_t> GetRequiredClobberedRegisters() override;
 		virtual bool AreArgumentRegistersSharedIndex() override;
 		virtual bool AreArgumentRegistersUsedForVarArgs() override;
 		virtual bool IsStackReservedForArgumentRegisters() override;

binaryninjacore.h

@@ -2833,6 +2833,8 @@ extern "C"
 		uint32_t* (*getCalleeSavedRegisters)(void* ctxt, size_t* count);
 		uint32_t* (*getIntegerArgumentRegisters)(void* ctxt, size_t* count);
 		uint32_t* (*getFloatArgumentRegisters)(void* ctxt, size_t* count);
+		uint32_t* (*getRequiredArgumentRegisters)(void* ctxt, size_t* count);
+		uint32_t* (*getRequiredClobberedRegisters)(void* ctxt, size_t* count);
 		void (*freeRegisterList)(void* ctxt, uint32_t* regs, size_t len);
 
 		bool (*areArgumentRegistersSharedIndex)(void* ctxt);
@@ -7496,6 +7498,8 @@ extern "C"
 
 	BINARYNINJACOREAPI uint32_t* BNGetIntegerArgumentRegisters(BNCallingConvention* cc, size_t* count);
 	BINARYNINJACOREAPI uint32_t* BNGetFloatArgumentRegisters(BNCallingConvention* cc, size_t* count);
+	BINARYNINJACOREAPI uint32_t* BNGetRequiredArgumentRegisters(BNCallingConvention* cc, size_t* count);
+	BINARYNINJACOREAPI uint32_t* BNGetRequiredClobberedRegisters(BNCallingConvention* cc, size_t* count);
 	BINARYNINJACOREAPI bool BNAreArgumentRegistersSharedIndex(BNCallingConvention* cc);
 	BINARYNINJACOREAPI bool BNAreArgumentRegistersUsedForVarArgs(BNCallingConvention* cc);
 	BINARYNINJACOREAPI bool BNIsStackReservedForArgumentRegisters(BNCallingConvention* cc);

callingconvention.cpp

@@ -39,6 +39,8 @@ CallingConvention::CallingConvention(Architecture* arch, const string& name)
 	cc.getCalleeSavedRegisters = GetCalleeSavedRegistersCallback;
 	cc.getIntegerArgumentRegisters = GetIntegerArgumentRegistersCallback;
 	cc.getFloatArgumentRegisters = GetFloatArgumentRegistersCallback;
+	cc.getRequiredArgumentRegisters = GetRequiredArgumentRegistersCallback;
+	cc.getRequiredClobberedRegisters = GetRequiredClobberedRegistersCallback;
 	cc.freeRegisterList = FreeRegisterListCallback;
 	cc.areArgumentRegistersSharedIndex = AreArgumentRegistersSharedIndexCallback;
 	cc.areArgumentRegistersUsedForVarArgs = AreArgumentRegistersUsedForVarArgsCallback;
@@ -119,6 +121,32 @@ uint32_t* CallingConvention::GetFloatArgumentRegistersCallback(void* ctxt, size_
 }
 
 
+uint32_t* CallingConvention::GetRequiredArgumentRegistersCallback(void* ctxt, size_t* count)
+{
+	CallbackRef<CallingConvention> cc(ctxt);
+	vector<uint32_t> regs = cc->GetRequiredArgumentRegisters();
+	*count = regs.size();
+
+	uint32_t* result = new uint32_t[regs.size()];
+	for (size_t i = 0; i < regs.size(); i++)
+		result[i] = regs[i];
+	return result;
+}
+
+
+uint32_t* CallingConvention::GetRequiredClobberedRegistersCallback(void* ctxt, size_t* count)
+{
+	CallbackRef<CallingConvention> cc(ctxt);
+	vector<uint32_t> regs = cc->GetRequiredClobberedRegisters();
+	*count = regs.size();
+
+	uint32_t* result = new uint32_t[regs.size()];
+	for (size_t i = 0; i < regs.size(); i++)
+		result[i] = regs[i];
+	return result;
+}
+
+
 void CallingConvention::FreeRegisterListCallback(void*, uint32_t* regs, size_t)
 {
 	delete[] regs;
@@ -284,6 +312,18 @@ vector<uint32_t> CallingConvention::GetFloatArgumentRegisters()
 }
 
 
+vector<uint32_t> CallingConvention::GetRequiredArgumentRegisters()
+{
+	return vector<uint32_t>();
+}
+
+
+vector<uint32_t> CallingConvention::GetRequiredClobberedRegisters()
+{
+	return vector<uint32_t>();
+}
+
+
 bool CallingConvention::AreArgumentRegistersSharedIndex()
 {
 	return false;
@@ -417,6 +457,28 @@ vector<uint32_t> CoreCallingConvention::GetFloatArgumentRegisters()
 }
 
 
+vector<uint32_t> CoreCallingConvention::GetRequiredArgumentRegisters()
+{
+	size_t count;
+	uint32_t* regs = BNGetRequiredArgumentRegisters(m_object, &count);
+	vector<uint32_t> result;
+	result.insert(result.end(), regs, &regs[count]);
+	BNFreeRegisterList(regs);
+	return result;
+}
+
+
+vector<uint32_t> CoreCallingConvention::GetRequiredClobberedRegisters()
+{
+	size_t count;
+	uint32_t* regs = BNGetRequiredClobberedRegisters(m_object, &count);
+	vector<uint32_t> result;
+	result.insert(result.end(), regs, &regs[count]);
+	BNFreeRegisterList(regs);
+	return result;
+}
+
+
 bool CoreCallingConvention::AreArgumentRegistersSharedIndex()
 {
 	return BNAreArgumentRegistersSharedIndex(m_object);

python/callingconvention.py

@@ -40,6 +40,8 @@ class CallingConvention:
 	callee_saved_regs = []
 	int_arg_regs = []
 	float_arg_regs = []
+	required_arg_regs = []
+	required_clobbered_regs = []
 	arg_regs_share_index = False
 	arg_regs_for_varargs = True
 	stack_reserved_for_arg_regs = False
@@ -70,6 +72,12 @@ def __init__(
 			    self._get_int_arg_regs
 			)
 			self._cb.getFloatArgumentRegisters = self._cb.getFloatArgumentRegisters.__class__(self._get_float_arg_regs)
+			self._cb.getRequiredArgumentRegisters = self._cb.getRequiredArgumentRegisters.__class__(
+				self._get_required_arg_regs
+			)
+			self._cb.getRequiredClobberedRegisters = self._cb.getRequiredClobberedRegisters.__class__(
+				self._get_required_clobbered_regs
+			)
 			self._cb.freeRegisterList = self._cb.freeRegisterList.__class__(self._free_register_list)
 			self._cb.areArgumentRegistersSharedIndex = self._cb.areArgumentRegistersSharedIndex.__class__(
 			    self._arg_regs_share_index
@@ -161,6 +169,26 @@ def __init__(
 			core.BNFreeRegisterList(regs)
 			self.__dict__["float_arg_regs"] = result
 
+			count = ctypes.c_ulonglong()
+			regs = core.BNGetRequiredArgumentRegisters(handle, count)
+			assert regs is not None, "core.BNGetRequiredArgumentRegisters returned None"
+			result = []
+			arch = self.arch
+			for i in range(0, count.value):
+				result.append(arch.get_reg_name(regs[i]))
+			core.BNFreeRegisterList(regs)
+			self.__dict__["required_arg_regs"] = result
+
+			count = ctypes.c_ulonglong()
+			regs = core.BNGetRequiredClobberedRegisters(handle, count)
+			assert regs is not None, "core.BNGetRequiredClobberedRegisters returned None"
+			result = []
+			arch = self.arch
+			for i in range(0, count.value):
+				result.append(arch.get_reg_name(regs[i]))
+			core.BNFreeRegisterList(regs)
+			self.__dict__["required_clobbered_regs"] = result
+
 			reg = core.BNGetIntegerReturnValueRegister(_handle)
 			if reg == 0xffffffff:
 				self.__dict__["int_return_reg"] = None
@@ -281,6 +309,36 @@ def _get_float_arg_regs(self, ctxt, count):
 			count[0] = 0
 			return None
 
+	def _get_required_arg_regs(self, ctxt, count):
+		try:
+			regs = self.__class__.required_arg_regs
+			count[0] = len(regs)
+			reg_buf = (ctypes.c_uint * len(regs))()
+			for i in range(0, len(regs)):
+				reg_buf[i] = self.arch.regs[regs[i]].index
+			result = ctypes.cast(reg_buf, ctypes.c_void_p)
+			self._pending_reg_lists[result.value] = (result, reg_buf)
+			return result.value
+		except:
+			log_error_for_exception("Unhandled Python exception in CallingConvention._get_required_arg_regs")
+			count[0] = 0
+			return None
+
+	def _get_required_clobbered_regs(self, ctxt, count):
+		try:
+			regs = self.__class__.required_clobbered_regs
+			count[0] = len(regs)
+			reg_buf = (ctypes.c_uint * len(regs))()
+			for i in range(0, len(regs)):
+				reg_buf[i] = self.arch.regs[regs[i]].index
+			result = ctypes.cast(reg_buf, ctypes.c_void_p)
+			self._pending_reg_lists[result.value] = (result, reg_buf)
+			return result.value
+		except:
+			log_error_for_exception("Unhandled Python exception in CallingConvention._get_required_clobbered_regs")
+			count[0] = 0
+			return None
+
 	def _free_register_list(self, ctxt, regs, count):
 		try:
 			buf = ctypes.cast(regs, ctypes.c_void_p)