Add system register list to architectures.

Global registers are commonly used for module data pointers in
architectures, which are global in nature but should still be eliminated
in high level ILs once the data pointers have been resolved. This means
there needs to be an additional type of global register for those that
should never be eliminated in any way, such as an MSR register or other
registers that contain important system state with side effects. To
allow for this, a list of system registers can be supplied by the
archtiecture.

4 files changed, 82 insertions, 0 deletions

diff --git a/architecture.cpp b/architecture.cpp
index 42323a6c..a3bd77fa 100644
--- a/architecture.cpp
+++ b/architecture.cpp
@@ -530,6 +530,19 @@ uint32_t* Architecture::GetGlobalRegistersCallback(void* ctxt, size_t* count)
 }
 
 
+uint32_t* Architecture::GetSystemRegistersCallback(void* ctxt, size_t* count)
+{
+	Architecture* arch = (Architecture*)ctxt;
+	vector<uint32_t> regs = arch->GetSystemRegisters();
+	*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;
+}
+
+
 char* Architecture::GetRegisterStackNameCallback(void* ctxt, uint32_t regStack)
 {
 	Architecture* arch = (Architecture*)ctxt;
@@ -756,6 +769,7 @@ void Architecture::Register(Architecture* arch)
 	callbacks.getStackPointerRegister = GetStackPointerRegisterCallback;
 	callbacks.getLinkRegister = GetLinkRegisterCallback;
 	callbacks.getGlobalRegisters = GetGlobalRegistersCallback;
+	callbacks.getSystemRegisters = GetSystemRegistersCallback;
 	callbacks.getRegisterStackName = GetRegisterStackNameCallback;
 	callbacks.getAllRegisterStacks = GetAllRegisterStacksCallback;
 	callbacks.getRegisterStackInfo = GetRegisterStackInfoCallback;
@@ -1035,12 +1049,24 @@ vector<uint32_t> Architecture::GetGlobalRegisters()
 }
 
 
+vector<uint32_t> Architecture::GetSystemRegisters()
+{
+	return vector<uint32_t>();
+}
+
+
 bool Architecture::IsGlobalRegister(uint32_t reg)
 {
 	return BNIsArchitectureGlobalRegister(m_object, reg);
 }
 
 
+bool Architecture::IsSystemRegister(uint32_t reg)
+{
+	return BNIsArchitectureSystemRegister(m_object, reg);
+}
+
+
 string Architecture::GetRegisterStackName(uint32_t regStack)
 {
 	char regStr[32];
@@ -1647,6 +1673,21 @@ vector<uint32_t> CoreArchitecture::GetGlobalRegisters()
 }
 
 
+vector<uint32_t> CoreArchitecture::GetSystemRegisters()
+{
+	size_t count;
+	uint32_t* regs = BNGetArchitectureSystemRegisters(m_object, &count);
+
+	vector<uint32_t> result;
+	result.reserve(count);
+	for (size_t i = 0; i < count; i++)
+		result.push_back(regs[i]);
+
+	BNFreeRegisterList(regs);
+	return result;
+}
+
+
 string CoreArchitecture::GetRegisterStackName(uint32_t regStack)
 {
 	char* name = BNGetArchitectureRegisterStackName(m_object, regStack);
@@ -2020,6 +2061,12 @@ vector<uint32_t> ArchitectureExtension::GetGlobalRegisters()
 }
 
 
