Use EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL if available

6 files changed, 55 insertions, 19 deletions

diff --git a/Application/Loader/Loader.c b/Application/Loader/Loader.c
index 4dc9e22..ed18f33 100644
--- a/Application/Loader/Loader.c
+++ b/Application/Loader/Loader.c
@@ -36,6 +36,7 @@ STATIC CHAR16* mDriverPaths[] = {
 	L"\\" EFIGUARD_DRIVER_FILENAME
 };
 
+STATIC EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *mTextInputEx = NULL;
 
 VOID
 EFIAPI
@@ -57,12 +58,19 @@ WaitForKey(
 	VOID
 	)
 {
-	EFI_INPUT_KEY Key = { 0, 0 };
+	EFI_KEY_DATA KeyData = { 0 };
 	UINTN Index = 0;
-	gBS->WaitForEvent(1, &gST->ConIn->WaitForKey, &Index);
-	gST->ConIn->ReadKeyStroke(gST->ConIn, &Key);
-
-	return Key.ScanCode != SCAN_ESC;
+	if (mTextInputEx != NULL)
+	{
+		gBS->WaitForEvent(1, &mTextInputEx->WaitForKeyEx, &Index);
+		mTextInputEx->ReadKeyStrokeEx(mTextInputEx, &KeyData);
+	}
+	else
+	{
+		gBS->WaitForEvent(1, &gST->ConIn->WaitForKey, &Index);
+		gST->ConIn->ReadKeyStroke(gST->ConIn, &KeyData.Key);
+	}
+	return KeyData.Key.ScanCode != SCAN_ESC;
 }
 
 #if CONFIGURE_DRIVER
