relicensing: Rewrite allocator, configuration table, and image unload PRs

uefi-test-runner/src/boot/misc.rs

@@ -184,29 +184,31 @@ fn test_uninstall_protocol_interface() {
 
 fn test_install_configuration_table() {
     // Get the current number of entries.
-    let count = system::with_config_table(|t| t.len());
+    let initial_table_count = system::with_config_table(|t| t.len());
 
     // Create the entry data.
-    let config = boot::allocate_pool(MemoryType::RUNTIME_SERVICES_DATA, 1)
-        .unwrap()
-        .as_ptr();
-    unsafe { config.write(42) };
+    let config: NonNull<u8> = boot::allocate_pool(MemoryType::RUNTIME_SERVICES_DATA, 1).unwrap();
+    unsafe { config.write(123u8) };
 
     // Install the table.
-    const ID: Guid = guid!("4bec53c4-5fc1-48a1-ab12-df214907d29f");
+    const TABLE_GUID: Guid = guid!("4bec53c4-5fc1-48a1-ab12-df214907d29f");
     unsafe {
-        boot::install_configuration_table(&ID, config.cast()).unwrap();
+        boot::install_configuration_table(&TABLE_GUID, config.as_ptr().cast()).unwrap();
     }
 
     // Verify the installation.
-    assert_eq!(count + 1, system::with_config_table(|t| t.len()));
+    assert_eq!(
+        initial_table_count + 1,
+        system::with_config_table(|t| t.len())
+    );
     system::with_config_table(|t| {
-        let config_entry = t.iter().find(|ct| ct.guid == ID).unwrap();
-        assert_eq!(unsafe { *(config_entry.address as *const u8) }, 42);
+        let config_entry = t.iter().find(|ct| ct.guid == TABLE_GUID).unwrap();
+        assert_eq!(unsafe { *config_entry.address.cast::<u8>() }, 123);
     });
 
-    // Uninstall the table.
+    // Uninstall the table and free the memory.
     unsafe {
-        boot::install_configuration_table(&ID, ptr::null()).unwrap();
+        boot::install_configuration_table(&TABLE_GUID, ptr::null()).unwrap();
+        boot::free_pool(config).unwrap();
     }
 }

uefi-test-runner/src/proto/media.rs

@@ -15,6 +15,9 @@ use uefi::proto::media::fs::SimpleFileSystem;
 use uefi::proto::media::partition::{MbrOsType, PartitionInfo};
 use uefi::runtime::{Daylight, Time, TimeParams};
 
+#[repr(align(8))]
+struct AlignedBuf([u8; 256]);
+
 /// Test directory entry iteration.
 fn test_existing_dir(directory: &mut Directory) {
     info!("Testing existing directory");
@@ -31,11 +34,9 @@ fn test_existing_dir(directory: &mut Directory) {
     let dir = RefCell::new(dir);
 
     assert_eq!(FileInfo::alignment(), 8);
-    #[repr(align(8))]
-    struct Buf([u8; 200]);
 
     // Backing memory to read the file info data into.
-    let mut stack_buf = Buf([0; 200]);
+    let mut stack_buf = AlignedBuf([0; 256]);
 
     // The file names that the test read from the directory.
     let entry_names = RefCell::new(vec![]);

uefi/src/allocator.rs

@@ -7,13 +7,12 @@
 //! services are not active, `alloc` will return a null pointer, and `dealloc`
 //! will panic.
 
-use core::alloc::{GlobalAlloc, Layout};
-use core::ptr::{self, NonNull};
-use core::sync::atomic::{AtomicU32, Ordering};
-
 use crate::boot;
 use crate::mem::memory_map::MemoryType;
 use crate::proto::loaded_image::LoadedImage;
+use core::alloc::{GlobalAlloc, Layout};
+use core::ptr::{self, NonNull};
+use core::sync::atomic::{AtomicU32, Ordering};
 
 /// Get the memory type to use for allocation.
 ///
@@ -48,9 +47,11 @@ fn get_memory_type() -> MemoryType {
 pub struct Allocator;
 
 unsafe impl GlobalAlloc for Allocator {
-    /// Allocate memory using [`boot::allocate_pool`]. The allocation is
-    /// of type [`MemoryType::LOADER_DATA`] for UEFI applications, [`MemoryType::BOOT_SERVICES_DATA`]
-    /// for UEFI boot drivers and [`MemoryType::RUNTIME_SERVICES_DATA`] for UEFI runtime drivers.
+    /// Allocate memory using [`boot::allocate_pool`]. The allocation's [memory
+    /// type] matches the current image's [data type].
+    ///
+    /// [memory type]: MemoryType
+    /// [data type]: LoadedImage::data_type
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         if !boot::are_boot_services_active() {
             return ptr::null_mut();

uefi/src/proto/loaded_image.rs

@@ -134,25 +134,21 @@ impl LoadedImage {
         self.0.image_size = size;
     }
 
-    /// Set the callback handler to unload the image.
-    ///
-    /// Drivers that wish to support unloading have to register their unload handler
-    /// using this protocol. It is responsible for cleaning up any resources the
-    /// image is using before returning. Unloading a driver is done with
-    /// [`boot::unload_image`].
+    /// Registers a cleanup function that is called when [`boot::unload_image`]
+    /// is called.
     ///
     /// # Safety
     ///
-    /// Only the driver that this [`LoadedImage`] is attached to should register an
-    /// unload handler.
+    /// The registered function must reside in memory that is not freed until
+    /// after the image is unloaded.
     ///
     /// [`boot::unload_image`]: crate::boot::unload_image
     pub unsafe fn set_unload(
         &mut self,
         unload: extern "efiapi" fn(image_handle: Handle) -> Status,
     ) {
-        type RawFn = unsafe extern "efiapi" fn(image_handle: uefi_raw::Handle) -> uefi_raw::Status;
-        let unload: RawFn = mem::transmute(unload);
+        let unload: unsafe extern "efiapi" fn(image_handle: uefi_raw::Handle) -> uefi_raw::Status =
+            mem::transmute(unload);
         self.0.unload = Some(unload);
     }
 
@@ -180,23 +176,13 @@ impl LoadedImage {
         (self.0.image_base, self.0.image_size)
     }
 
-    /// Get the memory type of the image's code sections.
-    ///
-    /// Normally the returned value is one of:
-    ///  - `MemoryType::LOADER_CODE` for UEFI applications
-    ///  - `MemoryType::BOOT_SERVICES_CODE` for UEFI boot drivers
-    ///  - `MemoryType::RUNTIME_SERVICES_CODE` for UEFI runtime drivers
+    /// Returns the memory type that the image's code sections were loaded as.
     #[must_use]
     pub const fn code_type(&self) -> MemoryType {
         self.0.image_code_type
     }
 
-    /// Get the memory type of the image's data sections.
-    ///
-    /// Normally the returned value is one of:
-    ///  - `MemoryType::LOADER_DATA` for UEFI applications
-    ///  - `MemoryType::BOOT_SERVICES_DATA` for UEFI boot drivers
-    ///  - `MemoryType::RUNTIME_SERVICES_DATA` for UEFI runtime drivers
+    /// Returns the memory type that the image's data sections were loaded as.
     #[must_use]
     pub const fn data_type(&self) -> MemoryType {
         self.0.image_data_type