From 798797e1e58cac1326a283bb042ada19dc283fd3 Mon Sep 17 00:00:00 2001
From: andre-braga <andrebraga@google.com>
Date: Tue, 11 Jun 2024 20:17:36 +0000
Subject: [PATCH] uefi: Add TryFrom u8 slice to DevicePath

Adds the capability to convert a byte slice to a DevicePath. This allows for
users to easily get a DevicePath from the data returned by UEFI variables, etc.
---
 uefi/CHANGELOG.md                 |  2 +-
 uefi/src/proto/device_path/mod.rs | 62 +++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md
index b33952820..fe4c90515 100644
--- a/uefi/CHANGELOG.md
+++ b/uefi/CHANGELOG.md
@@ -11,7 +11,7 @@
 - Added `table::{set_system_table, system_table_boot, system_table_runtime}`.
   This provides an initial API for global tables that do not require passing
   around a reference.
-- Added `TryFrom<&[u8]>` for `DevicePathHeader` and `DevicePathNode`.
+- Added `TryFrom<&[u8]>` for `DevicePathHeader`, `DevicePathNode` and `DevicePath`.
 - Added `ByteConversionError`.
 
 ## Changed
diff --git a/uefi/src/proto/device_path/mod.rs b/uefi/src/proto/device_path/mod.rs
index 087440bb5..b61268e08 100644
--- a/uefi/src/proto/device_path/mod.rs
+++ b/uefi/src/proto/device_path/mod.rs
@@ -397,6 +397,35 @@ impl DevicePath {
         total_size_in_bytes
     }
 
+    /// Calculate the size in bytes of the entire `DevicePath` starting
+    /// at `bytes`. This adds up each node's length, including the
+    /// end-entire node.
+    ///
+    /// # Errors
+    ///
+    /// The [`ByteConversionError::InvalidLength`] error will be returned
+    /// when the length of the given bytes slice cannot contain the full
+    /// [`DevicePath`] represented by the slice.
+    fn size_in_bytes_from_slice(mut bytes: &[u8]) -> Result<usize, ByteConversionError> {
+        let max_size_in_bytes = bytes.len();
+        let mut total_size_in_bytes: usize = 0;
+        loop {
+            let node = <&DevicePathNode>::try_from(bytes)?;
+            let node_size_in_bytes = usize::from(node.length());
+            total_size_in_bytes += node_size_in_bytes;
+            // Length of last processed node extends past the bytes slice.
+            if total_size_in_bytes > max_size_in_bytes {
+                return Err(ByteConversionError::InvalidLength);
+            }
+            if node.is_end_entire() {
+                break;
+            }
+            bytes = &bytes[node_size_in_bytes..];
+        }
+
+        Ok(total_size_in_bytes)
+    }
+
     /// Create a [`DevicePath`] reference from an opaque pointer.
     ///
     /// # Safety
@@ -488,6 +517,15 @@ impl PartialEq for DevicePath {
     }
 }
 
+impl<'a> TryFrom<&[u8]> for &'a DevicePath {
+    type Error = ByteConversionError;
+
+    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
+        let len = DevicePath::size_in_bytes_from_slice(bytes)?;
+        unsafe { Ok(&*ptr_meta::from_raw_parts(bytes.as_ptr().cast(), len)) }
+    }
+}
+
 #[cfg(feature = "alloc")]
 impl ToOwned for DevicePath {
     type Owned = Box<DevicePath>;
@@ -1017,4 +1055,28 @@ mod tests {
         raw_data[2] += 1;
         assert!(<&DevicePathNode>::try_from(raw_data.as_slice()).is_err());
     }
+
+    #[test]
+    fn test_device_path_nodes_from_bytes() {
+        let raw_data = create_raw_device_path();
+        let dp = <&DevicePath>::try_from(raw_data.as_slice()).unwrap();
+
+        // Check that the size is the sum of the nodes' lengths.
+        assert_eq!(mem::size_of_val(dp), 6 + 8 + 4 + 6 + 8 + 4);
+
+        // Check the list's node iter.
+        let nodes: Vec<_> = dp.node_iter().collect();
+        check_node(nodes[0], 0xa0, 0xb0, &[10, 11]);
+        check_node(nodes[1], 0xa1, 0xb1, &[20, 21, 22, 23]);
+        check_node(
+            nodes[2],
+            DeviceType::END.0,
+            DeviceSubType::END_INSTANCE.0,
+            &[],
+        );
+        check_node(nodes[3], 0xa2, 0xb2, &[30, 31]);
+        check_node(nodes[4], 0xa3, 0xb3, &[40, 41, 42, 43]);
+        // The end-entire node is not returned by the iterator.
+        assert_eq!(nodes.len(), 5);
+    }
 }