fix(eficore): unregister media loader on drop to prevent more opportunities for leaks (#50)

crates/boot/src/actions/chainload.rs

@@ -10,7 +10,6 @@ use eficore::loader::source::ImageSource;
 use eficore::loader::{ImageLoadRequest, ImageLoader};
 use eficore::media_loader::MediaLoaderHandle;
 use eficore::media_loader::constants::linux::LINUX_EFI_INITRD_MEDIA_GUID;
-use log::error;
 use uefi::CString16;
 use uefi::proto::loaded_image::LoadedImage;
 
@@ -99,19 +98,15 @@ pub fn chainload(context: Rc<SproutContext>, configuration: &ChainloadConfigurat
     // after the optional initrd has been unregistered.
     let result = uefi::boot::start_image(*image.handle());
 
-    // Unregister the initrd if it was registered.
-    if let Some(initrd_handle) = initrd_handle
-        && let Err(error) = initrd_handle.unregister()
-    {
-        error!("unable to unregister linux initrd: {}", error);
-    }
-
     // Assert there was no error starting the image.
     result.context("unable to start image")?;
 
     // Explicitly drop the options to clarify the lifetime.
     drop(options);
 
+    // Explicitly drop the initrd handle to clarify when it should be unregistered.
+    drop(initrd_handle);
+
     // Return control to sprout.
     Ok(())
 }

crates/boot/src/actions/edera.rs

@@ -15,7 +15,6 @@ use eficore::media_loader::{
         XEN_EFI_CONFIG_MEDIA_GUID, XEN_EFI_KERNEL_MEDIA_GUID, XEN_EFI_RAMDISK_MEDIA_GUID,
     },
 };
-use log::error;
 use uefi::Guid;
 
 /// Builds a configuration string for the Xen EFI stub using the specified `configuration`.
@@ -105,7 +104,7 @@ pub fn edera(context: Rc<SproutContext>, configuration: &EderaConfiguration) ->
     )
     .context("unable to register kernel media loader")?;
 
-    // Create a vector of media loaders to unregister on error.
+    // Create a vector of media loaders to drop them only after this function completes.
     let mut media_loaders = vec![config, kernel];
 
     // Register the initrd if it is provided.
@@ -127,12 +126,8 @@ pub fn edera(context: Rc<SproutContext>, configuration: &EderaConfiguration) ->
     )
     .context("unable to chainload to xen");
 
-    // Unregister the media loaders when an error happens.
-    for media_loader in media_loaders {
-        if let Err(error) = media_loader.unregister() {
-            error!("unable to unregister media loader: {}", error);
-        }
-    }
+    // Explicitly drop the media loaders to clarify when they should be unregistered.
+    drop(media_loaders);
 
     result
 }

crates/eficore/src/media_loader.rs

@@ -3,6 +3,7 @@ use alloc::vec::Vec;
 use anyhow::{Context, Result, bail};
 use core::ffi::c_void;
 use core::ptr;
+use log::error;
 use uefi::proto::device_path::DevicePath;
 use uefi::proto::device_path::build::DevicePathBuilder;
 use uefi::proto::device_path::build::media::Vendor;
@@ -255,7 +256,7 @@ impl MediaLoaderHandle {
 
     /// Unregisters a media loader from the UEFI stack.
     /// This will free the memory allocated by the passed data.
-    pub fn unregister(self) -> Result<()> {
+    fn unregister(&self) -> Result<()> {
         // SAFETY: We know that the media loader is registered if the handle is valid,
         // so we can safely uninstall it.
         // We should have allocated the pointers involved, so we can safely free them.
@@ -293,3 +294,14 @@ impl MediaLoaderHandle {
         Ok(())
     }
 }
+
+/// Implement drop for the handle to automatically unregister the media loader.
+impl Drop for MediaLoaderHandle {
+    fn drop(&mut self) {
+        // If unregister fails, print an error to the log.
+        // This may leak stuff, but the only other option is to panic.
+        if let Err(error) = self.unregister() {
+            error!("unable to unregister media loader: {}", error);
+        }
+    }
+}