+vector<uint32_t> ArchitectureExtension::GetSystemRegisters()
+{
+	return m_base->GetSystemRegisters();
+}
+
+
 string ArchitectureExtension::GetRegisterStackName(uint32_t regStack)
 {
 	return m_base->GetRegisterStackName(regStack);
@@ -2156,6 +2203,7 @@ void ArchitectureHook::Register(BNCustomArchitecture* callbacks)
 {
 	AddRefForRegistration();
 	m_object = BNRegisterArchitectureHook(m_base->GetObject(), callbacks);
+	BNFinalizeArchitectureHook(m_base->GetObject());
 }
 
 
diff --git a/binaryninjaapi.h b/binaryninjaapi.h
index af68eafd..d0e23d28 100644
--- a/binaryninjaapi.h
+++ b/binaryninjaapi.h
@@ -2064,6 +2064,7 @@ __attribute__ ((format (printf, 1, 2)))
 		static uint32_t GetStackPointerRegisterCallback(void* ctxt);
 		static uint32_t GetLinkRegisterCallback(void* ctxt);
 		static uint32_t* GetGlobalRegistersCallback(void* ctxt, size_t* count);
+		static uint32_t* GetSystemRegistersCallback(void* ctxt, size_t* count);
 
 		static char* GetRegisterStackNameCallback(void* ctxt, uint32_t regStack);
 		static uint32_t* GetAllRegisterStacksCallback(void* ctxt, size_t* count);
@@ -2151,6 +2152,8 @@ __attribute__ ((format (printf, 1, 2)))
 		virtual uint32_t GetLinkRegister();
 		virtual std::vector<uint32_t> GetGlobalRegisters();
 		bool IsGlobalRegister(uint32_t reg);
+		virtual std::vector<uint32_t> GetSystemRegisters();
+		bool IsSystemRegister(uint32_t reg);
 		std::vector<uint32_t> GetModifiedRegistersOnWrite(uint32_t reg);
 		uint32_t GetRegisterByName(const std::string& name);
 
@@ -2303,6 +2306,7 @@ __attribute__ ((format (printf, 1, 2)))
 		virtual uint32_t GetStackPointerRegister() override;
 		virtual uint32_t GetLinkRegister() override;
 		virtual std::vector<uint32_t> GetGlobalRegisters() override;
+		virtual std::vector<uint32_t> GetSystemRegisters() override;
 
 		virtual std::string GetRegisterStackName(uint32_t regStack) override;
 		virtual std::vector<uint32_t> GetAllRegisterStacks() override;
@@ -2377,6 +2381,7 @@ __attribute__ ((format (printf, 1, 2)))
 		virtual uint32_t GetStackPointerRegister() override;
 		virtual uint32_t GetLinkRegister() override;
 		virtual std::vector<uint32_t> GetGlobalRegisters() override;
+		virtual std::vector<uint32_t> GetSystemRegisters() override;
 
 		virtual std::string GetRegisterStackName(uint32_t regStack) override;
 		virtual std::vector<uint32_t> GetAllRegisterStacks() override;
diff --git a/binaryninjacore.h b/binaryninjacore.h
index b1a615f5..97cdf853 100644
--- a/binaryninjacore.h
+++ b/binaryninjacore.h
@@ -1456,6 +1456,7 @@ extern "C"
 		uint32_t (*getStackPointerRegister)(void* ctxt);
 		uint32_t (*getLinkRegister)(void* ctxt);
 		uint32_t* (*getGlobalRegisters)(void* ctxt, size_t* count);
+		uint32_t* (*getSystemRegisters)(void* ctxt, size_t* count);
 
 		char* (*getRegisterStackName)(void* ctxt, uint32_t regStack);
 		uint32_t* (*getAllRegisterStacks)(void* ctxt, size_t* count);
@@ -2673,6 +2674,7 @@ __attribute__ ((format (printf, 1, 2)))
 		BNArchitecture* base, BNCustomArchitecture* arch);
 	BINARYNINJACOREAPI void BNAddArchitectureRedirection(BNArchitecture* arch, BNArchitecture* from, BNArchitecture* to);
 	BINARYNINJACOREAPI BNArchitecture* BNRegisterArchitectureHook(BNArchitecture* base, BNCustomArchitecture* arch);
+	BINARYNINJACOREAPI void BNFinalizeArchitectureHook(BNArchitecture* base);
 
 	BINARYNINJACOREAPI char* BNGetArchitectureName(BNArchitecture* arch);
 	BINARYNINJACOREAPI BNEndianness BNGetArchitectureEndianness(BNArchitecture* arch);
@@ -2730,6 +2732,8 @@ __attribute__ ((format (printf, 1, 2)))
 	BINARYNINJACOREAPI uint32_t BNGetArchitectureLinkRegister(BNArchitecture* arch);
 	BINARYNINJACOREAPI uint32_t* BNGetArchitectureGlobalRegisters(BNArchitecture* arch, size_t* count);
 	BINARYNINJACOREAPI bool BNIsArchitectureGlobalRegister(BNArchitecture* arch, uint32_t reg);
