From 1d388e20630a27c57b2e054158cc0fbfde3b21fc Mon Sep 17 00:00:00 2001
From: Nicholas Bishop <nbishop@nbishop.net>
Date: Thu, 8 Aug 2024 10:09:55 -0400
Subject: [PATCH] runtime: Add freestanding set_virtual_address_map

---
 uefi/src/runtime.rs   | 45 +++++++++++++++++++++++++++++++++++++++++++
 uefi/src/table/mod.rs | 13 ++++++++++---
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/uefi/src/runtime.rs b/uefi/src/runtime.rs
index 2ee16c7de..db92de505 100644
--- a/uefi/src/runtime.rs
+++ b/uefi/src/runtime.rs
@@ -10,6 +10,7 @@ use crate::table::{self, Revision};
 use crate::{CStr16, Error, Result, Status, StatusExt};
 use core::mem;
 use core::ptr::{self, NonNull};
+use uefi_raw::table::boot::MemoryDescriptor;
 
 #[cfg(feature = "alloc")]
 use {
@@ -467,3 +468,47 @@ pub fn reset(reset_type: ResetType, status: Status, data: Option<&[u8]>) -> ! {
 
     unsafe { (rt.reset_system)(reset_type, status, size, data) }
 }
+
+/// Changes the runtime addressing mode of EFI firmware from physical to
+/// virtual. It is up to the caller to translate the old system table address
+/// to a new virtual address and provide it for this function.
+///
+/// If successful, this function will call [`set_system_table`] with
+/// `new_system_table_virtual_addr`.
+///
+/// [`set_system_table`]: table::set_system_table
+///
+/// # Safety
+///
+/// The caller must ensure the memory map is valid.
+///
+/// # Errors
+///
+/// * [`Status::UNSUPPORTED`]: either boot services haven't been exited, the
+///   firmware's addressing mode is already virtual, or the firmware does not
+///   support this operation.
+/// * [`Status::NO_MAPPING`]: `map` is missing a required range.
+/// * [`Status::NOT_FOUND`]: `map` contains an address that is not in the
+///   current memory map.
+pub unsafe fn set_virtual_address_map(
+    map: &mut [MemoryDescriptor],
+    new_system_table_virtual_addr: *const uefi_raw::table::system::SystemTable,
+) -> Result {
+    let rt = runtime_services_raw_panicking();
+    let rt = unsafe { rt.as_ref() };
+
+    // Unsafe Code Guidelines guarantees that there is no padding in an array or a slice
+    // between its elements if the element type is `repr(C)`, which is our case.
+    //
+    // See https://rust-lang.github.io/unsafe-code-guidelines/layout/arrays-and-slices.html
+    let map_size = core::mem::size_of_val(map);
+    let entry_size = core::mem::size_of::<MemoryDescriptor>();
+    let entry_version = MemoryDescriptor::VERSION;
+    let map_ptr = map.as_mut_ptr();
+    (rt.set_virtual_address_map)(map_size, entry_size, entry_version, map_ptr).to_result()?;
+
+    // Update the global system table pointer.
+    table::set_system_table(new_system_table_virtual_addr);
+
+    Ok(())
+}
diff --git a/uefi/src/table/mod.rs b/uefi/src/table/mod.rs
index cd837bba8..e8f54b396 100644
--- a/uefi/src/table/mod.rs
+++ b/uefi/src/table/mod.rs
@@ -33,14 +33,21 @@ pub(crate) fn system_table_raw_panicking() -> NonNull<uefi_raw::table::system::S
 /// Update the global system table pointer.
 ///
 /// This is called automatically in the `main` entry point as part of
-/// [`uefi::entry`]. It should not be called at any other point in time, unless
-/// the executable does not use [`uefi::entry`], in which case it should be
-/// called once before calling any other API in this crate.
+/// [`uefi::entry`].
+///
+/// It is also called by [`set_virtual_address_map`] to transition from a
+/// physical address to a virtual address.
+///
+/// This function should not be called at any other point in time, unless the
+/// executable does not use [`uefi::entry`], in which case it should be called
+/// once before calling any other API in this crate.
 ///
 /// # Safety
 ///
 /// This function should only be called as described above, and the
 /// `ptr` must be a valid [`SystemTable`].
+///
+/// [`set_virtual_address_map`]: uefi::runtime::set_virtual_address_map
 pub unsafe fn set_system_table(ptr: *const uefi_raw::table::system::SystemTable) {
     SYSTEM_TABLE.store(ptr.cast_mut(), Ordering::Release);
 }