Merge pull request #1525 from rust-osdev/bishop-dp-convert-free

Fix memory leaks in DevicePathFromText

Author: nicholasbishop

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

@@ -30,6 +30,17 @@ pub fn test() {
         )
         .expect("Failed to open DevicePathFromText protocol");
 
+        // Test round-trip conversion from path to text and back.
+        let device_path_string = device_path_to_text
+            .convert_device_path_to_text(&device_path, DisplayOnly(false), AllowShortcuts(false))
+            .unwrap();
+        assert_eq!(
+            *device_path_from_text
+                .convert_text_to_device_path(&device_path_string)
+                .unwrap(),
+            *device_path
+        );
+
         for path in device_path.node_iter() {
             info!(
                 "path: type={:?}, subtype={:?}, length={}",
@@ -47,7 +58,7 @@ pub fn test() {
             let convert = device_path_from_text
                 .convert_text_to_device_node(text)
                 .expect("Failed to convert text to device path");
-            assert_eq!(path, convert);
+            assert_eq!(*path, *convert);
         }
 
         // Get the `LoadedImageDevicePath`. Verify it start with the same nodes as

uefi/CHANGELOG.md

@@ -1,11 +1,18 @@
 # uefi - [Unreleased]
 
+## Added
+- Added `proto::device_path::PoolDevicePath` and
+  `proto::device_path::PoolDevicePathNode`.
+
 ## Changed
 - MSRV increased to 1.81.
 - `core::error::Error` impls are no longer gated by the `unstable` feature.
 - Fixed missing checks in the `TryFrom` conversion from `&DevicePathNode` to
   specific node types. The node type and subtype are now checked, and
   `NodeConversionError::DifferentType` is returned if they do not match.
+- **Breaking:** Fixed memory leaks in `DevicePathFromText` protocol. The methods
+  now return wrapper objects that free the device path / device path node on
+  drop.
 
 
 # uefi - 0.33.0 (2024-10-23)

uefi/src/mem/mod.rs

@@ -1,9 +1,36 @@
 //! Types, functions, traits, and other helpers to work with memory in UEFI
 //! libraries and applications.
 
+use crate::boot;
+use core::ptr::NonNull;
+
 pub mod memory_map;
+
 #[cfg(feature = "alloc")]
 pub(crate) mod util;
 
 #[cfg(feature = "alloc")]
 pub(crate) use util::*;
+
+/// Wrapper for memory allocated with UEFI's pool allocator. The memory is freed
+/// on drop.
+#[derive(Debug)]
+pub(crate) struct PoolAllocation(NonNull<u8>);
+
+impl PoolAllocation {
+    pub(crate) const fn new(ptr: NonNull<u8>) -> Self {
+        Self(ptr)
+    }
+
+    pub(crate) const fn as_ptr(&self) -> NonNull<u8> {
+        self.0
+    }
+}
+
+impl Drop for PoolAllocation {
+    fn drop(&mut self) {
+        // Ignore errors returned by `free_pool` since we can't propagate them
+        // from `drop`.
+        let _ = unsafe { boot::free_pool(self.0) };
+    }
+}

uefi/src/proto/device_path/text.rs

@@ -8,9 +8,10 @@
 // if there is insufficient memory. So we treat any NULL output as an
 // `OUT_OF_RESOURCES` error.
 
+use crate::mem::PoolAllocation;
 use crate::proto::device_path::{DevicePath, DevicePathNode};
 use crate::proto::unsafe_protocol;
-use crate::{boot, CStr16, Char16, Result, Status};
+use crate::{CStr16, Char16, Result, Status};
 use core::ops::Deref;
 use core::ptr::NonNull;
 use uefi_raw::protocol::device_path::{DevicePathFromTextProtocol, DevicePathToTextProtocol};
@@ -42,12 +43,12 @@ pub struct AllowShortcuts(pub bool);
 /// Wrapper for a string internally allocated from
 /// UEFI boot services memory.
 #[derive(Debug)]
-pub struct PoolString(NonNull<Char16>);
+pub struct PoolString(PoolAllocation);
 
 impl PoolString {
     fn new(text: *const Char16) -> Result<Self> {
         NonNull::new(text.cast_mut())
-            .map(Self)
+            .map(|p| Self(PoolAllocation::new(p.cast())))
             .ok_or(Status::OUT_OF_RESOURCES.into())
     }
 }
@@ -56,13 +57,31 @@ impl Deref for PoolString {
     type Target = CStr16;
 
     fn deref(&self) -> &Self::Target {
-        unsafe { CStr16::from_ptr(self.0.as_ptr()) }
+        unsafe { CStr16::from_ptr(self.0.as_ptr().as_ptr().cast()) }
     }
 }
 
-impl Drop for PoolString {
-    fn drop(&mut self) {
-        unsafe { boot::free_pool(self.0.cast()) }.expect("Failed to free pool [{addr:#?}]");
+/// Device path allocated from UEFI pool memory.
+#[derive(Debug)]
+pub struct PoolDevicePath(PoolAllocation);
+
+impl Deref for PoolDevicePath {
+    type Target = DevicePath;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { DevicePath::from_ffi_ptr(self.0.as_ptr().as_ptr().cast()) }
+    }
+}
+
+/// Device path node allocated from UEFI pool memory.
+#[derive(Debug)]
+pub struct PoolDevicePathNode(PoolAllocation);
+
+impl Deref for PoolDevicePathNode {
+    type Target = DevicePathNode;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { DevicePathNode::from_ffi_ptr(self.0.as_ptr().as_ptr().cast()) }
     }
 }
 
@@ -144,14 +163,12 @@ impl DevicePathFromText {
     pub fn convert_text_to_device_node(
         &self,
         text_device_node: &CStr16,
-    ) -> Result<&DevicePathNode> {
+    ) -> Result<PoolDevicePathNode> {
         unsafe {
             let ptr = (self.0.convert_text_to_device_node)(text_device_node.as_ptr().cast());
-            if ptr.is_null() {
-                Err(Status::OUT_OF_RESOURCES.into())
-            } else {
-                Ok(DevicePathNode::from_ffi_ptr(ptr.cast()))
-            }
+            NonNull::new(ptr.cast_mut())
+                .map(|p| PoolDevicePathNode(PoolAllocation::new(p.cast())))
+                .ok_or(Status::OUT_OF_RESOURCES.into())
         }
     }
 
@@ -165,14 +182,12 @@ impl DevicePathFromText {
     /// memory for the conversion.
     ///
     /// [`OUT_OF_RESOURCES`]: Status::OUT_OF_RESOURCES
-    pub fn convert_text_to_device_path(&self, text_device_path: &CStr16) -> Result<&DevicePath> {
+    pub fn convert_text_to_device_path(&self, text_device_path: &CStr16) -> Result<PoolDevicePath> {
         unsafe {
             let ptr = (self.0.convert_text_to_device_path)(text_device_path.as_ptr().cast());
-            if ptr.is_null() {
-                Err(Status::OUT_OF_RESOURCES.into())
-            } else {
-                Ok(DevicePath::from_ffi_ptr(ptr.cast()))
-            }
+            NonNull::new(ptr.cast_mut())
+                .map(|p| PoolDevicePath(PoolAllocation::new(p.cast())))
+                .ok_or(Status::OUT_OF_RESOURCES.into())
         }
     }
 }