Add callee saved registers on calling conventions to specify which registers are assumed to be preserved (used in heuristics to avoid incomplete data flow cloberring these registers)

4 files changed, 62 insertions, 0 deletions

diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index 901ea65d..b3f1c87a 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -3597,6 +3597,7 @@ namespace BinaryNinja
 		static void FreeCallback(void* ctxt);
 
 		static uint32_t* GetCallerSavedRegistersCallback(void* ctxt, size_t* count);
+		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 void FreeRegisterListCallback(void* ctxt, uint32_t* regs);
@@ -3624,6 +3625,7 @@ namespace BinaryNinja
 		std::string GetName() const;
 
 		virtual std::vector<uint32_t> GetCallerSavedRegisters();
+		virtual std::vector<uint32_t> GetCalleeSavedRegisters();
 
 		virtual std::vector<uint32_t> GetIntegerArgumentRegisters();
 		virtual std::vector<uint32_t> GetFloatArgumentRegisters();
@@ -3650,6 +3652,7 @@ namespace BinaryNinja
 		CoreCallingConvention(BNCallingConvention* cc);
 
 		virtual std::vector<uint32_t> GetCallerSavedRegisters() override;
+		virtual std::vector<uint32_t> GetCalleeSavedRegisters() override;
 
 		virtual std::vector<uint32_t> GetIntegerArgumentRegisters() override;
 		virtual std::vector<uint32_t> GetFloatArgumentRegisters() override;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index 0916112f..90d0ec2b 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1455,6 +1455,7 @@ extern "C"
 		void (*freeObject)(void* ctxt);
 
 		uint32_t* (*getCallerSavedRegisters)(void* ctxt, size_t* count);
+		uint32_t* (*getCalleeSavedRegisters)(void* ctxt, size_t* count);
 		uint32_t* (*getIntegerArgumentRegisters)(void* ctxt, size_t* count);
 		uint32_t* (*getFloatArgumentRegisters)(void* ctxt, size_t* count);
 		void (*freeRegisterList)(void* ctxt, uint32_t* regs);
@@ -3228,6 +3229,7 @@ extern "C"
 	BINARYNINJACOREAPI BNArchitecture* BNGetCallingConventionArchitecture(BNCallingConvention* cc);
 	BINARYNINJACOREAPI char* BNGetCallingConventionName(BNCallingConvention* cc);
 	BINARYNINJACOREAPI uint32_t* BNGetCallerSavedRegisters(BNCallingConvention* cc, size_t* count);
+	BINARYNINJACOREAPI uint32_t* BNGetCalleeSavedRegisters(BNCallingConvention* cc, size_t* count);
 
 	BINARYNINJACOREAPI uint32_t* BNGetIntegerArgumentRegisters(BNCallingConvention* cc, size_t* count);
 	BINARYNINJACOREAPI uint32_t* BNGetFloatArgumentRegisters(BNCallingConvention* cc, size_t* count);
diff --git a/callingconvention.cpp b/callingconvention.cpp
index a56b8f11..9e162ed9 100644
--- a/callingconvention.cpp
+++ b/callingconvention.cpp
@@ -36,6 +36,7 @@ CallingConvention::CallingConvention(Architecture* arch, const string& name)
 	cc.context = this;
 	cc.freeObject = FreeCallback;
 	cc.getCallerSavedRegisters = GetCallerSavedRegistersCallback;
+	cc.getCalleeSavedRegisters = GetCalleeSavedRegistersCallback;
 	cc.getIntegerArgumentRegisters = GetIntegerArgumentRegistersCallback;
 	cc.getFloatArgumentRegisters = GetFloatArgumentRegistersCallback;
 	cc.freeRegisterList = FreeRegisterListCallback;
@@ -77,6 +78,19 @@ uint32_t* CallingConvention::GetCallerSavedRegistersCallback(void* ctxt, size_t*
 }
 
 
+uint32_t* CallingConvention::GetCalleeSavedRegistersCallback(void* ctxt, size_t* count)
+{
+	CallingConvention* cc = (CallingConvention*)ctxt;
+	vector<uint32_t> regs = cc->GetCalleeSavedRegisters();
+	*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::GetIntegerArgumentRegistersCallback(void* ctxt, size_t* count)
 {
 	CallingConvention* cc = (CallingConvention*)ctxt;
@@ -234,6 +248,12 @@ vector<uint32_t> CallingConvention::GetCallerSavedRegisters()
 }
 
 
+vector<uint32_t> CallingConvention::GetCalleeSavedRegisters()
+{
+	return vector<uint32_t>();
+}
+
+
 vector<uint32_t> CallingConvention::GetIntegerArgumentRegisters()
 {
 	return vector<uint32_t>();
@@ -336,6 +356,17 @@ vector<uint32_t> CoreCallingConvention::GetCallerSavedRegisters()
 }
 
 
+vector<uint32_t> CoreCallingConvention::GetCalleeSavedRegisters()
+{
+	size_t count;
+	uint32_t* regs = BNGetCalleeSavedRegisters(m_object, &count);
+	vector<uint32_t> result;
+	result.insert(result.end(), regs, &regs[count]);
+	BNFreeRegisterList(regs);
+	return result;
+}
+
+
 vector<uint32_t> CoreCallingConvention::GetIntegerArgumentRegisters()
 {
 	size_t count;
diff --git a/python/callingconvention.py b/python/callingconvention.py
index 6b906184..f4349816 100644
--- a/python/callingconvention.py
+++ b/python/callingconvention.py
@@ -34,6 +34,7 @@ from binaryninja import range
 class CallingConvention(object):
 	name = None
 	caller_saved_regs = []
+	callee_saved_regs = []
 	int_arg_regs = []
 	float_arg_regs = []
 	arg_regs_share_index = False
@@ -56,6 +57,7 @@ class CallingConvention(object):
 			self._cb = core.BNCustomCallingConvention()
 			self._cb.context = 0
 			self._cb.getCallerSavedRegisters = self._cb.getCallerSavedRegisters.__class__(self._get_caller_saved_regs)
+			self._cb.getCalleeSavedRegisters = self._cb.getCalleeSavedRegisters.__class__(self._get_callee_saved_regs)
 			self._cb.getIntegerArgumentRegisters = self._cb.getIntegerArgumentRegisters.__class__(self._get_int_arg_regs)
 			self._cb.getFloatArgumentRegisters = self._cb.getFloatArgumentRegisters.__class__(self._get_float_arg_regs)
 			self._cb.freeRegisterList = self._cb.freeRegisterList.__class__(self._free_register_list)
@@ -91,6 +93,15 @@ class CallingConvention(object):
 			self.__dict__["caller_saved_regs"] = result
 
 			count = ctypes.c_ulonglong()
+			regs = core.BNGetCalleeSavedRegisters(self.handle, count)
+			result = []
+			arch = self.arch
+			for i in range(0, count.value):
+				result.append(arch.get_reg_name(regs[i]))
+			core.BNFreeRegisterList(regs, count.value)
+			self.__dict__["callee_saved_regs"] = result
+
+			count = ctypes.c_ulonglong()
 			regs = core.BNGetIntegerArgumentRegisters(self.handle, count)
 			result = []
 			arch = self.arch
@@ -171,6 +182,21 @@ class CallingConvention(object):
 			count[0] = 0
 			return None
 
+	def _get_callee_saved_regs(self, ctxt, count):
+		try:
+			regs = self.__class__.callee_saved_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.log_error(traceback.format_exc())
+			count[0] = 0
+			return None
+
 	def _get_int_arg_regs(self, ctxt, count):
 		try:
 			regs = self.__class__.int_arg_regs