1 files changed, 88 insertions, 24 deletions

diff --git a/EfiGuardDxe/util.c b/EfiGuardDxe/util.c
index 17686ef..78424e5 100644
--- a/EfiGuardDxe/util.c
+++ b/EfiGuardDxe/util.c
@@ -119,6 +119,40 @@ PrintKernelPatchInfo(
 	}
 }
 
+VOID
+EFIAPI
+DisableWriteProtect(
+	OUT BOOLEAN *WpEnabled,
+	OUT BOOLEAN *CetEnabled
+	)
+{
+	CONST UINTN Cr0 = AsmReadCr0();
+	*WpEnabled = (Cr0 & CR0_WP) != 0;
+	*CetEnabled = (AsmReadCr4() & CR4_CET) != 0;
+
+	if (*WpEnabled)
+	{
+		if (*CetEnabled)
+			AsmDisableCet();
+		AsmWriteCr0(Cr0 & ~CR0_WP);
+	}
+}
+
+VOID
+EFIAPI
+EnableWriteProtect(
+	IN BOOLEAN WpEnabled,
+	IN BOOLEAN CetEnabled
+	)
+{
+	if (WpEnabled)
+	{
+		AsmWriteCr0(AsmReadCr0() | CR0_WP);
+		if (CetEnabled)
+			AsmEnableCet();
+	}
+}
+
 VOID*
 EFIAPI
 CopyWpMem(
@@ -127,16 +161,12 @@ CopyWpMem(
 	IN UINTN Length
 	)
 {
-	CONST UINTN Cr0 = AsmReadCr0();
-	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
-	if (WpSet)
-		AsmWriteCr0(Cr0 & ~CR0_WP);
+	BOOLEAN WpEnabled, CetEnabled;
+	DisableWriteProtect(&WpEnabled, &CetEnabled);
 
 	VOID* Result = CopyMem(Destination, Source, Length);
 	
-	if (WpSet)
-		AsmWriteCr0(Cr0);
-
+	EnableWriteProtect(WpEnabled, CetEnabled);
 	return Result;
 }
 
@@ -148,16 +178,12 @@ SetWpMem(
 	IN UINT8 Value
 	)
 {
-	CONST UINTN Cr0 = AsmReadCr0();
-	CONST BOOLEAN WpSet = (Cr0 & CR0_WP) != 0;
-	if (WpSet)
-		AsmWriteCr0(Cr0 & ~CR0_WP);
+	BOOLEAN WpEnabled, CetEnabled;
+	DisableWriteProtect(&WpEnabled, &CetEnabled);
 
 	VOID* Result = SetMem(Destination, Length, Value);
 	
-	if (WpSet)
-		AsmWriteCr0(Cr0);
-
+	EnableWriteProtect(WpEnabled, CetEnabled);
 	return Result;
 }
 
@@ -199,6 +225,43 @@ StrniCmp(
 	return UpperFirstChar - UpperSecondChar;
 }
 
+CONST CHAR16*
+EFIAPI
+StriStr(
+	IN CONST CHAR16 *String1,
+	IN CONST CHAR16 *String2
+	)
+{
+	if (*String2 == L'\0')
+		return String1;
+
+	while (*String1 != L'\0')
+	{
+		CONST CHAR16* FirstMatch = String1;
+		CONST CHAR16* String2Ptr = String2;
+		CHAR16 String1Char = CharToUpper(*String1);
+		CHAR16 String2Char = CharToUpper(*String2Ptr);
+
+		while (String1Char == String2Char && String1Char != L'\0')
+		{
+			String1++;
+			String2Ptr++;
+
+			String1Char = CharToUpper(*String1);
+			String2Char = CharToUpper(*String2Ptr);
+		}
+
+		if (String2Char == L'\0')
+			return FirstMatch;
+
+		if (String1Char == L'\0')
+			return NULL;
+
+		String1 = FirstMatch + 1;
+	}
+	return NULL;
+}
+
 BOOLEAN
 EFIAPI
 WaitForKey(
@@ -206,27 +269,28 @@ WaitForKey(
 	)
 {
 	// Hack: because we call this at TPL_NOTIFY in ExitBootServices, we cannot use WaitForEvent()
-	// in that scenario because it requires TPL == TPL_APPLICATION. So check the TPL
+	// in that scenario because it requires TPL <= TPL_APPLICATION. So check the TPL
 	CONST EFI_TPL Tpl = EfiGetCurrentTpl();
 
-	EFI_INPUT_KEY Key = { 0, 0 };
+	EFI_KEY_DATA KeyData = { 0 };
 	EFI_STATUS Status = EFI_NOT_READY;
 
 	while (Status == EFI_NOT_READY)
 	{
-		// Can we call WaitForEvent()?
 		UINTN Index = 0;
-		if (Tpl == TPL_APPLICATION)
-			gBS->WaitForEvent(1, &gST->ConIn->WaitForKey, &Index); // Yep
+		if (Tpl <= TPL_APPLICATION)
+			gBS->WaitForEvent(1, (VOID**)(gTextInputEx != NULL ? gTextInputEx->WaitForKeyEx : gST->ConIn->WaitForKey), &Index);
 		else
-			RtlStall(1); // Nope; burn CPU. // TODO: find a way to parallelize this to achieve GeForce FX 5800 temperatures
+			RtlStall(1); // WaitForEvent() unavailable, burn CPU
 
 		// At TPL_APPLICATION, we will always get EFI_SUCCESS (barring hardware failures). At higher TPLs we may also get EFI_NOT_READY
-		Status = gST->ConIn->ReadKeyStroke(gST->ConIn, &Key);
+		Status = gTextInputEx != NULL
+			? gTextInputEx->ReadKeyStrokeEx(gTextInputEx, &KeyData)
+			: gST->ConIn->ReadKeyStroke(gST->ConIn, &KeyData.Key);
 	}
 
 	ASSERT_EFI_ERROR(Status);
-	return (BOOLEAN)(Key.ScanCode != SCAN_ESC);
+	return KeyData.Key.ScanCode != SCAN_ESC;
 }
 
 INT32
@@ -475,8 +539,8 @@ BacktrackToFunctionStart(
 
 	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)
+		// If the function entry specifies indirection, get the address of its master function entry
+		while ((FunctionEntry->u.UnwindData & RUNTIME_FUNCTION_INDIRECT) != 0)
 		{
 			FunctionEntry = (PIMAGE_RUNTIME_FUNCTION_ENTRY)(FunctionEntry->u.UnwindData + ImageBase - 1);
 		}