EfiDSEFix: improve g_CiOptions address validationv1.2.1

- Verify expected lengths of instructions
- Verify CipInitialize is in PAGE
- Verify g_CiOptions is in either .data or CiPolicy

Fixes #31 (regression due to KB5003173 fix)

3 files changed, 79 insertions, 20 deletions

diff --git a/Application/EfiDSEFix/src/EfiDSEFix.cpp b/Application/EfiDSEFix/src/EfiDSEFix.cpp
index 324018d..bfbe561 100644
--- a/Application/EfiDSEFix/src/EfiDSEFix.cpp
+++ b/Application/EfiDSEFix/src/EfiDSEFix.cpp
@@ -75,7 +75,7 @@ static
 LONG
 QueryCiOptions(
 	_In_ PVOID MappedBase,
-	_In_ ULONG_PTR KernelBase,
+	_In_ ULONG_PTR CiDllBase,
 	_Out_ PULONG_PTR gCiOptionsAddress
 	)
 {
@@ -85,6 +85,10 @@ QueryCiOptions(
 	LONG Relative = 0;
 	hde64s hs;
 
+	const PIMAGE_NT_HEADERS NtHeaders = RtlImageNtHeader(MappedBase);
+	if (NtHeaders == nullptr)
+		return 0;
+
 	const PUCHAR CiInitialize = static_cast<PUCHAR>(GetProcedureAddress(reinterpret_cast<ULONG_PTR>(MappedBase), "CiInitialize"));
 	if (CiInitialize == nullptr)
 		return 0;
@@ -95,8 +99,12 @@ QueryCiOptions(
 		ULONG j = 0;
 		do
 		{
+			hde64_disasm(CiInitialize + i, &hs);
+			if (hs.flags & F_ERROR)
+				break;
+
 			// call CipInitialize
-			const BOOLEAN IsCall = CiInitialize[i] == 0xE8;
+			const BOOLEAN IsCall = hs.len == 5 && CiInitialize[i] == 0xE8;
 			if (IsCall)
 				j++;
 
@@ -104,16 +112,15 @@ QueryCiOptions(
 			{
 				Relative = *reinterpret_cast<PLONG>(CiInitialize + i + 1);
 
-				// KB5003173 added a new 'call wil_InitializeFeatureStaging' to CiInitialize that we need to skip
-				const PUCHAR CallTarget = CiInitialize + i + 5 + Relative;
-				hde64_disasm(CallTarget, &hs);
-				if ((hs.flags & F_ERROR) == 0 && hs.len >= 4 && hs.len <= 6) // wil_InitializeFeatureStaging: 3, __security_init_cookie: 7, CipInitialize: 5
+				// Check the call target to skip calls to __security_init_cookie, wil_InitializeFeatureStaging, and other stuff in INIT. CipInitialize is in PAGE.
+				const PUCHAR CallTarget = CiInitialize + i + hs.len + Relative;
+				if (AddressIsInSection(static_cast<PUCHAR>(MappedBase), CallTarget, NtHeaders, "PAGE"))
+				{
 					break;
+				}
+				Relative = 0;
 			}
 
-			hde64_disasm(CiInitialize + i, &hs);
-			if (hs.flags & F_ERROR)
-				break;
 			i += hs.len;
 
 		} while (i < 256);
@@ -123,39 +130,54 @@ QueryCiOptions(
 		i = 0;
 		do
 		{
+			hde64_disasm(CiInitialize + i, &hs);
+			if (hs.flags & F_ERROR)
+				break;
+
 			// jmp CipInitialize
-			if (CiInitialize[i] == 0xE9)
+			if (hs.len == 5 && CiInitialize[i] == 0xE9)
 			{
 				Relative = *reinterpret_cast<PLONG>(CiInitialize + i + 1);
 				break;
 			}
-			hde64_disasm(CiInitialize + i, &hs);
-			if (hs.flags & F_ERROR)
-				break;
+
 			i += hs.len;
 
 		} while (i < 256);
 	}
 
-	const PUCHAR CipInitialize = CiInitialize + i + 5 + Relative;
+	if (Relative == 0)
+		return 0;
+
+	const PUCHAR CipInitialize = CiInitialize + i + hs.len + Relative;
+	if (!AddressIsInSection(static_cast<PUCHAR>(MappedBase), CipInitialize, NtHeaders, "PAGE"))
+		return 0;
+
 	i = 0;
 	do
 	{
-		if (*reinterpret_cast<PUSHORT>(CipInitialize + i) == 0x0d89)
+		hde64_disasm(CipInitialize + i, &hs);
+		if (hs.flags & F_ERROR)
+			break;
+
+		if (hs.len == 6 && *reinterpret_cast<PUSHORT>(CipInitialize + i) == 0x0d89) // mov g_CiOptions, ecx
 		{
 			Relative = *reinterpret_cast<PLONG>(CipInitialize + i + 2);
 			break;
 		}
-		hde64_disasm(CipInitialize + i, &hs);
-		if (hs.flags & F_ERROR)
-			break;
+
 		i += hs.len;
 
 	} while (i < 256);
 
-	const PUCHAR MappedCiOptions = CipInitialize + i + 6 + Relative;
+	const PUCHAR MappedCiOptions = CipInitialize + i + hs.len + Relative;
+
+	// g_CiOptions is in .data or (newer builds) "CiPolicy"
+	if (!AddressIsInSection(static_cast<PUCHAR>(MappedBase), MappedCiOptions, NtHeaders, ".data") &&
+		!AddressIsInSection(static_cast<PUCHAR>(MappedBase), MappedCiOptions, NtHeaders, "CiPolicy"))
+		return 0;
 
-	*gCiOptionsAddress = KernelBase + MappedCiOptions - static_cast<PUCHAR>(MappedBase);
+	*gCiOptionsAddress = CiDllBase + MappedCiOptions - static_cast<PUCHAR>(MappedBase);
 
 	return Relative;
 }
diff --git a/Application/EfiDSEFix/src/EfiDSEFix.h b/Application/EfiDSEFix/src/EfiDSEFix.h
index 0a99540..e7cb109 100644
--- a/Application/EfiDSEFix/src/EfiDSEFix.h
+++ b/Application/EfiDSEFix/src/EfiDSEFix.h
@@ -44,6 +44,14 @@ MapFileSectionView(
 	_Out_ PSIZE_T ViewSize
 	);
 
+BOOLEAN
+AddressIsInSection(
+	_In_ PUCHAR ImageBase,
+	_In_ PUCHAR Address,
+	_In_ PIMAGE_NT_HEADERS NtHeaders,
+	_In_ PCCH SectionName
+	);
+
 PVOID
 GetProcedureAddress(
 	_In_ ULONG_PTR DllBase,
diff --git a/Application/EfiDSEFix/src/pe.cpp b/Application/EfiDSEFix/src/pe.cpp
index 2187537..97d2333 100644
--- a/Application/EfiDSEFix/src/pe.cpp
+++ b/Application/EfiDSEFix/src/pe.cpp
@@ -110,6 +110,35 @@ MapFileSectionView(
 	return Status;
 }
 
+BOOLEAN
+AddressIsInSection(
+	_In_ PUCHAR ImageBase,
+	_In_ PUCHAR Address,
+	_In_ PIMAGE_NT_HEADERS NtHeaders,
+	_In_ PCCH SectionName
+	)
+{
+	if (ImageBase > Address)
+		return FALSE;
+	if (Address >= ImageBase + NtHeaders->OptionalHeader.SizeOfImage)
+		return FALSE;
+
+	const ULONG Rva = static_cast<ULONG>(static_cast<ULONG_PTR>(Address - ImageBase));
+	PIMAGE_SECTION_HEADER Section = IMAGE_FIRST_SECTION(NtHeaders);
+	const USHORT NumberOfSections = NtHeaders->FileHeader.NumberOfSections;
+	for (USHORT i = 0; i < NumberOfSections; ++i)
+	{
+		if (Section->VirtualAddress <= Rva &&
+			Section->VirtualAddress + Section->Misc.VirtualSize > Rva)
+		{
+			if (strncmp(reinterpret_cast<PCCH>(Section->Name), SectionName, IMAGE_SIZEOF_SHORT_NAME) == 0)
+				return TRUE;
+		}
+		Section++;
+	}
+	return FALSE;
+}
+
 PVOID
 GetProcedureAddress(
 	_In_ ULONG_PTR DllBase,