Use PE runtime function tables for finding function start addressesv1.1.1

6 files changed, 74 insertions, 50 deletions

diff --git a/EfiGuardDxe/PatchBootmgr.c b/EfiGuardDxe/PatchBootmgr.c
index aa27135..5c3c481 100644
--- a/EfiGuardDxe/PatchBootmgr.c
+++ b/EfiGuardDxe/PatchBootmgr.c
@@ -300,7 +300,7 @@ PatchBootManager(
 	// Found signature; backtrack to function start
 	// Note: pOriginalAddress is a pointer to a (function) pointer, because the original address depends on the type of boot manager we are patching.
 	VOID **pOriginalAddress = PatchingBootmgrEfi ? &gOriginalBootmgrImgArchStartBootApplication : &gOriginalBootmgfwImgArchStartBootApplication;
-	*pOriginalAddress = (VOID*)BacktrackToFunctionStart(Found, MAX((UINT8*)ImageBase + CodeSection->VirtualAddress, Found - 1024));
+	*pOriginalAddress = (VOID*)BacktrackToFunctionStart((UINT8*)ImageBase, NtHeaders, Found);
 	CONST VOID* OriginalAddress = *pOriginalAddress;
 	if (OriginalAddress == NULL)
 	{
diff --git a/EfiGuardDxe/PatchNtoskrnl.c b/EfiGuardDxe/PatchNtoskrnl.c
index 1ddb0c8..ac8f181 100644
--- a/EfiGuardDxe/PatchNtoskrnl.c
+++ b/EfiGuardDxe/PatchNtoskrnl.c
@@ -110,8 +110,7 @@ DisablePatchGuard(
 	}
 
 	// Backtrack to function start
-	UINT8* KeInitAmd64SpecificState = BacktrackToFunctionStart(KeInitAmd64SpecificStatePatternAddress,
-		(UINT8*)(KeInitAmd64SpecificStatePatternAddress - StartVa));
+	UINT8* KeInitAmd64SpecificState = BacktrackToFunctionStart(ImageBase, NtHeaders, KeInitAmd64SpecificStatePatternAddress);
 	if (KeInitAmd64SpecificState == NULL)
 	{
 		PRINT_KERNEL_PATCH_MSG(L"    Failed to find KeInitAmd64SpecificState%S.\r\n",
@@ -202,8 +201,7 @@ DisablePatchGuard(
 	}
 
 	// Backtrack to function start
-	UINT8* CcInitializeBcbProfiler = BacktrackToFunctionStart(CcInitializeBcbProfilerPatternAddress,
-		(UINT8*)(CcInitializeBcbProfilerPatternAddress - StartVa));
+	UINT8* CcInitializeBcbProfiler = BacktrackToFunctionStart(ImageBase, NtHeaders, CcInitializeBcbProfilerPatternAddress);
 	if (CcInitializeBcbProfiler == NULL)
 	{
 		PRINT_KERNEL_PATCH_MSG(L"    Failed to find %S%S.\r\n",
@@ -249,8 +247,7 @@ DisablePatchGuard(
 		}
 
 		// Backtrack to function start
-		ExpLicenseWatchInitWorker = BacktrackToFunctionStart(ExpLicenseWatchInitWorkerPatternAddress,
-			(UINT8*)(ExpLicenseWatchInitWorkerPatternAddress - StartVa));
+		ExpLicenseWatchInitWorker = BacktrackToFunctionStart(ImageBase, NtHeaders, ExpLicenseWatchInitWorkerPatternAddress);
 		if (ExpLicenseWatchInitWorker == NULL)
 		{
 			PRINT_KERNEL_PATCH_MSG(L"    Failed to find ExpLicenseWatchInitWorker%S.\r\n",
@@ -279,8 +276,7 @@ DisablePatchGuard(
 		PRINT_KERNEL_PATCH_MSG(L"    Found KiVerifyScopesExecute pattern at 0x%llX.\r\n", (UINTN)KiVerifyScopesExecutePatternAddress);
 
 		// Backtrack to function start
-		KiVerifyScopesExecute = BacktrackToFunctionStart(KiVerifyScopesExecutePatternAddress,
-			(UINT8*)(KiVerifyScopesExecutePatternAddress - StartVa));
+		KiVerifyScopesExecute = BacktrackToFunctionStart(ImageBase, NtHeaders, KiVerifyScopesExecutePatternAddress);
 		if (KiVerifyScopesExecute == NULL)
 		{
 			PRINT_KERNEL_PATCH_MSG(L"    Failed to find KiVerifyScopesExecute.\r\n");
@@ -352,10 +348,8 @@ DisablePatchGuard(
 		}
 
 		// Backtrack to function start
-		KiMcaDeferredRecoveryServiceCallers[0] = BacktrackToFunctionStart(KiMcaDeferredRecoveryServiceCallers[0],
-			(UINT8*)(KiMcaDeferredRecoveryServiceCallers[0] - StartVa));
-		KiMcaDeferredRecoveryServiceCallers[1] = BacktrackToFunctionStart(KiMcaDeferredRecoveryServiceCallers[1],
-			(UINT8*)(KiMcaDeferredRecoveryServiceCallers[1] - StartVa));
+		KiMcaDeferredRecoveryServiceCallers[0] = BacktrackToFunctionStart(ImageBase, NtHeaders, KiMcaDeferredRecoveryServiceCallers[0]);
+		KiMcaDeferredRecoveryServiceCallers[1] = BacktrackToFunctionStart(ImageBase, NtHeaders, KiMcaDeferredRecoveryServiceCallers[1]);
 		if (KiMcaDeferredRecoveryServiceCallers[0] == NULL || KiMcaDeferredRecoveryServiceCallers[1] == NULL)
 		{
 			PRINT_KERNEL_PATCH_MSG(L"    Failed to find KiMcaDeferredRecoveryService callers.\r\n");
diff --git a/EfiGuardDxe/PatchWinload.c b/EfiGuardDxe/PatchWinload.c
index 2182fb9..a36f8b6 100644
--- a/EfiGuardDxe/PatchWinload.c
+++ b/EfiGuardDxe/PatchWinload.c
@@ -191,7 +191,7 @@ PatchImgpValidateImageHash(
 	}
 
 	// Backtrack to function start
-	CONST UINT8* ImgpValidateImageHash = BacktrackToFunctionStart(AndMinusFortyOneAddress, CodeStartVa);
+	CONST UINT8* ImgpValidateImageHash = BacktrackToFunctionStart(ImageBase, NtHeaders, AndMinusFortyOneAddress);
 	if (ImgpValidateImageHash == NULL)
 	{
 		Print(L"    Failed to find %S!ImgpValidateImageHash%S.\r\n",
@@ -329,7 +329,7 @@ PatchImgpFilterValidationFailure(
 	}
 
 	// Backtrack to function start
-	CONST UINT8* ImgpFilterValidationFailure = BacktrackToFunctionStart(LeaIntegrityFailureAddress, LeaIntegrityFailureAddress - Length);
+	CONST UINT8* ImgpFilterValidationFailure = BacktrackToFunctionStart(ImageBase, NtHeaders, LeaIntegrityFailureAddress);
 	if (ImgpFilterValidationFailure == NULL)
 	{
 		Print(L"    Failed to find %S!ImgpFilterValidationFailure%S.\r\n",
@@ -380,7 +380,7 @@ FindOslFwpKernelSetupPhase1(
 		if (!EFI_ERROR(Status))
 		{
 			// Found signature; backtrack to function start
-			*OslFwpKernelSetupPhase1Address = BacktrackToFunctionStart(Found, Found - 0x400);
+			*OslFwpKernelSetupPhase1Address = BacktrackToFunctionStart(ImageBase, NtHeaders, Found);
 			if (*OslFwpKernelSetupPhase1Address != NULL)
 			{
 				Print(L"\r\nFound OslFwpKernelSetupPhase1 at 0x%llX.\r\n", (UINTN)(*OslFwpKernelSetupPhase1Address));
@@ -479,7 +479,7 @@ FindOslFwpKernelSetupPhase1(
 		return EFI_NOT_FOUND;
 	}
 
-	CONST UINT8* EfipGetRsdt = BacktrackToFunctionStart(LeaEfiAcpiTableGuidAddress, LeaEfiAcpiTableGuidAddress - Length);
+	CONST UINT8* EfipGetRsdt = BacktrackToFunctionStart(ImageBase, NtHeaders, LeaEfiAcpiTableGuidAddress);
 	if (EfipGetRsdt == NULL)
 	{
 		Print(L"    Failed to find EfipGetRsdt.\r\n");
@@ -516,7 +516,7 @@ FindOslFwpKernelSetupPhase1(
 				OperandAddress == (UINTN)EfipGetRsdt)
 			{
 				// Calculate the distance from the start of the function to the instruction. OslFwpKernelSetupPhase1 will always have the shortest distance
-				CONST UINTN StartOfFunction = (UINTN)BacktrackToFunctionStart((UINT8*)InstructionAddress, (UINT8*)InstructionAddress - Length);
+				CONST UINTN StartOfFunction = (UINTN)BacktrackToFunctionStart(ImageBase, NtHeaders, (UINT8*)InstructionAddress);
 				CONST UINTN Distance = InstructionAddress - StartOfFunction;
 				if (Distance < ShortestDistanceToCall)
 				{
diff --git a/EfiGuardDxe/pe.h b/EfiGuardDxe/pe.h
index cf47119..1104dfd 100644
--- a/EfiGuardDxe/pe.h
+++ b/EfiGuardDxe/pe.h
@@ -40,6 +40,8 @@ typedef EFI_IMAGE_EXPORT_DIRECTORY *PEFI_IMAGE_EXPORT_DIRECTORY;
 #define VS_VERSION_INFO									1
 #define VS_FF_DEBUG										(0x00000001L)
 
+#define RUNTIME_FUNCTION_INDIRECT						0x1
+
 #define IMAGE32(NtHeaders) ((NtHeaders)->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC)
 #define IMAGE64(NtHeaders) ((NtHeaders)->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC)
 
@@ -173,6 +175,21 @@ typedef struct _VS_VERSIONINFO
 
 
 //
+// Function table entry data
+//
+typedef struct _RUNTIME_FUNCTION
+{
+	UINT32 BeginAddress;
+	UINT32 EndAddress;
+	union
+	{
+		UINT32 UnwindInfoAddress;
+		UINT32 UnwindData;
+	} u;
+} RUNTIME_FUNCTION, *PRUNTIME_FUNCTION;
+
+
+//
 // Function declarations
 //
 PEFI_IMAGE_NT_HEADERS
diff --git a/EfiGuardDxe/util.c b/EfiGuardDxe/util.c
index 6420a56..6d45675 100644
--- a/EfiGuardDxe/util.c
+++ b/EfiGuardDxe/util.c
@@ -78,6 +78,7 @@ AppendKernelPatchMessage(
 VOID
 EFIAPI
 PrintKernelPatchInfo(
+	VOID
 	)
 {
 	ASSERT(gST->ConOut != NULL);
@@ -101,6 +102,7 @@ PrintKernelPatchInfo(
 BOOLEAN
 EFIAPI
 WaitForKey(
+	VOID
 	)
 {
 	// Hack: because we call this at TPL_NOTIFY in ExitBootServices, we cannot use WaitForEvent()
@@ -339,42 +341,51 @@ ZydisInit(
 UINT8*
 EFIAPI
 BacktrackToFunctionStart(
-	IN CONST UINT8* StartAddress,
-	IN CONST UINT8* LowerBound
+	IN CONST UINT8* ImageBase,
+	IN PEFI_IMAGE_NT_HEADERS NtHeaders,
+	IN CONST UINT8* AddressInFunction
 	)
 {
 	// Test for null. This allows callers to do 'FindPattern(..., &Address); X = Backtrack(Address, ...)' with a single failure branch
-	if (StartAddress == NULL)
+	if (AddressInFunction == NULL)
+		return NULL;
+	if (NtHeaders->OptionalHeader.NumberOfRvaAndSizes <= EFI_IMAGE_DIRECTORY_ENTRY_EXCEPTION)
 		return NULL;
 
-	ASSERT(StartAddress > LowerBound);
+	CONST PRUNTIME_FUNCTION FunctionTable = (PRUNTIME_FUNCTION)(ImageBase + NtHeaders->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_EXCEPTION].VirtualAddress);
+	CONST UINT32 FunctionTableSize = NtHeaders->OptionalHeader.DataDirectory[EFI_IMAGE_DIRECTORY_ENTRY_EXCEPTION].Size;
+	if (FunctionTableSize == 0)
+		return NULL;
 
-	UINT8 *Address;
-	BOOLEAN Found = FALSE;
-	for (Address = (UINT8*)StartAddress; Address >= LowerBound; --Address)
+	// Do a binary search until we find the function that contains our address
+	CONST UINT32 RelativeAddress = (UINT32)(AddressInFunction - ImageBase);
+	PRUNTIME_FUNCTION FunctionEntry = NULL;
+	INT32 Low = 0;
+	INT32 High = (FunctionTableSize / sizeof(RUNTIME_FUNCTION)) - 1;
+	
+	while (High >= Low)
 	{
-		if ((*(Address - 1) == 0xCC ||										// Previous byte is int 3 padding, or
-			(*(Address - 2) == 0x90 && *(Address - 1) == 0x90) ||			// Previous 2 bytes are nop padding, or
-			(*(Address - 4) == 0x00 && *(Address - 3) == 0x00 &&			// Previous 4+ bytes are 00 padding (rare, only happens at start of a section), or
-				*(Address - 2) == 0x00 && *(Address - 1) == 0x00) ||
-			(*(Address - 1) == 0xC3 && *(Address - 3) != 0x8D)				// Previous byte is 'ret', or
-#if defined(MDE_CPU_IA32) || defined(_M_IX86)
-			|| (*(Address - 3) == 0xC2 && *(Address - 1) == 0x00)			// Previous 3 bytes are 'ret XX' (x86)
-#endif
-			)
-			&&																// *and*
-			(*Address == 0x40 || *Address == 0x55 ||						// Current byte is either 'push [ebp|ebx|rbp|rbx]', 'mov REG, XX' or 'sub REG, XX'
-			(Address < StartAddress && *Address == 0x44 && *(Address + 1) == 0x89) ||
-			(Address < StartAddress && *Address == 0x48 && *(Address + 1) == 0x83) ||
-			(Address < StartAddress && *Address == 0x48 && *(Address + 1) == 0x89) ||
-			(Address < StartAddress && *Address == 0x48 && *(Address + 1) == 0x8B) ||
-			(Address < StartAddress && *Address == 0x49 && *(Address + 1) == 0x89) ||
-			(Address < StartAddress && *Address == 0x4C && *(Address + 1) == 0x8B)))
-		{
-			Found = TRUE;
+		CONST INT32 Middle = (Low + High) >> 1;
+		FunctionEntry = &FunctionTable[Middle];
+
+		if (RelativeAddress < FunctionEntry->BeginAddress)
+			High = Middle - 1;
+		else if (RelativeAddress >= FunctionEntry->EndAddress)
+			Low = Middle + 1;
+		else
 			break;
+	}
+
+	if (High >= Low)
+	{
+		// If the function entry specifies indirection, get the address of the master function entry
+		if ((FunctionEntry->u.UnwindData & RUNTIME_FUNCTION_INDIRECT) != 0)
+		{
+			FunctionEntry = (PRUNTIME_FUNCTION)(FunctionEntry->u.UnwindData + ImageBase - 1);
 		}
+		
+		return (UINT8*)ImageBase + FunctionEntry->BeginAddress;
 	}
 
-	return Found ? Address : NULL;
+	return NULL;
 }
diff --git a/EfiGuardDxe/util.h b/EfiGuardDxe/util.h
index f25c3d6..c306a0e 100644
--- a/EfiGuardDxe/util.h
+++ b/EfiGuardDxe/util.h
@@ -42,6 +42,7 @@ AppendKernelPatchMessage(
 VOID
 EFIAPI
 PrintKernelPatchInfo(
+	VOID
 	);
 
 //
@@ -51,6 +52,7 @@ PrintKernelPatchInfo(
 BOOLEAN
 EFIAPI
 WaitForKey(
+	VOID
 	);
 
 //
@@ -108,13 +110,13 @@ ZydisInit(
 	);
 
 //
-// Finds the start of a function given an address within it, scanning downwards.
-// Returns NULL if StartAddress is NULL (this simplifies error checking logic in calling functions).
-// Returns NULL if LowerBound is reached and no function boundary was found.
+// Finds the start of a function given an address within it.
+// Returns NULL if AddressInFunction is NULL (this simplifies error checking logic in calling functions).
 //
 UINT8*
 EFIAPI
 BacktrackToFunctionStart(
-	IN CONST UINT8* StartAddress,
-	IN CONST UINT8* LowerBound
+	IN CONST UINT8* ImageBase,
+	IN PEFI_IMAGE_NT_HEADERS NtHeaders,
+	IN CONST UINT8* AddressInFunction
 	);