@@ -82,12 +90,20 @@ PromptInput(
 	{
 		SelectedChar = CHAR_NULL;
 
-		EFI_INPUT_KEY Key = { 0, 0 };
+		EFI_KEY_DATA KeyData = { 0 };
 		UINTN Index = 0;
-		gBS->WaitForEvent(1, &gST->ConIn->WaitForKey, &Index);
-		gST->ConIn->ReadKeyStroke(gST->ConIn, &Key);
+		if (mTextInputEx != NULL)
+		{
+			gBS->WaitForEvent(1, &mTextInputEx->WaitForKeyEx, &Index);
+			mTextInputEx->ReadKeyStrokeEx(mTextInputEx, &KeyData);
+		}
+		else
+		{
+			gBS->WaitForEvent(1, &gST->ConIn->WaitForKey, &Index);
+			gST->ConIn->ReadKeyStroke(gST->ConIn, &KeyData.Key);
+		}
 
-		if (Key.UnicodeChar == CHAR_LINEFEED || Key.UnicodeChar == CHAR_CARRIAGE_RETURN)
+		if (KeyData.Key.UnicodeChar == CHAR_LINEFEED || KeyData.Key.UnicodeChar == CHAR_CARRIAGE_RETURN)
 		{
 			SelectedChar = DefaultSelection;
 			break;
@@ -95,9 +111,9 @@ PromptInput(
 
 		for (UINTN i = 0; i < NumAcceptedChars; ++i)
 		{
-			if (Key.UnicodeChar == AcceptedChars[i])
+			if (KeyData.Key.UnicodeChar == AcceptedChars[i])
 			{
-				SelectedChar = Key.UnicodeChar;
+				SelectedChar = KeyData.Key.UnicodeChar;
 				break;
 			}
 		}
@@ -589,6 +605,11 @@ UefiMain(
 	gBS->SetWatchdogTimer(0, 0, 0, NULL);
 
 	//
+	// Query the console input handle for the Simple Text Input Ex protocol
+	//
+	gBS->HandleProtocol(gST->ConsoleInHandle, &gEfiSimpleTextInputExProtocolGuid, (VOID **)&mTextInputEx);
+
+	//
 	// Locate, load, start and configure the driver
 	//
 	CONST EFI_STATUS DriverStatus = StartAndConfigureDriver(ImageHandle, SystemTable);
diff --git a/Application/Loader/Loader.inf b/Application/Loader/Loader.inf
index b400fe0..0acb5f2 100644
--- a/Application/Loader/Loader.inf
+++ b/Application/Loader/Loader.inf
@@ -55,6 +55,7 @@
   gEfiUsbIoProtocolGuid                            ## CONSUMES
   gEfiFirmwareVolume2ProtocolGuid                  ## CONSUMES
   gEfiSimpleTextInProtocolGuid                     ## CONSUMES
+  gEfiSimpleTextInputExProtocolGuid                ## CONSUMES
   gEfiSimpleTextOutProtocolGuid                    ## CONSUMES
   gEfiLegacyBiosProtocolGuid                       ## CONSUMES
 
diff --git a/EfiGuardDxe/EfiGuardDxe.c b/EfiGuardDxe/EfiGuardDxe.c
index 0af6bd9..1dd72ff 100644
--- a/EfiGuardDxe/EfiGuardDxe.c
+++ b/EfiGuardDxe/EfiGuardDxe.c
@@ -54,6 +54,7 @@ EFI_HANDLE gBootmgfwHandle = NULL;
 //
 // EFI runtime globals
 //
+EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL* gTextInputEx = NULL;
 EFI_EVENT gEfiExitBootServicesEvent = NULL;
 BOOLEAN gEfiAtRuntime = FALSE;
 EFI_EVENT gEfiVirtualNotifyEvent = NULL;
@@ -567,6 +568,11 @@ EfiGuardInitialize(
 	}
 
 	//
+	// Query the console input handle for the Simple Text Input Ex protocol
+	//
+	gBS->HandleProtocol(gST->ConsoleInHandle, &gEfiSimpleTextInputExProtocolGuid, (VOID **)&gTextInputEx);
+
+	//
 	// Install EfiGuard driver protocol
 	//
 	Status = gBS->InstallProtocolInterface(&gImageHandle,
diff --git a/EfiGuardDxe/EfiGuardDxe.h b/EfiGuardDxe/EfiGuardDxe.h
index 5d87513..e1e1d53 100644
--- a/EfiGuardDxe/EfiGuardDxe.h
+++ b/EfiGuardDxe/EfiGuardDxe.h
@@ -36,6 +36,11 @@ extern EFIGUARD_CONFIGURATION_DATA gDriverConfig;
 extern EFI_HANDLE gBootmgfwHandle;
 
 //
+// Simple Text Input Ex protocol pointer. May be NULL
+//
+extern EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL* gTextInputEx;
+
+//
 // TRUE if ExitBootServices() has been called
 //
 extern BOOLEAN gEfiAtRuntime;
diff --git a/EfiGuardDxe/EfiGuardDxe.inf b/EfiGuardDxe/EfiGuardDxe.inf
index 2ab9719..438fff8 100644
--- a/EfiGuardDxe/EfiGuardDxe.inf
+++ b/EfiGuardDxe/EfiGuardDxe.inf
@@ -53,6 +53,8 @@
   gEfiDevicePathUtilitiesProtocolGuid              ## CONSUMES
   gEfiLoadedImageProtocolGuid                      ## CONSUMES
   gEfiShellProtocolGuid                            ## SOMETIMES_CONSUMES
+  gEfiSimpleTextInProtocolGuid                     ## SOMETIMES_CONSUMES
+  gEfiSimpleTextInputExProtocolGuid                ## SOMETIMES_CONSUMES
 
 [Guids]
   gEfiGlobalVariableGuid                           ## SOMETIMES_PRODUCES
diff --git a/EfiGuardDxe/util.c b/EfiGuardDxe/util.c
index e7d88db..6ff8524 100644
--- a/EfiGuardDxe/util.c
+++ b/EfiGuardDxe/util.c
@@ -231,27 +231,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, 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