Rollup merge of #144350 - Ayush1325:uefi-io, r=tgross35,nicholasbishop

std: sys: io: io_slice: Add UEFI types

UEFI networking APIs do support vectored read/write. While the types for UDP4, UDP6, TCP4 and TCP6 are defined separately, they are essentially the same C struct. So we can map IoSlice and IoSliceMut to have the same binary representation.

Since all UEFI networking types for read/write are DSTs, `IoSlice` and `IoSliceMut` will need to be copied to the end of the transmit/receive structures. So having the same binary representation just allows us to do a single memcpy instead of having to loop and set the DST.

cc ``@nicholasbishop``

3 files changed, 196 insertions, 0 deletions

diff --git a/library/std/src/sys/io/io_slice/uefi.rs b/library/std/src/sys/io/io_slice/uefi.rs
new file mode 100644
index 00000000000..909cfbea0b7
--- /dev/null
+++ b/library/std/src/sys/io/io_slice/uefi.rs
@@ -0,0 +1,74 @@
+//! A buffer type used with `Write::write_vectored` for UEFI Networking APIs. Vectored writing to
+//! File is not supported as of UEFI Spec 2.11.
+
+use crate::marker::PhantomData;
+use crate::slice;
+
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub struct IoSlice<'a> {
+    len: u32,
+    data: *const u8,
+    _p: PhantomData<&'a [u8]>,
+}
+
+impl<'a> IoSlice<'a> {
+    #[inline]
+    pub fn new(buf: &'a [u8]) -> IoSlice<'a> {
+        let len = buf.len().try_into().unwrap();
+        Self { len, data: buf.as_ptr(), _p: PhantomData }
+    }
+
+    #[inline]
+    pub fn advance(&mut self, n: usize) {
+        self.len = u32::try_from(n)
+            .ok()
+            .and_then(|n| self.len.checked_sub(n))
+            .expect("advancing IoSlice beyond its length");
+        unsafe { self.data = self.data.add(n) };
+    }
+
+    #[inline]
+    pub const fn as_slice(&self) -> &'a [u8] {
+        unsafe { slice::from_raw_parts(self.data, self.len as usize) }
+    }
+}
+
+#[repr(C)]
+pub struct IoSliceMut<'a> {
+    len: u32,
+    data: *mut u8,
+    _p: PhantomData<&'a mut [u8]>,
+}
+
+impl<'a> IoSliceMut<'a> {
+    #[inline]
+    pub fn new(buf: &'a mut [u8]) -> IoSliceMut<'a> {
+        let len = buf.len().try_into().unwrap();
+        Self { len, data: buf.as_mut_ptr(), _p: PhantomData }
+    }
+
+    #[inline]
+    pub fn advance(&mut self, n: usize) {
+        self.len = u32::try_from(n)
+            .ok()
+            .and_then(|n| self.len.checked_sub(n))
+            .expect("advancing IoSlice beyond its length");
+        unsafe { self.data = self.data.add(n) };
+    }
+
+    #[inline]
+    pub fn as_slice(&self) -> &[u8] {
+        unsafe { slice::from_raw_parts(self.data, self.len as usize) }
+    }
+
+    #[inline]
+    pub const fn into_slice(self) -> &'a mut [u8] {
+        unsafe { slice::from_raw_parts_mut(self.data, self.len as usize) }
+    }
+
+    #[inline]
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        unsafe { slice::from_raw_parts_mut(self.data, self.len as usize) }
+    }
+}
diff --git a/library/std/src/sys/io/mod.rs b/library/std/src/sys/io/mod.rs
index 4d0365d42fd..ae75f4d97b4 100644
--- a/library/std/src/sys/io/mod.rs
+++ b/library/std/src/sys/io/mod.rs
@@ -11,6 +11,9 @@ mod io_slice {
         } else if #[cfg(target_os = "wasi")] {
             mod wasi;
             pub use wasi::*;
+        } else if #[cfg(target_os = "uefi")] {
+            mod uefi;
+            pub use uefi::*;
         } else {
             mod unsupported;
             pub use unsupported::*;
diff --git a/library/std/src/sys/pal/uefi/tests.rs b/library/std/src/sys/pal/uefi/tests.rs
index 38658cc4e9a..49e75a1a70d 100644
--- a/library/std/src/sys/pal/uefi/tests.rs
+++ b/library/std/src/sys/pal/uefi/tests.rs
@@ -1,5 +1,6 @@
 use super::alloc::*;
 use super::time::*;
+use crate::io::{IoSlice, IoSliceMut};
 use crate::time::Duration;
 
 #[test]
@@ -39,3 +40,121 @@ fn epoch() {
     };
     assert_eq!(system_time_internal::uefi_time_to_duration(t), Duration::new(0, 0));
 }
