From 8b673b91c650deec5876f57da961a6727f77f6f0 Mon Sep 17 00:00:00 2001
From: Matthijs Lavrijsen <mattiwatti@gmail.com>
Date: Sun, 26 Mar 2023 20:05:35 +0200
Subject: [PATCH] RtlSleep: wait for a timer event instead of stalling

---
 EfiGuardDxe/EfiGuardDxe.c |  2 +-
 EfiGuardDxe/util.c        | 38 ++++++++++++++++++++++++++++----------
 EfiGuardDxe/util.h        | 11 ++++++++++-
 3 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/EfiGuardDxe/EfiGuardDxe.c b/EfiGuardDxe/EfiGuardDxe.c
index 81cfb00..8b17f7e 100644
--- a/EfiGuardDxe/EfiGuardDxe.c
+++ b/EfiGuardDxe/EfiGuardDxe.c
@@ -381,7 +381,7 @@ ExitBootServicesEvent(
 			PrintKernelPatchInfo();
 
 			// Give time for user to register their loss and allow for the grieving process to set in
-			RtlSleep(2000);
+			RtlStall(2000);
 
 			// Prompt user to ask what they want to do
 			Print(L"\r\nPress any key to continue anyway, or press ESC to reboot.\r\n");
diff --git a/EfiGuardDxe/util.c b/EfiGuardDxe/util.c
index 4c87f89..e435757 100644
--- a/EfiGuardDxe/util.c
+++ b/EfiGuardDxe/util.c
@@ -8,28 +8,46 @@
 #include <Library/UefiBootServicesTableLib.h>
 
 #ifndef ZYDIS_DISABLE_FORMATTER
-
 #include <Library/PrintLib.h>
 #include <Zycore/Format.h>
 
 STATIC ZydisFormatterFunc DefaultInstructionFormatter;
-
 #endif
 
-//
-// When debugging, we can choose between poor debugging facilities (VirtualBox) or poor performance and Windows compatibility (QEMU).
-// (I guess there is also the closed source thing with the horrible user interface that installs 50 drivers on the host (VMware))
-// This is a bandaid to make Print() calls readable ...for a while... when using VirtualBox or a live machine with no debugger
-//
+
 EFI_STATUS
-EFIAPI
 RtlSleep(
 	IN UINTN Milliseconds
 	)
 {
 	ASSERT(gBS != NULL);
-	ASSERT(gBS->Stall != NULL);
 
+	// Create a timer event, set its timeout, and wait for it
+	EFI_EVENT TimerEvent;
+	EFI_STATUS Status = gBS->CreateEvent(EVT_TIMER, 0, NULL, NULL, &TimerEvent);
+	if (EFI_ERROR(Status))
+		return RtlStall(Milliseconds); // Fall back to stalling CPU
+
+	gBS->SetTimer(TimerEvent,
+				TimerRelative,
+				EFI_TIMER_PERIOD_MILLISECONDS(Milliseconds));
+
+	UINTN Index;
+	Status = gBS->WaitForEvent(1, &TimerEvent, &Index);
+	if (EFI_ERROR(Status))
+		Status = RtlStall(Milliseconds);
+
+	gBS->CloseEvent(TimerEvent);
+	return Status;
+}
+
+EFI_STATUS
+EFIAPI
+RtlStall(
+	IN UINTN Milliseconds
+	)
+{
+	ASSERT(gBS != NULL);
 	return gBS->Stall(Milliseconds * 1000);
 }
 
@@ -189,7 +207,7 @@ WaitForKey(
 		if (Tpl == TPL_APPLICATION)
 			gBS->WaitForEvent(1, &gST->ConIn->WaitForKey, &Index); // Yep
 		else
-			RtlSleep(1); // Nope; burn CPU. // TODO: find a way to parallelize this to achieve GeForce FX 5800 temperatures
+			RtlStall(1); // Nope; burn CPU. // TODO: find a way to parallelize this to achieve GeForce FX 5800 temperatures
 
 		// 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);
diff --git a/EfiGuardDxe/util.h b/EfiGuardDxe/util.h
index 89c13e6..f67432e 100644
--- a/EfiGuardDxe/util.h
+++ b/EfiGuardDxe/util.h
@@ -10,7 +10,8 @@
 
 
 //
-// Stalls CPU for N milliseconds
+// Waits for a timer event for N milliseconds.
+// Requires current TPL to be TPL_APPLICATION.
 //
 EFI_STATUS
 EFIAPI
@@ -18,6 +19,14 @@ RtlSleep(
 	IN UINTN Milliseconds
 	);
 
+//
+// Stalls CPU for N milliseconds.
+//
+EFI_STATUS
+RtlStall(
+	IN UINTN Milliseconds
+	);
+
 // 
 // Prints info about a loaded image
 //