From c9fc8249756305a2cf50a288167658df595d56df Mon Sep 17 00:00:00 2001
From: Nicholas Bishop <nbishop@nbishop.net>
Date: Wed, 31 Jul 2024 00:09:09 -0400
Subject: [PATCH 1/2] boot: Add freestanding version of raise_tpl

This comes with its own version of the `TplGuard` struct. This one has no
`BootServices` reference and no lifetime param.
---
 uefi/src/boot.rs | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/uefi/src/boot.rs b/uefi/src/boot.rs
index 8bc961a63..e3946b6c3 100644
--- a/uefi/src/boot.rs
+++ b/uefi/src/boot.rs
@@ -20,7 +20,7 @@ use {
 pub use uefi::table::boot::{
     AllocateType, LoadImageSource, OpenProtocolAttributes, OpenProtocolParams, SearchType,
 };
-pub use uefi_raw::table::boot::MemoryType;
+pub use uefi_raw::table::boot::{MemoryType, Tpl};
 
 /// Global image handle. This is only set by [`set_image_handle`], and it is
 /// only read by [`image_handle`].
@@ -58,6 +58,30 @@ fn boot_services_raw_panicking() -> NonNull<uefi_raw::table::boot::BootServices>
     NonNull::new(st.boot_services).expect("boot services are not active")
 }
 
+/// Raises a task's priority level and returns a [`TplGuard`].
+///
+/// The effect of calling `raise_tpl` with a `Tpl` that is below the current
+/// one (which, sadly, cannot be queried) is undefined by the UEFI spec,
+/// which also warns against remaining at high `Tpl`s for a long time.
+///
+/// This function returns an RAII guard that will automatically restore the
+/// original `Tpl` when dropped.
+///
+/// # Safety
+///
+/// Raising a task's priority level can affect other running tasks and
+/// critical processes run by UEFI. The highest priority level is the
+/// most dangerous, since it disables interrupts.
+#[must_use]
+pub unsafe fn raise_tpl(tpl: Tpl) -> TplGuard {
+    let bt = boot_services_raw_panicking();
+    let bt = unsafe { bt.as_ref() };
+
+    TplGuard {
+        old_tpl: (bt.raise_tpl)(tpl),
+    }
+}
+
 /// Allocates memory pages from the system.
 ///
 /// UEFI OS loaders should allocate memory of the type `LoaderData`.
@@ -451,3 +475,22 @@ impl<P: Protocol + ?Sized> ScopedProtocol<P> {
         self.interface.map(|mut p| unsafe { p.as_mut() })
     }
 }
+
+/// RAII guard for task priority level changes.
+///
+/// Will automatically restore the former task priority level when dropped.
+#[derive(Debug)]
+pub struct TplGuard {
+    old_tpl: Tpl,
+}
+
+impl Drop for TplGuard {
+    fn drop(&mut self) {
+        let bt = boot_services_raw_panicking();
+        let bt = unsafe { bt.as_ref() };
+
+        unsafe {
+            (bt.restore_tpl)(self.old_tpl);
+        }
+    }
+}

From 7b5351295b44b296d0f13e5838ff096a4decb735 Mon Sep 17 00:00:00 2001
From: Nicholas Bishop <nbishop@nbishop.net>
Date: Wed, 31 Jul 2024 13:50:26 -0400
Subject: [PATCH 2/2] test-runner: Add raise_tpl test

---
 uefi-test-runner/src/boot/misc.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/uefi-test-runner/src/boot/misc.rs b/uefi-test-runner/src/boot/misc.rs
index 2513b3f61..f78a8ea97 100644
--- a/uefi-test-runner/src/boot/misc.rs
+++ b/uefi-test-runner/src/boot/misc.rs
@@ -13,6 +13,7 @@ use uefi::{boot, guid, Event, Guid, Identify};
 
 pub fn test(st: &SystemTable<Boot>) {
     let bt = st.boot_services();
+    test_tpl();
     info!("Testing timer...");
     test_timer(bt);
     info!("Testing events...");
@@ -28,6 +29,12 @@ pub fn test(st: &SystemTable<Boot>) {
     test_install_configuration_table(st);
 }
 
+fn test_tpl() {
+    info!("Testing watchdog...");
+    // There's no way to query the TPL, so we can't assert that this does anything.
+    let _guard = unsafe { boot::raise_tpl(Tpl::NOTIFY) };
+}
+
 fn test_timer(bt: &BootServices) {
     let timer_event = unsafe { bt.create_event(EventType::TIMER, Tpl::APPLICATION, None, None) }
         .expect("Failed to create TIMER event");