From 63536b58536277d6c837e2b0025777385a56548d Mon Sep 17 00:00:00 2001
From: Zichuan Li <34680029+river-li@users.noreply.github.com>
Date: Thu, 25 Jul 2024 17:37:40 -0400
Subject: Change types for IDTR

UEFI PEI modules often use IDTR to fetch the pointer to the pointer to
EFI_PEI_SERVICES. We can leverage binja's global register type for this
---
 platform/efi/platform_efi.cpp | 75 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)

(limited to 'platform/efi/platform_efi.cpp')

diff --git a/platform/efi/platform_efi.cpp b/platform/efi/platform_efi.cpp
index eb743b96..4f30cc35 100644
--- a/platform/efi/platform_efi.cpp
+++ b/platform/efi/platform_efi.cpp
@@ -11,11 +11,16 @@ Ref<Platform> g_efiX86Windows, g_efiX64Windows, g_efiArm64Windows;
 
 class EFIX86Platform : public Platform
 {
+	uint32_t m_idtr;
+	Ref<Type> m_idtrtype;
+
 public:
 	EFIX86Platform(Architecture* arch) : Platform(arch, "efi-x86")
 	{
 		Ref<CallingConvention> cc;
 
+		m_idtr = arch->GetRegisterByName("idtr");
+
 		cc = arch->GetCallingConventionByName("cdecl");
 		if (cc)
 		{
@@ -50,16 +55,36 @@ public:
 			return g_efiX86;
 		return nullptr;
 	}
+
+	virtual void BinaryViewInit(BinaryView* view) override
+	{
+		if (!m_idtrtype)
+			m_idtrtype = Type::NamedType(QualifiedName("IDTR32"), GetTypeByName(QualifiedName("IDTR32")));
+	}
+
+	virtual Ref<Type> GetGlobalRegisterType(uint32_t reg) override
+	{
+		if (reg == m_idtr)
+		    return m_idtrtype;
+
+		return nullptr;
+	}
+
 };
 
 
 class EFIX86WindowsPlatform : public Platform
 {
+	uint32_t m_idtr;
+	Ref<Type> m_idtrtype;
+
 public:
 	EFIX86WindowsPlatform(Architecture* arch) : Platform(arch, "efi-windows-x86")
 	{
 		Ref<CallingConvention> cc;
 
+		m_idtr = arch->GetRegisterByName("idtr");
+
 		cc = arch->GetCallingConventionByName("cdecl");
 		if (cc)
 		{
@@ -94,16 +119,34 @@ public:
 			return g_efiX86Windows;
 		return nullptr;
 	}
+
+	virtual void BinaryViewInit(BinaryView* view) override
+	{
+		if (!m_idtrtype)
+			m_idtrtype = Type::NamedType(QualifiedName("IDTR32"), GetTypeByName(QualifiedName("IDTR32")));
+	}
+
+	virtual Ref<Type> GetGlobalRegisterType(uint32_t reg) override
+	{
+		if (reg == m_idtr)
+		    return m_idtrtype;
+
+		return nullptr;
+	}
 };
 
 
 class EFIX64Platform : public Platform
 {
+	uint32_t m_idtr;
+	Ref<Type> m_idtrtype;
+
 public:
 	EFIX64Platform(Architecture* arch) : Platform(arch, "efi-x86_64")
 	{
 		Ref<CallingConvention> cc;
 
+		m_idtr = arch->GetRegisterByName("idtr");
 		cc = arch->GetCallingConventionByName("win64");
 		if (cc)
 		{
@@ -123,16 +166,34 @@ public:
 			return g_efiX64;
 		return nullptr;
 	}
+
+	virtual void BinaryViewInit(BinaryView* view) override
+	{
+		if (!m_idtrtype)
+			m_idtrtype = Type::NamedType(QualifiedName("IDTR64"), GetTypeByName(QualifiedName("IDTR64")));
+	}
+
+	virtual Ref<Type> GetGlobalRegisterType(uint32_t reg) override
+	{
+		if (reg == m_idtr)
+		    return m_idtrtype;
+
+		return nullptr;
+	}
 };
 
 
 class EFIX64WindowsPlatform : public Platform
 {
+	uint32_t m_idtr;
+	Ref<Type> m_idtrtype;
+
 public:
 	EFIX64WindowsPlatform(Architecture* arch) : Platform(arch, "efi-windows-x86_64")
 	{
 		Ref<CallingConvention> cc;
 
+		m_idtr = arch->GetRegisterByName("idtr");
 		cc = arch->GetCallingConventionByName("win64");
 		if (cc)
 		{
@@ -152,6 +213,20 @@ public:
 			return g_efiX64Windows;
 		return nullptr;
 	}
+
+	virtual void BinaryViewInit(BinaryView* view) override
+	{
+		if (!m_idtrtype)
+			m_idtrtype = Type::NamedType(QualifiedName("IDTR64"), GetTypeByName(QualifiedName("IDTR64")));
+	}
+
+	virtual Ref<Type> GetGlobalRegisterType(uint32_t reg) override
+	{
+		if (reg == m_idtr)
+		    return m_idtrtype;
+
+		return nullptr;
+	}
 };
 
 
-- 
cgit v1.3.1