Rework memory_attributes_table to remove the MEMORY_ATTRIBUTES_TABLE static

patina_dxe_core/src/allocator.rs

@@ -790,7 +790,7 @@ pub fn terminate_memory_map(map_key: usize) -> Result<(), EfiError> {
 
 pub fn install_memory_type_info_table(system_table: &mut EfiSystemTable) -> Result<(), EfiError> {
     let table_ptr = NonNull::from(GCD.memory_type_info_table()).cast::<c_void>().as_ptr();
-    config_tables::core_install_configuration_table(guids::MEMORY_TYPE_INFORMATION, table_ptr, system_table)
+    config_tables::core_install_configuration_table(guids::MEMORY_TYPE_INFORMATION, table_ptr, system_table).map(|_| ())
 }
 
 fn process_hob_allocations(hob_list: &HobList) {

patina_dxe_core/src/config_tables.rs

@@ -10,7 +10,11 @@ pub(crate) mod debug_image_info_table;
 pub(crate) mod memory_attributes_table;
 
 use alloc::{boxed::Box, vec};
-use core::{ffi::c_void, ptr::slice_from_raw_parts_mut};
+use core::{
+    ffi::c_void,
+    ptr::{NonNull, slice_from_raw_parts_mut},
+    slice::from_raw_parts,
+};
 use patina::error::EfiError;
 use r_efi::efi;
 
@@ -36,15 +40,17 @@ extern "efiapi" fn install_configuration_table(table_guid: *mut efi::Guid, table
 
     match core_install_configuration_table(table_guid, table, st) {
         Err(err) => err.into(),
-        Ok(()) => efi::Status::SUCCESS,
+        Ok(_) => efi::Status::SUCCESS,
     }
 }
 
+/// Install a configuration table in the system table, replacing any existing table with the same GUID.
+/// If a table is replaced or deleted, a pointer to the old table is returned.
 pub fn core_install_configuration_table(
     vendor_guid: efi::Guid,
     vendor_table: *mut c_void,
     efi_system_table: &mut EfiSystemTable,
-) -> Result<(), EfiError> {
+) -> Result<Option<NonNull<c_void>>, EfiError> {
     let system_table = efi_system_table.as_mut();
     //if a table is already present, reconstruct it from the pointer and length in the st.
     let old_cfg_table = if system_table.configuration_table.is_null() {
@@ -60,6 +66,7 @@ pub fn core_install_configuration_table(
         Some(ct_slice_box)
     };
 
+    let mut old_vendor_table_ptr = None;
     // construct the new table contents as a vector.
     let new_table = match old_cfg_table {
         Some(cfg_table) => {
@@ -70,15 +77,17 @@ pub fn core_install_configuration_table(
                 // vendor_table is not null; we are adding or modifying an entry.
                 if let Some(entry) = existing_entry {
                     //entry exists, modify it.
+                    old_vendor_table_ptr = NonNull::new(entry.vendor_table);
                     entry.vendor_table = vendor_table;
                 } else {
                     //entry doesn't exist, add it.
                     current_table.push(efi::ConfigurationTable { vendor_guid, vendor_table });
                 }
             } else {
                 //vendor_table is none; we are deleting an entry.
-                if let Some(_entry) = existing_entry {
+                if let Some(entry) = existing_entry {
                     //entry exists, we can delete it
+                    old_vendor_table_ptr = NonNull::new(entry.vendor_table);
                     current_table.retain(|x| x.vendor_guid != vendor_guid);
                 } else {
                     //entry does not exist, we can't delete it. We have to put the original box back
@@ -121,9 +130,87 @@ pub fn core_install_configuration_table(
     //signal the table guid as an event group
     EVENT_DB.signal_group(vendor_guid);
 
-    Ok(())
+    Ok(old_vendor_table_ptr)
+}
+
+/// Returns the pointer to a configuration table for the specified guid, if it exists.
+pub fn get_configuration_table(table_guid: &efi::Guid) -> Option<NonNull<c_void>> {
+    let st_guard = SYSTEM_TABLE.lock();
+    let st = st_guard.as_ref()?;
+
+    let system_table = st.as_ref();
+    if system_table.configuration_table.is_null() || system_table.number_of_table_entries == 0 {
+        return None;
+    }
+
+    // Safety: system table exists, and configuration is non-null, and number_of_table_entries is non-zero.
+    let ct_slice = unsafe { from_raw_parts(system_table.configuration_table, system_table.number_of_table_entries) };
+
+    for entry in ct_slice {
+        if entry.vendor_guid == *table_guid {
+            return NonNull::new(entry.vendor_table);
+        }
+    }
+    None
 }
 
 pub fn init_config_tables_support(bs: &mut efi::BootServices) {
     bs.install_configuration_table = install_configuration_table;
 }
+
+#[cfg(test)]
+mod tests {
+    use patina::base::guid;
+
+    use crate::{systemtables::init_system_table, test_support};
+
+    use super::*;
+
+    fn with_locked_state<F: Fn() + std::panic::RefUnwindSafe>(f: F) {
+        test_support::with_global_lock(|| {
+            unsafe {
+                test_support::init_test_gcd(None);
+                test_support::reset_allocators();
+                init_system_table();
+            }
+            f();
+        })
+        .unwrap();
+    }
+
+    #[test]
+    fn install_configuration_table_should_install_table() {
+        with_locked_state(|| {
+            let guid: efi::Guid = guid::Guid::from_string("78926ab0-af16-49e4-8e05-115aafbca1df").to_efi_guid();
+            let table = 0x12345678u32 as *mut c_void;
+
+            assert!(get_configuration_table(&guid).is_none());
+
+            assert_eq!(install_configuration_table(&guid as *const _ as *mut _, table), efi::Status::SUCCESS);
+            assert_eq!(get_configuration_table(&guid).unwrap().as_ptr(), table);
+        });
+    }
+
+    #[test]
+    fn delete_config_table_should_return_ptr() {
+        with_locked_state(|| {
+            let guid: efi::Guid = guid::Guid::from_string("78926ab0-af16-49e4-8e05-115aafbca1df").to_efi_guid();
+            let table = 0x12345678u32 as *mut c_void;
+
+            assert_eq!(install_configuration_table(&guid as *const _ as *mut _, table), efi::Status::SUCCESS);
+
+            assert_eq!(get_configuration_table(&guid).unwrap().as_ptr(), table);
+
+            assert_eq!(
+                core_install_configuration_table(
+                    guid,
+                    core::ptr::null_mut(),
+                    &mut *SYSTEM_TABLE.lock().as_mut().unwrap()
+                ),
+                Ok(Some(NonNull::new(table).unwrap()))
+            );
+
+            assert!(get_configuration_table(&guid).is_none());
+        });
+    }
+}

patina_dxe_core/src/config_tables/memory_attributes_table.rs

@@ -9,33 +9,30 @@
 extern crate alloc;
 use alloc::vec::Vec;
 
-use core::{
-    ffi::c_void,
-    fmt::Debug,
-    mem::size_of,
-    slice,
-    sync::atomic::{AtomicBool, AtomicPtr, Ordering},
-};
+#[cfg(not(test))]
+use spin::Once;
+
+use core::{ffi::c_void, fmt::Debug, mem::size_of, slice};
 
 use crate::{
     allocator::{MemoryDescriptorSlice, core_allocate_pool, core_free_pool, get_memory_map_descriptors},
-    config_tables::core_install_configuration_table,
+    config_tables::{core_install_configuration_table, get_configuration_table},
     events::EVENT_DB,
     gcd::MemoryProtectionPolicy,
     systemtables,
 };
 use r_efi::efi;
 
-// We cache the MAT here because we need to free it in whenever we get a new runtime code/data allocation
-static MEMORY_ATTRIBUTES_TABLE: AtomicPtr<c_void> = AtomicPtr::new(core::ptr::null_mut());
-
 // create a wrapper struct so that we can create an install method on it. That way, we can have the install function
 // be a no-op until after ReadyToBoot
 pub struct MemoryAttributesTable(*mut efi::MemoryAttributesTable);
 
 // this is a flag to indicate that we have passed ReadyToBoot and can install the MAT on the next runtime memory
-// allocation/deallocation
-static POST_RTB: AtomicBool = AtomicBool::new(false);
+// allocation/deallocation.
+#[cfg(not(test))]
+static POST_RTB: Once<()> = Once::new();
+#[cfg(test)]
+static POST_RTB: crate::test_support::TestOnce<()> = crate::test_support::TestOnce::new();
 
 impl MemoryAttributesTable {
     ///
@@ -55,7 +52,7 @@ impl MemoryAttributesTable {
     /// ```
     ///
     pub fn install() {
-        if POST_RTB.load(Ordering::Relaxed) {
+        if POST_RTB.is_completed() {
             core_install_memory_attributes_table()
         }
     }
@@ -98,50 +95,51 @@ pub fn init_memory_attributes_table_support() {
 // After this point, subsequent runtime memory allocations/deallocations will create new MAT tables
 extern "efiapi" fn core_install_memory_attributes_table_event_wrapper(event: efi::Event, _context: *mut c_void) {
     core_install_memory_attributes_table();
-    // now we want to capture any future runtime memory changes, so we will mark that ReadyToBoot has occurred
-    // and the install callback will be invoked on the next runtime memory allocation
-    POST_RTB.store(true, Ordering::Relaxed);
 
     if let Err(status) = EVENT_DB.close_event(event) {
         log::error!("Failed to close MAT ready to boot event with status {status:#X?}. This should be okay.");
     }
 }
 
 pub fn core_install_memory_attributes_table() {
-    let mut st_guard = systemtables::SYSTEM_TABLE.lock();
-    let st = st_guard.as_mut().expect("System table support not initialized");
-
-    let current_ptr = MEMORY_ATTRIBUTES_TABLE.load(Ordering::Relaxed);
-    if current_ptr.is_null() {
-        // we need to install an empty configuration table the first time here, because core_install_configuration_table
-        // may allocate runtime memory. Because it actually gets installed we need to allocate one here, it will be
-        // freed below when we install the real MAT. If we don't allocate this on the heap, we may have undefined
-        // behavior with a stack pointer that goes out of scope
-        match core_allocate_pool(efi::BOOT_SERVICES_DATA, size_of::<efi::MemoryAttributesTable>()) {
-            Ok(empty_ptr) => {
-                if let Some(empty_mat) = unsafe { (empty_ptr as *mut efi::MemoryAttributesTable).as_mut() } {
-                    *empty_mat = efi::MemoryAttributesTable {
-                        version: 0,
-                        number_of_entries: 0,
-                        descriptor_size: 0,
-                        reserved: 0,
-                        entry: [],
-                    };
-                    MEMORY_ATTRIBUTES_TABLE.store(empty_ptr, Ordering::Relaxed);
-
-                    if let Err(status) =
-                        core_install_configuration_table(efi::MEMORY_ATTRIBUTES_TABLE_GUID, empty_ptr, st)
-                    {
-                        log::error!("Failed to create a null MAT table with status {status:#X?}, cannot create MAT");
-                        return;
+    if !POST_RTB.is_completed() {
+        if get_configuration_table(&efi::MEMORY_ATTRIBUTES_TABLE_GUID).is_none() {
+            // we need to install an empty configuration table the first time here, because core_install_configuration_table
+            // may allocate runtime memory. Because it actually gets installed we need to allocate one here, it will be
+            // freed below when we install the real MAT. If we don't allocate this on the heap, we may have undefined
+            // behavior with a stack pointer that goes out of scope
+            match core_allocate_pool(efi::BOOT_SERVICES_DATA, size_of::<efi::MemoryAttributesTable>()) {
+                Ok(empty_ptr) => {
+                    if let Some(empty_mat) = unsafe { (empty_ptr as *mut efi::MemoryAttributesTable).as_mut() } {
+                        *empty_mat = efi::MemoryAttributesTable {
+                            version: 0,
+                            number_of_entries: 0,
+                            descriptor_size: 0,
+                            reserved: 0,
+                            entry: [],
+                        };
+                        let mut st_guard = systemtables::SYSTEM_TABLE.lock();
+                        let st = st_guard.as_mut().expect("System table support not initialized");
+
+                        if let Err(status) =
+                            core_install_configuration_table(efi::MEMORY_ATTRIBUTES_TABLE_GUID, empty_ptr, st)
+                        {
+                            log::error!(
+                                "Failed to create a null MAT table with status {status:#X?}, cannot create MAT"
+                            );
+                            return;
+                        }
                     }
                 }
-            }
-            Err(err) => {
-                log::error!("Failed to allocate memory for a null MAT! Status {err:#X?}");
-                return;
+                Err(err) => {
+                    log::error!("Failed to allocate memory for a null MAT! Status {err:#X?}");
+                    return;
+                }
             }
         }
+        // now we want to capture any future runtime memory changes, so we will mark that ReadyToBoot has occurred
+        // and the install callback will be invoked on the next runtime memory allocation
+        POST_RTB.call_once(|| {});
     }
 
     // get the GCD memory map descriptors and filter out the non-runtime sections
@@ -212,6 +210,9 @@ pub fn core_install_memory_attributes_table() {
                     mat_desc_list.len() * size_of::<efi::MemoryDescriptor>(),
                 );
 
+                let mut st_guard = systemtables::SYSTEM_TABLE.lock();
+                let st = st_guard.as_mut().expect("System table support not initialized");
+
                 match core_install_configuration_table(efi::MEMORY_ATTRIBUTES_TABLE_GUID, void_ptr, st) {
                     Err(status) => {
                         log::error!("Failed to install MAT table! Status {status:#X?}");
@@ -220,17 +221,13 @@ pub fn core_install_memory_attributes_table() {
                         }
                         return;
                     }
-
-                    Ok(_) => {
+                    Ok(Some(current_ptr)) => {
                         // free the old MAT table if we have one
-                        let current_ptr = MEMORY_ATTRIBUTES_TABLE.load(Ordering::Relaxed);
-                        if !current_ptr.is_null()
-                            && let Err(err) = core_free_pool(current_ptr)
-                        {
+                        if let Err(err) = core_free_pool(current_ptr.as_ptr()) {
                             log::error!("Error freeing previous MAT pointer: {err:#X?}");
                         }
-                        MEMORY_ATTRIBUTES_TABLE.store(void_ptr, Ordering::Relaxed);
                     }
+                    Ok(None) => (),
                 }
             }
 
@@ -256,8 +253,7 @@ mod tests {
 
     fn with_locked_state<F: Fn() + std::panic::RefUnwindSafe>(f: F) {
         test_support::with_global_lock(|| {
-            POST_RTB.store(false, Ordering::Relaxed);
-            MEMORY_ATTRIBUTES_TABLE.store(core::ptr::null_mut(), Ordering::Relaxed);
+            POST_RTB.reset();
 
             unsafe {
                 test_support::init_test_gcd(None);
@@ -266,8 +262,7 @@ mod tests {
             }
             f();
 
-            POST_RTB.store(false, Ordering::Relaxed);
-            MEMORY_ATTRIBUTES_TABLE.store(core::ptr::null_mut(), Ordering::Relaxed);
+            POST_RTB.reset();
         })
         .unwrap();
     }
@@ -284,6 +279,7 @@ mod tests {
     fn test_memory_attributes_table_generation() {
         with_locked_state(|| {
             // Create a vector to store the allocated pages
+
             let mut allocated_pages = Vec::new();
             let mut entry_count = 0;
 
@@ -325,25 +321,24 @@ mod tests {
             }
 
             // before we create the MAT, we expect MEMORY_ATTRIBUTES_TABLE to be None
-            assert!(MEMORY_ATTRIBUTES_TABLE.load(Ordering::Relaxed).is_null());
+            assert!(get_configuration_table(&efi::MEMORY_ATTRIBUTES_TABLE_GUID).is_none());
 
             // Create a dummy event
             let dummy_event: efi::Event = core::ptr::null_mut();
 
             // Ensure POST_RTB is false before the event
-            assert!(!POST_RTB.load(Ordering::Relaxed));
+            assert!(!POST_RTB.is_completed());
 
             // Call the event wrapper
             core_install_memory_attributes_table_event_wrapper(dummy_event, core::ptr::null_mut());
 
             // Check if POST_RTB is set after the event
-            assert!(POST_RTB.load(Ordering::Relaxed));
+            assert!(POST_RTB.is_completed());
 
             // Check if MEMORY_ATTRIBUTES_TABLE is set after installation
-            assert!(!MEMORY_ATTRIBUTES_TABLE.load(Ordering::Relaxed).is_null());
-            let mat_ptr = MEMORY_ATTRIBUTES_TABLE.load(Ordering::Relaxed);
+            let mat_ptr = get_configuration_table(&efi::MEMORY_ATTRIBUTES_TABLE_GUID).unwrap();
             unsafe {
-                let mat = &*(mat_ptr as *const _ as *const efi::MemoryAttributesTable);
+                let mat = &*(mat_ptr.as_ptr() as *const efi::MemoryAttributesTable);
 
                 assert_eq!(mat.version, efi::MEMORY_ATTRIBUTES_TABLE_VERSION);
                 // we have one extra entry here because init_system_table allocates runtime pages