std: sys: net: uefi: tcp4: Implement write

A blocking implementation of tcp4 write.

Signed-off-by: Ayush Singh <ayush@beagleboard.org>

3 files changed, 53 insertions, 6 deletions

diff --git a/library/std/src/sys/net/connection/uefi/mod.rs b/library/std/src/sys/net/connection/uefi/mod.rs
index 46d67c8e510..2d465bd0efa 100644
--- a/library/std/src/sys/net/connection/uefi/mod.rs
+++ b/library/std/src/sys/net/connection/uefi/mod.rs
@@ -7,7 +7,7 @@ use crate::time::Duration;
 mod tcp;
 pub(crate) mod tcp4;
 
-pub struct TcpStream(#[expect(dead_code)] tcp::Tcp);
+pub struct TcpStream(tcp::Tcp);
 
 impl TcpStream {
     pub fn connect(addr: io::Result<&SocketAddr>) -> io::Result<TcpStream> {
@@ -54,12 +54,13 @@ impl TcpStream {
         false
     }
 
-    pub fn write(&self, _: &[u8]) -> io::Result<usize> {
-        unsupported()
+    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
+        self.0.write(buf)
     }
 
-    pub fn write_vectored(&self, _: &[IoSlice<'_>]) -> io::Result<usize> {
-        unsupported()
+    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)
     }
 
     pub fn is_write_vectored(&self) -> bool {
diff --git a/library/std/src/sys/net/connection/uefi/tcp.rs b/library/std/src/sys/net/connection/uefi/tcp.rs
index f87accdc41d..9c3462e3468 100644
--- a/library/std/src/sys/net/connection/uefi/tcp.rs
+++ b/library/std/src/sys/net/connection/uefi/tcp.rs
@@ -3,7 +3,7 @@ use crate::io;
 use crate::net::SocketAddr;
 
 pub(crate) enum Tcp {
-    V4(#[expect(dead_code)] tcp4::Tcp4),
+    V4(tcp4::Tcp4),
 }
 
 impl Tcp {
@@ -18,4 +18,10 @@ impl Tcp {
             SocketAddr::V6(_) => todo!(),
         }
     }
+
+    pub(crate) fn write(&self, buf: &[u8]) -> io::Result<usize> {
+        match self {
+            Self::V4(client) => client.write(buf),
+        }
+    }
 }
diff --git a/library/std/src/sys/net/connection/uefi/tcp4.rs b/library/std/src/sys/net/connection/uefi/tcp4.rs
index f7ca373b52b..d0f0d27d975 100644
--- a/library/std/src/sys/net/connection/uefi/tcp4.rs
+++ b/library/std/src/sys/net/connection/uefi/tcp4.rs
@@ -88,6 +88,46 @@ impl Tcp4 {
         }
     }
 
+    pub(crate) fn write(&self, buf: &[u8]) -> io::Result<usize> {
+        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::<crate::ffi::c_void>().cast_mut(),
+        };
+        let mut tx_data = tcp4::TransmitData {
+            push: r_efi::efi::Boolean::FALSE,
+            urgent: r_efi::efi::Boolean::FALSE,
+            data_length: data_len,
+            fragment_count: 1,
+            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::<tcp4::TransmitData<0>>(),
+            },
+        };
+
+        let r = unsafe { ((*protocol).transmit)(protocol, &mut token) };
+        if r.is_error() {
+            return Err(io::Error::from_raw_os_error(r.as_usize()));
+        }
+
+        self.wait_for_flag();
+
+        if completion_token.status.is_error() {
+            Err(io::Error::from_raw_os_error(completion_token.status.as_usize()))
+        } else {
+            Ok(data_len as usize)
+        }
+    }
+
     unsafe fn create_evt(&self) -> io::Result<helpers::OwnedEvent> {
         self.flag.store(false, Ordering::Relaxed);
         helpers::OwnedEvent::new(