diff --git a/library/std/src/sys/net/connection/uefi/mod.rs b/library/std/src/sys/net/connection/uefi/mod.rs index 16e3487a1747f..9f11a3e7af545 100644 --- a/library/std/src/sys/net/connection/uefi/mod.rs +++ b/library/std/src/sys/net/connection/uefi/mod.rs @@ -77,12 +77,11 @@ impl TcpStream { } pub fn write_vectored(&self, buf: &[IoSlice<'_>]) -> io::Result { - // 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 { diff --git a/library/std/src/sys/net/connection/uefi/tcp.rs b/library/std/src/sys/net/connection/uefi/tcp.rs index aac97007bbfe5..16283e64fb35a 100644 --- a/library/std/src/sys/net/connection/uefi/tcp.rs +++ b/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, + ) -> io::Result { + match self { + Self::V4(client) => client.write_vectored(buf, timeout), + } + } + pub(crate) fn read(&self, buf: &mut [u8], timeout: Option) -> io::Result { match self { Self::V4(client) => client.read(buf, timeout), diff --git a/library/std/src/sys/net/connection/uefi/tcp4.rs b/library/std/src/sys/net/connection/uefi/tcp4.rs index 75862ff247b4f..fd364c09502a4 100644 --- a/library/std/src/sys/net/connection/uefi/tcp4.rs +++ b/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}; @@ -108,11 +108,7 @@ impl Tcp4 { } pub(crate) fn write(&self, buf: &[u8], timeout: Option) -> io::Result { - let evt = unsafe { self.create_evt() }?; - let completion_token = - tcp4::CompletionToken { event: evt.as_ptr(), status: Status::SUCCESS }; let data_len = u32::try_from(buf.len()).unwrap_or(u32::MAX); - let fragment = tcp4::FragmentData { fragment_length: data_len, fragment_buffer: buf.as_ptr().cast::().cast_mut(), @@ -125,14 +121,63 @@ impl Tcp4 { fragment_table: [fragment], }; - let protocol = self.protocol.as_ptr(); - let mut token = tcp4::IoToken { - completion_token, - packet: tcp4::IoTokenPacket { - tx_data: (&raw mut tx_data).cast::>(), - }, + self.write_inner((&raw mut tx_data).cast(), timeout).map(|_| data_len as usize) + } + + pub(crate) fn write_vectored( + &self, + buf: &[IoSlice<'_>], + timeout: Option, + ) -> io::Result { + let mut data_length = 0u32; + let mut fragment_count = 0u32; + + // Calculate how many IoSlice in buf can be transmitted. + for i in buf { + // IoSlice length is always <= u32::MAX in UEFI. + match data_length + .checked_add(u32::try_from(i.as_slice().len()).expect("value is stored as a u32")) + { + Some(x) => data_length = x, + None => break, + } + fragment_count += 1; + } + + let tx_data_size = size_of::>() + + size_of::() * (fragment_count as usize); + let mut tx_data = helpers::UefiBox::::new(tx_data_size); + tx_data.write(tcp4::TransmitData { + push: r_efi::efi::Boolean::FALSE, + urgent: r_efi::efi::Boolean::FALSE, + data_length, + fragment_count, + fragment_table: [], + }); + unsafe { + // 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, + ); }; + self.write_inner(tx_data.as_mut_ptr(), timeout).map(|_| data_length as usize) + } + + fn write_inner( + &self, + tx_data: *mut tcp4::TransmitData, + timeout: Option, + ) -> io::Result<()> { + let evt = unsafe { self.create_evt() }?; + let completion_token = + tcp4::CompletionToken { event: evt.as_ptr(), status: Status::SUCCESS }; + + let protocol = self.protocol.as_ptr(); + let mut token = tcp4::IoToken { completion_token, packet: tcp4::IoTokenPacket { tx_data } }; + let r = unsafe { ((*protocol).transmit)(protocol, &mut token) }; if r.is_error() { return Err(io::Error::from_raw_os_error(r.as_usize())); @@ -143,7 +188,7 @@ impl Tcp4 { if completion_token.status.is_error() { Err(io::Error::from_raw_os_error(completion_token.status.as_usize())) } else { - Ok(data_len as usize) + Ok(()) } } diff --git a/library/std/src/sys/pal/uefi/helpers.rs b/library/std/src/sys/pal/uefi/helpers.rs index b50574de937aa..5735406959e86 100644 --- a/library/std/src/sys/pal/uefi/helpers.rs +++ b/library/std/src/sys/pal/uefi/helpers.rs @@ -12,6 +12,7 @@ use r_efi::efi::{self, Guid}; use r_efi::protocols::{device_path, device_path_to_text, service_binding, shell}; +use crate::alloc::Layout; use crate::ffi::{OsStr, OsString}; use crate::io::{self, const_error}; use crate::marker::PhantomData; @@ -765,3 +766,35 @@ 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]) } + +/// This type is intended for use with ZSTs. Since such types are unsized, a reference to such types +/// is not valid in Rust. Thus, only pointers should be used when interacting with such types. +pub(crate) struct UefiBox { + inner: NonNull, + size: usize, +} + +impl UefiBox { + pub(crate) fn new(len: usize) -> Self { + assert!(len >= size_of::()); + // UEFI always expects types to be 8 byte aligned. + let layout = Layout::from_size_align(len, 8).unwrap(); + let ptr = unsafe { crate::alloc::alloc(layout) }; + Self { inner: NonNull::new(ptr.cast()).unwrap(), size: len } + } + + pub(crate) fn write(&mut self, data: T) { + unsafe { self.inner.write(data) } + } + + pub(crate) fn as_mut_ptr(&mut self) -> *mut T { + self.inner.as_ptr().cast() + } +} + +impl Drop for UefiBox { + fn drop(&mut self) { + let layout = Layout::from_size_align(self.size, 8).unwrap(); + unsafe { crate::alloc::dealloc(self.inner.as_ptr().cast(), layout) }; + } +}