+	BINARYNINJACOREAPI uint32_t* BNGetArchitectureSystemRegisters(BNArchitecture* arch, size_t* count);
+	BINARYNINJACOREAPI bool BNIsArchitectureSystemRegister(BNArchitecture* arch, uint32_t reg);
 	BINARYNINJACOREAPI uint32_t BNGetArchitectureRegisterByName(BNArchitecture* arch, const char* name);
 
 	BINARYNINJACOREAPI char* BNGetArchitectureRegisterStackName(BNArchitecture* arch, uint32_t regStack);
diff --git a/python/architecture.py b/python/architecture.py
index a6e1a083..ad295c91 100644
--- a/python/architecture.py
+++ b/python/architecture.py
@@ -122,6 +122,7 @@ class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 	stack_pointer = None
 	link_reg = None
 	global_regs = []
+	system_regs = []
 	flags = []
 	flag_write_types = []
 	semantic_flag_classes = []
@@ -194,6 +195,7 @@ class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 			self._get_stack_pointer_register)
 		self._cb.getLinkRegister = self._cb.getLinkRegister.__class__(self._get_link_register)
 		self._cb.getGlobalRegisters = self._cb.getGlobalRegisters.__class__(self._get_global_registers)
+		self._cb.getSystemRegisters = self._cb.getSystemRegisters.__class__(self._get_system_registers)
 		self._cb.getRegisterStackName = self._cb.getRegisterStackName.__class__(self._get_register_stack_name)
 		self._cb.getAllRegisterStacks = self._cb.getAllRegisterStacks.__class__(self._get_all_register_stacks)
 		self._cb.getRegisterStackInfo = self._cb.getRegisterStackInfo.__class__(self._get_register_stack_info)
@@ -360,6 +362,7 @@ class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 			self._semantic_class_for_flag_write_type[self._flag_write_types[write_type]] = sem_class_index
 
 		self.__dict__["global_regs"] = self.__class__.global_regs
+		self.__dict__["system_regs"] = self.__class__.system_regs
 
 		self._intrinsics = {}
 		self._intrinsics_by_index = {}
@@ -926,6 +929,20 @@ class Architecture(with_metaclass(_ArchitectureMetaClass, object)):
 			count[0] = 0
 			return None
 
+	def _get_system_registers(self, ctxt, count):
+		try:
+			count[0] = len(self.system_regs)
+			reg_buf = (ctypes.c_uint * len(self.system_regs))()
+			for i in range(0, len(self.system_regs)):
+				reg_buf[i] = self._all_regs[self.system_regs[i]]
+			result = ctypes.cast(reg_buf, ctypes.c_void_p)
+			self._pending_reg_lists[result.value] = (result, reg_buf)
+			return result.value
+		except KeyError:
+			log.log_error(traceback.format_exc())
+			count[0] = 0
+			return None
+
 	def _get_register_stack_name(self, ctxt, reg_stack):
 		try:
 			if reg_stack in self._reg_stacks_by_index:
@@ -2231,6 +2248,13 @@ class CoreArchitecture(Architecture):
 		core.BNFreeRegisterList(regs)
 
 		count = ctypes.c_ulonglong()
+		regs = core.BNGetArchitectureSystemRegisters(self.handle, count)
+		self.__dict__["system_regs"] = []
+		for i in range(0, count.value):
+			self.system_regs.append(core.BNGetArchitectureRegisterName(self.handle, regs[i]))
+		core.BNFreeRegisterList(regs)
+
+		count = ctypes.c_ulonglong()
 		regs = core.BNGetAllArchitectureRegisterStacks(self.handle, count)
 		self._all_reg_stacks = {}
 		self._reg_stacks_by_index = {}
@@ -2701,6 +2725,7 @@ class ArchitectureHook(CoreArchitecture):
 	def register(self):
 		self.__class__._registered_cb = self._cb
 		self.handle = core.BNRegisterArchitectureHook(self._base_arch.handle, self._cb)
+		core.BNFinalizeArchitectureHook(self._base_arch.handle)
 
 	@property
 	def base_arch(self):