library: std: sys: net: uefi: tcp: Implement write_vectored

- A working vectored write implementation for TCP4.
- Also introduces a small helper UefiBox intended to be used with heap
  allocated UEFI DSTs.
- Tested on OVMF

Signed-off-by: Ayush Singh <[email protected]>

Author: Ayush1325

library/std/src/sys/net/connection/uefi/mod.rs

@@ -77,12 +77,11 @@ impl TcpStream {
     }
 
     pub fn write_vectored(&self, buf: &[IoSlice<'_>]) -> io::Result<usize> {
-        // FIXME: UEFI does support vectored write, so implement that.
-        crate::io::default_write_vectored(|b| self.write(b), buf)
+        self.inner.write_vectored(buf, self.write_timeout()?)
     }
 
     pub fn is_write_vectored(&self) -> bool {
-        false
+        true
     }
 
     pub fn peer_addr(&self) -> io::Result<SocketAddr> {

library/std/src/sys/net/connection/uefi/tcp.rs

@@ -1,5 +1,5 @@
 use super::tcp4;
-use crate::io;
+use crate::io::{self, IoSlice};
 use crate::net::SocketAddr;
 use crate::ptr::NonNull;
 use crate::sys::{helpers, unsupported};
@@ -28,6 +28,16 @@ impl Tcp {
         }
     }
 
+    pub(crate) fn write_vectored(
+        &self,
+        buf: &[IoSlice<'_>],
+        timeout: Option<Duration>,
+    ) -> io::Result<usize> {
+        match self {
+            Self::V4(client) => client.write_vectored(buf, timeout),
+        }
+    }
+
     pub(crate) fn read(&self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<usize> {
         match self {
             Self::V4(client) => client.read(buf, timeout),

library/std/src/sys/net/connection/uefi/tcp4.rs

@@ -1,7 +1,7 @@
 use r_efi::efi::{self, Status};
 use r_efi::protocols::tcp4;
 
-use crate::io;
+use crate::io::{self, IoSlice};
 use crate::net::SocketAddrV4;
 use crate::ptr::NonNull;
 use crate::sync::atomic::{AtomicBool, Ordering};
@@ -147,6 +147,62 @@ impl Tcp4 {
         }
     }
 
+    pub(crate) fn write_vectored(
+        &self,
+        buf: &[IoSlice<'_>],
+        timeout: Option<Duration>,
+    ) -> io::Result<usize> {
+        let evt = unsafe { self.create_evt() }?;
+        let completion_token =
+            tcp4::CompletionToken { event: evt.as_ptr(), status: Status::SUCCESS };
+        let mut data_len = 0u32;
+        let mut fragment_count = 0u32;
+
+        // Calculate how many IoSlice in buf can be transmitted.
+        for i in buf {
+            match data_len.checked_add(i.as_slice().len() as u32) {
+                Some(x) => data_len = x,
+                None => break,
+            }
+            fragment_count += 1;
+        }
+
+        let tx_data_size = crate::mem::size_of::<tcp4::TransmitData<0>>()
+            + crate::mem::size_of::<tcp4::FragmentData>() * (fragment_count as usize);
+        let mut tx_data = helpers::UefiBox::<tcp4::TransmitData>::new(tx_data_size);
+        unsafe {
+            (*tx_data.as_mut_ptr()).push = r_efi::efi::Boolean::FALSE;
+            (*tx_data.as_mut_ptr()).urgent = r_efi::efi::Boolean::FALSE;
+            (*tx_data.as_mut_ptr()).data_length = data_len;
+            (*tx_data.as_mut_ptr()).fragment_count = fragment_count;
+            // SAFETY: IoSlice and FragmentData are guaranteed to have same layout.
+            crate::ptr::copy_nonoverlapping(
+                buf.as_ptr().cast(),
+                (*tx_data.as_mut_ptr()).fragment_table.as_mut_ptr(),
+                fragment_count as usize,
+            );
+        };
+
+        let protocol = self.protocol.as_ptr();
+        let mut token = tcp4::IoToken {
+            completion_token,
+            packet: tcp4::IoTokenPacket { tx_data: tx_data.as_mut_ptr() },
+        };
+
+        let r = unsafe { ((*protocol).transmit)(protocol, &mut token) };
+        if r.is_error() {
+            return Err(io::Error::from_raw_os_error(r.as_usize()));
+        }
+
+        unsafe { self.wait_or_cancel(timeout, &mut token.completion_token) }?;
+
+        if completion_token.status.is_error() {
+            Err(io::Error::from_raw_os_error(completion_token.status.as_usize()))
+        } else {
+            Ok(data_len as usize)
+        }
+    }
+
     pub(crate) fn read(&self, buf: &mut [u8], timeout: Option<Duration>) -> io::Result<usize> {
         let evt = unsafe { self.create_evt() }?;
         let completion_token =

library/std/src/sys/pal/uefi/helpers.rs

@@ -765,3 +765,19 @@ pub(crate) const fn ipv4_to_r_efi(addr: crate::net::Ipv4Addr) -> efi::Ipv4Addres
 pub(crate) const fn ipv4_from_r_efi(ip: efi::Ipv4Address) -> crate::net::Ipv4Addr {
     crate::net::Ipv4Addr::new(ip.addr[0], ip.addr[1], ip.addr[2], ip.addr[3])
 }
+
+#[repr(transparent)]
+pub(crate) struct UefiBox<T> {
+    inner: Box<[u8]>,
+    phantom: PhantomData<T>,
+}
+
+impl<T> UefiBox<T> {
+    pub(crate) fn new(len: usize) -> Self {
+        Self { inner: vec![0u8; len].into_boxed_slice(), phantom: PhantomData }
+    }
+
+    pub(crate) fn as_mut_ptr(&mut self) -> *mut T {
+        self.inner.as_mut_ptr().cast()
+    }
+}