From a57928e1ba44bcbb3754a84f43312af8e91a13c1 Mon Sep 17 00:00:00 2001
From: Philipp Schuster <phip1611@gmail.com>
Date: Wed, 3 May 2023 21:39:26 +0200
Subject: [PATCH] cstr16: add method to get the underlying bytes

---
 CHANGELOG.md                               |  4 ++
 uefi-macros/src/lib.rs                     |  2 +-
 uefi/src/data_types/strs.rs                | 60 ++++++++++++++++++----
 uefi/src/proto/string/unicode_collation.rs |  4 +-
 4 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c3b0c51e1..be01fc593 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,8 @@
   - `From<&CStr16>` for `String`
   - `From<&CString16>` for `String`
 - Added `RuntimeServices::get_variable_boxed` (requires the `alloc` feature).
+- Added `CStr16::as_bytes`
+- Added `AsRef<[u8]>` and `Borrow<[u8]>` for `Cstr8` and `CStr16`.
 
 ### Changed
 
@@ -44,6 +46,8 @@
 - The `MEMORY_DESCRIPTOR_VERSION` constant has been moved to
   `MemoryDescriptor::VERSION`.
 - The `Revision` struct's one field is now public.
+- Renamed `CStr8::to_bytes` to `CStr8::as_bytes` and changed the semantics:
+  The trailing null character is now always included in the returned slice.
 
 ## uefi-macros - [Unreleased]
 
diff --git a/uefi-macros/src/lib.rs b/uefi-macros/src/lib.rs
index 00ffe9bbb..666c90439 100644
--- a/uefi-macros/src/lib.rs
+++ b/uefi-macros/src/lib.rs
@@ -262,7 +262,7 @@ pub fn entry(args: TokenStream, input: TokenStream) -> TokenStream {
 /// assert_eq!(cstr8!().to_u16_slice_with_nul(), [0]);
 /// assert_eq!(cstr8!("").to_u16_slice_with_nul(), [0]);
 /// // Non-empty string
-/// assert_eq!(cstr8!("test").to_bytes_with_nul(), [116, 101, 115, 116, 0]);
+/// assert_eq!(cstr8!("test").as_bytes(), [116, 101, 115, 116, 0]);
 /// ```
 #[proc_macro]
 pub fn cstr8(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
diff --git a/uefi/src/data_types/strs.rs b/uefi/src/data_types/strs.rs
index 38644976e..25db6afe9 100644
--- a/uefi/src/data_types/strs.rs
+++ b/uefi/src/data_types/strs.rs
@@ -1,6 +1,7 @@
 use super::chars::{Char16, Char8, NUL_16, NUL_8};
 use super::UnalignedSlice;
 use crate::polyfill::maybe_uninit_slice_assume_init_ref;
+use core::borrow::Borrow;
 use core::ffi::CStr;
 use core::iter::Iterator;
 use core::mem::MaybeUninit;
@@ -118,16 +119,10 @@ impl CStr8 {
         self.0.as_ptr()
     }
 
-    /// Converts this CStr8 to a slice of bytes without the terminating null byte.
+    /// Returns the underlying bytes as slice including the terminating null
+    /// character.
     #[must_use]
-    pub fn to_bytes(&self) -> &[u8] {
-        let chars = self.to_bytes_with_nul();
-        &chars[..chars.len() - 1]
-    }
-
-    /// Converts this CStr8 to a slice of bytes containing the trailing null byte.
-    #[must_use]
-    pub const fn to_bytes_with_nul(&self) -> &[u8] {
+    pub const fn as_bytes(&self) -> &[u8] {
         unsafe { &*(&self.0 as *const [Char8] as *const [u8]) }
     }
 }
@@ -147,6 +142,18 @@ impl fmt::Display for CStr8 {
     }
 }
 
+impl AsRef<[u8]> for CStr8 {
+    fn as_ref(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+
+impl Borrow<[u8]> for CStr8 {
+    fn borrow(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+
 impl<StrType: AsRef<str> + ?Sized> EqStrUntilNul<StrType> for CStr8 {
     fn eq_str_until_nul(&self, other: &StrType) -> bool {
         let other = other.as_ref();
@@ -389,6 +396,25 @@ impl CStr16 {
         }
         Ok(())
     }
+
+    /// Returns the underlying bytes as slice including the terminating null
+    /// character.
+    #[must_use]
+    pub fn as_bytes(&self) -> &[u8] {
+        unsafe { slice::from_raw_parts(self.0.as_ptr().cast(), self.num_bytes()) }
+    }
+}
+
+impl AsRef<[u8]> for CStr16 {
+    fn as_ref(&self) -> &[u8] {
+        self.as_bytes()
+    }
+}
+
+impl Borrow<[u8]> for CStr16 {
+    fn borrow(&self) -> &[u8] {
+        self.as_bytes()
+    }
 }
 
 #[cfg(feature = "alloc")]
@@ -522,6 +548,14 @@ mod tests {
         assert!(msg.eq_str_until_nul(cstr8));
     }
 
+    #[test]
+    fn test_cstr8_as_bytes() {
+        let string: &CStr8 = cstr8!("a");
+        assert_eq!(string.as_bytes(), &[b'a', 0]);
+        assert_eq!(<CStr8 as AsRef<[u8]>>::as_ref(string), &[b'a', 0]);
+        assert_eq!(<CStr8 as Borrow<[u8]>>::borrow(string), &[b'a', 0]);
+    }
+
     #[test]
     fn test_cstr16_num_bytes() {
         let s = CStr16::from_u16_with_nul(&[65, 66, 67, 0]).unwrap();
@@ -617,6 +651,14 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_cstr16_as_bytes() {
+        let string: &CStr16 = cstr16!("a");
+        assert_eq!(string.as_bytes(), &[b'a', 0, 0, 0]);
+        assert_eq!(<CStr16 as AsRef<[u8]>>::as_ref(string), &[b'a', 0, 0, 0]);
+        assert_eq!(<CStr16 as Borrow<[u8]>>::borrow(string), &[b'a', 0, 0, 0]);
+    }
+
     // Code generation helper for the compare tests of our CStrX types against "str" and "String"
     // from the standard library.
     #[allow(non_snake_case)]
diff --git a/uefi/src/proto/string/unicode_collation.rs b/uefi/src/proto/string/unicode_collation.rs
index d2e2a4821..c008f3f31 100644
--- a/uefi/src/proto/string/unicode_collation.rs
+++ b/uefi/src/proto/string/unicode_collation.rs
@@ -103,12 +103,12 @@ impl UnicodeCollation {
         fat: &CStr8,
         buf: &'a mut [u16],
     ) -> Result<&'a CStr16, StrConversionError> {
-        if buf.len() < fat.to_bytes_with_nul().len() {
+        if buf.len() < fat.as_bytes().len() {
             return Err(StrConversionError::BufferTooSmall);
         }
         (self.fat_to_str)(
             self,
-            fat.to_bytes_with_nul().len(),
+            fat.as_bytes().len(),
             fat.as_ptr(),
             buf.as_ptr() as *mut _,
         );