diff --git a/crates/boot/src/actions/chainload.rs b/crates/boot/src/actions/chainload.rs index fdebfc0..adcdca8 100644 --- a/crates/boot/src/actions/chainload.rs +++ b/crates/boot/src/actions/chainload.rs @@ -9,7 +9,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; @@ -94,19 +93,15 @@ pub fn chainload(context: Rc, 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(()) } diff --git a/crates/boot/src/actions/edera.rs b/crates/boot/src/actions/edera.rs index 147ef60..1acb94c 100644 --- a/crates/boot/src/actions/edera.rs +++ b/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, 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, 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 } diff --git a/crates/eficore/src/media_loader.rs b/crates/eficore/src/media_loader.rs index 245737f..f0bd793 100644 --- a/crates/eficore/src/media_loader.rs +++ b/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); + } + } +}