+
+// UEFI IoSlice and IoSliceMut Tests
+//
+// Strictly speaking, vectored read/write types for UDP4, UDP6, TCP4, TCP6 are defined
+// separately in the UEFI Spec. However, they have the same signature. These tests just ensure
+// that `IoSlice` and `IoSliceMut` are compatible with the vectored types for all the
+// networking protocols.
+
+unsafe fn to_slice<T>(val: &T) -> &[u8] {
+    let len = size_of_val(val);
+    unsafe { crate::slice::from_raw_parts(crate::ptr::from_ref(val).cast(), len) }
+}
+
+#[test]
+fn io_slice_single() {
+    let mut data = [0, 1, 2, 3, 4];
+
+    let tcp4_frag = r_efi::protocols::tcp4::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let tcp6_frag = r_efi::protocols::tcp6::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let udp4_frag = r_efi::protocols::udp4::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let udp6_frag = r_efi::protocols::udp6::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let io_slice = IoSlice::new(&data);
+
+    unsafe {
+        assert_eq!(to_slice(&io_slice), to_slice(&tcp4_frag));
+        assert_eq!(to_slice(&io_slice), to_slice(&tcp6_frag));
+        assert_eq!(to_slice(&io_slice), to_slice(&udp4_frag));
+        assert_eq!(to_slice(&io_slice), to_slice(&udp6_frag));
+    }
+}
+
+#[test]
+fn io_slice_mut_single() {
+    let mut data = [0, 1, 2, 3, 4];
+
+    let tcp4_frag = r_efi::protocols::tcp4::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let tcp6_frag = r_efi::protocols::tcp6::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let udp4_frag = r_efi::protocols::udp4::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let udp6_frag = r_efi::protocols::udp6::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let io_slice_mut = IoSliceMut::new(&mut data);
+
+    unsafe {
+        assert_eq!(to_slice(&io_slice_mut), to_slice(&tcp4_frag));
+        assert_eq!(to_slice(&io_slice_mut), to_slice(&tcp6_frag));
+        assert_eq!(to_slice(&io_slice_mut), to_slice(&udp4_frag));
+        assert_eq!(to_slice(&io_slice_mut), to_slice(&udp6_frag));
+    }
+}
+
+#[test]
+fn io_slice_multi() {
+    let mut data = [0, 1, 2, 3, 4];
+
+    let tcp4_frag = r_efi::protocols::tcp4::FragmentData {
+        fragment_length: data.len().try_into().unwrap(),
+        fragment_buffer: data.as_mut_ptr().cast(),
+    };
+    let rhs =
+        [tcp4_frag.clone(), tcp4_frag.clone(), tcp4_frag.clone(), tcp4_frag.clone(), tcp4_frag];
+    let lhs = [
+        IoSlice::new(&data),
+        IoSlice::new(&data),
+        IoSlice::new(&data),
+        IoSlice::new(&data),
+        IoSlice::new(&data),
+    ];
+
+    unsafe {
+        assert_eq!(to_slice(&lhs), to_slice(&rhs));
+    }
+}
+
+#[test]
+fn io_slice_basic() {
+    let data = [0, 1, 2, 3, 4];
+    let mut io_slice = IoSlice::new(&data);
+
+    assert_eq!(data, io_slice.as_slice());
+    io_slice.advance(2);
+    assert_eq!(&data[2..], io_slice.as_slice());
+}
+
+#[test]
+fn io_slice_mut_basic() {
+    let data = [0, 1, 2, 3, 4];
+    let mut data_clone = [0, 1, 2, 3, 4];
+    let mut io_slice_mut = IoSliceMut::new(&mut data_clone);
+
+    assert_eq!(data, io_slice_mut.as_slice());
+    assert_eq!(data, io_slice_mut.as_mut_slice());
+
+    io_slice_mut.advance(2);
+    assert_eq!(&data[2..], io_slice_mut.into_slice());
+}