new builder lol

Author: phip1611

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

@@ -1,8 +1,11 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 
-use uefi::boot::ScopedProtocol;
+use uefi::boot::{OpenProtocolAttributes, OpenProtocolParams, ScopedProtocol};
+use uefi::proto::device_path::DevicePath;
+use uefi::proto::device_path::text::{AllowShortcuts, DisplayOnly};
 use uefi::proto::shell::Shell;
 use uefi::{Error, Status, boot, cstr16};
+use uefi_raw::protocol::shell::ShellProtocol;
 
 /// Test `current_dir()` and `set_current_dir()`
 pub fn test_current_dir(shell: &ScopedProtocol<Shell>) {

uefi/src/proto/device_path/build.rs

@@ -2,8 +2,11 @@
 
 //! Utilities for creating new [`DevicePaths`].
 //!
-//! This module contains [`DevicePathBuilder`], as well as submodules
-//! containing types for building each type of device path node.
+//! This module provides builders to construct device paths, as well as
+//! submodules containing types for building each type of device path node.
+//!
+//! - [`DevicePathBuilder`] to construct device paths into pre-allocated buffers
+//! - [`OwnedDevicePathBuilder`] for construction on the Rust heap
 //!
 //! [`DevicePaths`]: DevicePath
 
@@ -15,7 +18,10 @@ use core::fmt::{self, Display, Formatter};
 use core::mem::MaybeUninit;
 
 #[cfg(feature = "alloc")]
-use alloc::vec::Vec;
+use {
+    alloc::{boxed::Box, vec::Vec},
+    core::mem,
+};
 
 /// A builder for [`DevicePaths`].
 ///
@@ -74,14 +80,14 @@ use alloc::vec::Vec;
 /// ```
 #[derive(Debug)]
 pub struct DevicePathBuilder<'a> {
-    storage: BuilderStorage<'a>,
+    storage: BuilderStorageRef<'a>,
 }
 
 impl<'a> DevicePathBuilder<'a> {
     /// Create a builder backed by a statically-sized buffer.
     pub const fn with_buf(buf: &'a mut [MaybeUninit<u8>]) -> Self {
         Self {
-            storage: BuilderStorage::Buf { buf, offset: 0 },
+            storage: BuilderStorageRef::Buf { buf, offset: 0 },
         }
     }
 
@@ -92,7 +98,7 @@ impl<'a> DevicePathBuilder<'a> {
     pub fn with_vec(v: &'a mut Vec<u8>) -> Self {
         v.clear();
         Self {
-            storage: BuilderStorage::Vec(v),
+            storage: BuilderStorageRef::Vec(v),
         }
     }
 
@@ -107,15 +113,15 @@ impl<'a> DevicePathBuilder<'a> {
         let node_size = usize::from(node.size_in_bytes()?);
 
         match &mut self.storage {
-            BuilderStorage::Buf { buf, offset } => {
+            BuilderStorageRef::Buf { buf, offset } => {
                 node.write_data(
                     buf.get_mut(*offset..*offset + node_size)
                         .ok_or(BuildError::BufferTooSmall)?,
                 );
                 *offset += node_size;
             }
             #[cfg(feature = "alloc")]
-            BuilderStorage::Vec(vec) => {
+            BuilderStorageRef::Vec(vec) => {
                 let old_size = vec.len();
                 vec.reserve(node_size);
                 let buf = &mut vec.spare_capacity_mut()[..node_size];
@@ -138,11 +144,11 @@ impl<'a> DevicePathBuilder<'a> {
         let this = self.push(&end::Entire)?;
 
         let data: &[u8] = match &this.storage {
-            BuilderStorage::Buf { buf, offset } => unsafe {
+            BuilderStorageRef::Buf { buf, offset } => unsafe {
                 maybe_uninit_slice_assume_init_ref(&buf[..*offset])
             },
             #[cfg(feature = "alloc")]
-            BuilderStorage::Vec(vec) => vec,
+            BuilderStorageRef::Vec(vec) => vec,
         };
 
         let ptr: *const () = data.as_ptr().cast();
@@ -152,7 +158,7 @@ impl<'a> DevicePathBuilder<'a> {
 
 /// Reference to the backup storage for [`DevicePathBuilder`]
 #[derive(Debug)]
-enum BuilderStorage<'a> {
+enum BuilderStorageRef<'a> {
     Buf {
         buf: &'a mut [MaybeUninit<u8>],
         offset: usize,
@@ -194,6 +200,58 @@ impl Display for BuildError {
 
 impl core::error::Error for BuildError {}
 
+/// Variant of [`DevicePathBuilder`] to construct a boxed [`DevicePath`].
+///
+/// Using this builder is equivalent to calling [`DevicePath::to_boxed`] on a
+/// device path constructed by the normal builder.
+#[derive(Debug)]
+#[cfg(feature = "alloc")]
+pub struct OwnedDevicePathBuilder {
+    vec: Vec<u8>,
+}
+
+#[cfg(feature = "alloc")]
+impl OwnedDevicePathBuilder {
+    /// Creates a new builder.
+    pub fn new() -> Self {
+        let vec = Vec::new();
+        Self { vec }
+    }
+
+    /// Add a node to the device path.
+    ///
+    /// An error will be returned if an [`END_ENTIRE`] node is passed to
+    /// this function, as that node will be added when [`Self::finalize`] is
+    /// called.
+    ///
+    /// [`END_ENTIRE`]: uefi::proto::device_path::DeviceSubType::END_ENTIRE
+    pub fn push(mut self, node: &dyn BuildNode) -> Result<Self, BuildError> {
+        let node_size = usize::from(node.size_in_bytes()?);
+
+        let old_size = self.vec.len();
+        self.vec.reserve(node_size);
+        let buf = &mut self.vec.spare_capacity_mut()[..node_size];
+        node.write_data(buf);
+        unsafe {
+            self.vec.set_len(old_size + node_size);
+        }
+        Ok(self)
+    }
+
+    /// Add an [`END_ENTIRE`] node and return the resulting [`DevicePath`].
+    ///
+    /// This method consumes the builder.
+    ///
+    /// [`END_ENTIRE`]: uefi::proto::device_path::DeviceSubType::END_ENTIRE
+    pub fn finalize(self) -> Result<Box<DevicePath>, BuildError> {
+        let this = self.push(&end::Entire)?;
+        let boxed = this.vec.into_boxed_slice();
+        // SAFETY: This is safe as a DevicePath has the same layout.
+        let dvp = unsafe { mem::transmute(boxed) };
+        Ok(dvp)
+    }
+}
+
 /// Trait for types that can be used to build a node via
 /// [`DevicePathBuilder::push`].
 ///