Enable Fastboot over TCP

Bug: 328786603
Change-Id: I42a3f03772377ecac7762a805c6e70b516a36989

7 files changed, 622 insertions, 8 deletions

diff --git a/gbl/efi/BUILD b/gbl/efi/BUILD
index 3ea964b..7f20ec0 100644
--- a/gbl/efi/BUILD
+++ b/gbl/efi/BUILD
@@ -26,8 +26,10 @@ rust_binary(
         "src/android_boot.rs",
         "src/avb.rs",
         "src/error.rs",
+        "src/fastboot.rs",
         "src/fuchsia_boot.rs",
         "src/main.rs",
+        "src/net.rs",
         "src/riscv64.rs",
         "src/utils.rs",
     ],
@@ -46,6 +48,7 @@ rust_binary(
         "@gbl//libbootconfig",
         "@gbl//libbootimg",
         "@gbl//libefi",
+        "@gbl//libfastboot",
         "@gbl//libfdt",
         "@gbl//libgbl",
         "@gbl//libstorage",
diff --git a/gbl/efi/src/error.rs b/gbl/efi/src/error.rs
index 60558b9..4d3048b 100644
--- a/gbl/efi/src/error.rs
+++ b/gbl/efi/src/error.rs
@@ -17,8 +17,10 @@ use boot::BootError;
 use bootconfig::BootConfigError;
 use bootimg::ImageError;
 use efi::EfiError;
+use fastboot::TransportError;
 use fdt::FdtError;
 use gbl_storage::StorageError;
+use smoltcp::socket::tcp::{ListenError, RecvError, SendError};
 use zbi::ZbiError;
 
 /// Error types for EFI application.
@@ -32,6 +34,8 @@ pub enum EfiAppError {
     NoFdt,
     NotFound,
     NoZbiImage,
+    PeerClosed,
+    Timeout,
     Unsupported,
 }
 
@@ -114,8 +118,12 @@ composite_enum! {
         EfiError(EfiError),
         FdtError(FdtError),
         ImageError(ImageError),
+        ListenError(ListenError),
+        RecvError(RecvError),
+        SendError(SendError),
         SlotVerifyError(SlotVerifyError<'static>),
         StorageError(StorageError),
+        TransportError(TransportError),
         ZbiError(ZbiError),
     }
 }
diff --git a/gbl/efi/src/fastboot.rs b/gbl/efi/src/fastboot.rs
new file mode 100644
index 0000000..8789e04
--- /dev/null
+++ b/gbl/efi/src/fastboot.rs
@@ -0,0 +1,132 @@
+// Copyright 2024, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This EFI application implements a demo for booting Android/Fuchsia from disk. See
+// bootable/libbootloader/gbl/README.md for how to run the demo. See comments of
+// `android_boot:android_boot_demo()` and `fuchsia_boot:fuchsia_boot_demo()` for
+// supported/unsupported features at the moment.
+
+use crate::error::Result;
+use crate::net::{with_efi_network, EfiTcpSocket};
+use core::fmt::Write;
+use core::str::Split;
+use efi::{efi_print, efi_println, EfiEntry};
+use fastboot::{CommandError, Fastboot, FastbootImplementation, TcpStream, TransportError};
+
+const DEFAULT_TIMEOUT_MS: u64 = 5_000;
+const FASTBOOT_TCP_PORT: u16 = 5554;
+
+struct EfiFastbootTcpTransport<'a, 'b, 'c> {
+    transport_error: &'c mut Result<()>,
+    socket: &'c mut EfiTcpSocket<'a, 'b>,
+}
+
+impl TcpStream for EfiFastbootTcpTransport<'_, '_, '_> {
+    /// Reads to `out` for exactly `out.len()` number bytes from the TCP connection.
+    fn read_exact(&mut self, out: &mut [u8]) -> core::result::Result<(), TransportError> {
+        *self.transport_error = self.socket.receive_exact(out, DEFAULT_TIMEOUT_MS);
+        self.transport_error.as_ref().map_err(|_| TransportError::Others(""))?;
+        Ok(())
+    }
+
+    /// Sends exactly `data.len()` number bytes from `data` to the TCP connection.
+    fn write_exact(&mut self, data: &[u8]) -> core::result::Result<(), TransportError> {
+        *self.transport_error = self.socket.send_exact(data, DEFAULT_TIMEOUT_MS);
+        self.transport_error.as_ref().map_err(|_| TransportError::Others(""))?;
+        Ok(())
+    }
+}
+
+/// TODO(b/328786603): Placeholder only. It'll be replaced by a generic GBL Fastboot implementation
+/// in a separate library.
+pub struct GblFastboot {}
+
+impl FastbootImplementation for GblFastboot {
+    fn get_var(
+        &mut self,
+        _: &str,
+        _: Split<char>,
+        _: &mut [u8],
+    ) -> core::result::Result<usize, CommandError> {
+        Err("Not found".into())
+    }
+
+    fn get_var_all<F>(&mut self, _: F) -> core::result::Result<(), CommandError>
+    where
+        F: FnMut(&str, &[&str], &str),
+    {
+        Ok(())
+    }
+}
+
+/// Internal helper for performing Fastboot over TCP.
+fn fastboot_tcp_usb(
+    socket: &mut EfiTcpSocket,
+    efi_entry: &EfiEntry,
+    download_buffer: &mut [u8],
+) -> Result<()> {
+    let mut gbl_fastboot = GblFastboot {};
+    efi_println!(efi_entry, "Listening for Fastboot over TCP...");
+    efi_println!(efi_entry, "IP addresses:");
+    socket.interface().ip_addrs().iter().for_each(|v| {
+        efi_println!(efi_entry, "\t{}", v.address());
+    });
+    loop {
+        socket.listen(FASTBOOT_TCP_PORT)?;
+        let listen_start = EfiTcpSocket::timestamp(0);
+        loop {
+            socket.poll();
+            if socket.check_active() {
+                // Has connection.
+                efi_println!(
+                    efi_entry,
+                    "Connection from {}",
+                    socket.get_socket().remote_endpoint().unwrap()
+                );
+
+                let mut transport_error = Ok(());
+                let mut transport = EfiFastbootTcpTransport {
+                    transport_error: &mut transport_error,
+                    socket: socket,
+                };
+                let mut fastboot = Fastboot::new(&mut download_buffer[..]);
+                let _ = fastboot.run_tcp_session(&mut transport, &mut gbl_fastboot);
+                efi_println!(efi_entry, "Fastboot TCP session ends. {:?}", transport_error);
+                break;
+            } else if EfiTcpSocket::timestamp(listen_start) > DEFAULT_TIMEOUT_MS {
+                // Reset once in a while in case a remote client disconnects in the middle of
+                // TCP handshake and leaves the socket in a half open state.
+                break;
+            }
+
+            // Perform Fastboot over USB here.
+        }
+    }
+}
+
+/// Runs Fastboot.
+pub fn run_fastboot(efi_entry: &EfiEntry) -> Result<()> {
+    // TODO(b/328786603): Figure out where to get download buffer size.
+    let mut download_buffer = vec![0u8; 512 * 1024 * 1024];
+    match with_efi_network(efi_entry, |socket| -> Result<()> {
+        fastboot_tcp_usb(socket, efi_entry, &mut download_buffer[..])
+    })? {
+        Err(e) => {
+            efi_println!(efi_entry, "Failed to initilaize network {:?}", e);
+            // Perform Fastboot over USB here.
+        }
+        _ => {}
+    };
+    Ok(())
+}
diff --git a/gbl/efi/src/main.rs b/gbl/efi/src/main.rs
index 92e2e29..aad07fb 100644
--- a/gbl/efi/src/main.rs
+++ b/gbl/efi/src/main.rs
@@ -39,7 +39,9 @@ mod riscv64;
 mod android_boot;
 mod avb;
 mod error;
+mod fastboot;
 mod fuchsia_boot;
+mod net;
 
 fn main(image_handle: *mut core::ffi::c_void, systab_ptr: *mut EfiSystemTable) -> Result<()> {
     // SAFETY: Called only once here upon EFI app entry.
@@ -54,7 +56,7 @@ fn main(image_handle: *mut core::ffi::c_void, systab_ptr: *mut EfiSystemTable) -
     match wait_key_stroke(&entry, 'f', 2000) {
         Ok(true) => {
             efi_println!(entry, "'f' pressed.");
-            // TODO(b/328786603): Implement fastboot.
+            fastboot::run_fastboot(&entry)?;
         }
         _ => {}
     }
diff --git a/gbl/efi/src/net.rs b/gbl/efi/src/net.rs
new file mode 100644
index 0000000..0c994f6
--- /dev/null
+++ b/gbl/efi/src/net.rs
@@ -0,0 +1,467 @@
+// Copyright 2024, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use alloc::boxed::Box;
+use alloc::vec::Vec;
+use core::fmt::Write;
+use core::sync::atomic::{AtomicU64, Ordering};
+use efi::defs::{
+    EfiEvent, EfiMacAddress, EFI_STATUS_ALREADY_STARTED, EFI_STATUS_NOT_STARTED,
+    EFI_TIMER_DELAY_TIMER_PERIODIC,
+};
+use efi::{
+    efi_print, efi_println, DeviceHandle, EfiEntry, EventNotify, EventType, Protocol,
+    SimpleNetworkProtocol, Tpl,
+};
+//use fastboot::{Fastboot, FastbootImplementation, FormattedBytes, TcpStream, TransportError};
+use crate::error::{EfiAppError, Result};
+use crate::utils::{get_device_path, loop_with_timeout, ms_to_100ns};
+use smoltcp::iface::{Config, Interface, SocketSet};
+use smoltcp::phy;
+use smoltcp::phy::{Device, DeviceCapabilities, Medium};
+use smoltcp::socket::tcp;
+use smoltcp::time::Instant;
+use smoltcp::wire::{EthernetAddress, IpAddress, IpCidr, Ipv6Address};
+
+/// Maintains a timestamp needed by smoltcp network. It's updated periodically during timer event.
+static NETWORK_TIMESTAMP: AtomicU64 = AtomicU64::new(0);
+/// Ethernet frame size for frame pool.
+const ETHERNET_FRAME_SIZE: usize = 1536;
+// Update period in milliseconds for `NETWORK_TIMESTAMP`.
+const NETWORK_TIMESTAMP_UPDATE_PERIOD: u64 = 50;
+// Size of the socket tx/rx application data buffer.
+const SOCKET_TX_RX_BUFFER: usize = 64 * 1024;
+
+/// Performs a shutdown and restart of the simple network protocol.
+fn reset_simple_network<'a>(snp: &Protocol<'a, SimpleNetworkProtocol>) -> Result<()> {
+    match snp.shutdown() {
+        Err(e) if !e.is_efi_err(EFI_STATUS_NOT_STARTED) => return Err(e.into()),
+        _ => {}
+    };
+
+    match snp.start() {
+        Err(e) if !e.is_efi_err(EFI_STATUS_ALREADY_STARTED) => {
+            return Err(e.into());
+        }
+        _ => {}
+    };
+    snp.initialize(0, 0).unwrap();
+    Ok(snp.reset(true)?)
+}
+
+/// `EfiNetworkDevice` manages a frame pool and handles receiving/sending network frames.
+pub struct EfiNetworkDevice<'a> {
+    protocol: Protocol<'a, SimpleNetworkProtocol>,
+    rx_frame: Box<[u8; ETHERNET_FRAME_SIZE]>,
+    tx_frames: Vec<*mut [u8; ETHERNET_FRAME_SIZE]>,
+    tx_frame_curr: usize, // Circular next index into tx_frames.
+    efi_entry: &'a EfiEntry,
+}
+
+impl<'a> EfiNetworkDevice<'a> {
+    /// Creates an new instance. Allocates `extra_tx_frames+1` number of TX frames.
+    pub fn new(
+        protocol: Protocol<'a, SimpleNetworkProtocol>,
+        extra_tx_frames: usize,
+        efi_entry: &'a EfiEntry,
+    ) -> Self {
+        let mut ret = Self {
+            protocol: protocol,
+            rx_frame: Box::new([0u8; ETHERNET_FRAME_SIZE]),
+            tx_frames: vec![core::ptr::null_mut(); extra_tx_frames + 1],
+            tx_frame_curr: 0,
+            efi_entry: efi_entry,
+        };
+        ret.tx_frames
+            .iter_mut()
+            .for_each(|v| *v = Box::into_raw(Box::new([0u8; ETHERNET_FRAME_SIZE])));
+        ret
+    }
+}
+
+impl Drop for EfiNetworkDevice<'_> {
+    fn drop(&mut self) {
+        if let Err(e) = self.protocol.shutdown() {
+            if !e.is_efi_err(EFI_STATUS_NOT_STARTED) {
+                // If shutdown fails, the protocol might still be operating on transmit buffers,
+                // which can cause undefined behavior. Thus we need to panic.
+                panic!("Failed to shutdown EFI network. {:?}", e);
+            }
+        }
+
+        // Deallocate TX frames.
+        self.tx_frames.iter_mut().for_each(|v| {
+            // SAFETY:
+            // Each pointer is created by `Box::new()` in `EfiNetworkDevice::new()`. Thus the
+            // pointer is valid and layout matches.
+            let _ = unsafe { Box::from_raw(v) };
+        });
+    }
+}
+
+// Implements network device trait backend for the `smoltcp` crate.
+impl<'a> Device for EfiNetworkDevice<'a> {
+    type RxToken<'b> = RxToken<'b> where Self: 'b;
+    type TxToken<'b> = TxToken<'a, 'b> where Self: 'b;
+
+    fn capabilities(&self) -> DeviceCapabilities {
+        // Taken from upstream example.
+        let mut res: DeviceCapabilities = Default::default();
+        res.max_transmission_unit = 65535;
+        res.medium = Medium::Ethernet;
+        res
+    }
+
+    fn receive(&mut self, _: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
+        let mut recv_size = self.rx_frame.len();
+        // Receive the next packet from the device.
+        self.protocol
+            .receive(None, Some(&mut recv_size), &mut self.rx_frame[..], None, None, None)
+            .ok()?;
+        match recv_size > 0 {
+            true => Some((
+                RxToken(&mut self.rx_frame[..recv_size]),
+                TxToken {
+                    protocol: &self.protocol,
+                    tx_frames: &mut self.tx_frames[..],
+                    curr: &mut self.tx_frame_curr,
+                    efi_entry: self.efi_entry,
+                },
+            )),
+            _ => None,
+        }
+    }
+
+    fn transmit(&mut self, _: Instant) -> Option<Self::TxToken<'_>> {
+        Some(TxToken {
+            protocol: &self.protocol,
+            tx_frames: &mut self.tx_frames[..],
+            curr: &mut self.tx_frame_curr,
+            efi_entry: self.efi_entry,
+        })
+    }
+}
+
+/// In smoltcp, a `RxToken` is used to receive/process a frame when consumed.
+pub struct RxToken<'a>(&'a mut [u8]);
+
+impl phy::RxToken for RxToken<'_> {
+    fn consume<R, F>(self, f: F) -> R
+    where
+        F: FnOnce(&mut [u8]) -> R,
+    {
+        f(self.0)
+    }
+}
+
+/// In smoltcp, a `TxToken` is used to transmit a frame when consumed.
+pub struct TxToken<'a: 'b, 'b> {
+    tx_frames: &'b mut [*mut [u8; ETHERNET_FRAME_SIZE]],
+    curr: &'b mut usize,
+    protocol: &'b Protocol<'a, SimpleNetworkProtocol>,
+    efi_entry: &'b EfiEntry,
+}
+
+impl TxToken<'_, '_> {
+    /// Tries to allocate a send buffer.
+    fn try_get_buffer(&mut self) -> Option<*mut [u8; ETHERNET_FRAME_SIZE]> {
+        let mut ptr: *mut core::ffi::c_void = core::ptr::null_mut();
+        let mut interrupt_status = 0u32;
+        // Recyle a buffer or take one from `tx_frames`.
+        match self.protocol.get_status(Some(&mut interrupt_status), Some(&mut ptr)) {
+            Ok(()) if self.tx_frames.contains(&(ptr as *mut _)) => Some(ptr as *mut _),
+            _ if *self.curr < self.tx_frames.len() => {
+                // If we can't recycle a buffer, see if we can take one from the pool.
+                let res = *self.curr;
+                *self.curr = *self.curr + 1;
+                Some(self.tx_frames[res])
+            }
+            _ => None,
+        }
+    }
+}
+
+impl phy::TxToken for TxToken<'_, '_> {
+    fn consume<R, F>(mut self, len: usize, f: F) -> R
+    where
+        F: FnOnce(&mut [u8]) -> R,
+    {
+        loop {
+            match loop_with_timeout(self.efi_entry, 5000, || self.try_get_buffer().ok_or(false)) {
+                Ok(Some(send_buffer)) => {
+                    // SAFETY:
+                    // * The pointer is confirmed to come from one of `self.tx_frames`. It's
+                    //   created via `Box::new()` in `EfiNetworkDevice::new()`. Thus it is properly
+                    //   aligned, dereferenceable and initialized.
+                    // * The pointer is either recycled from `self.protocol.get_status` or newly
+                    //   allocated from `self.tx_frames`. Thus There's no other references to it.
+                    // * The reference is only used for passing to `f` and goes out of scope
+                    //   immediately after.
+                    let result = f(&mut unsafe { send_buffer.as_mut() }.unwrap()[..len]);
+
+                    // SAFETY:
+                    // * `send_buffer` comes from `EfiNetworkDevice::tx_frames`. It has a valid
+                    //   length at least `len`. `EfiNetworkDevice` shuts down network on drop. Thus
+                    //   the transmit buffer remains valid throughout the operation of the network
+                    //   protocol.
+                    // * `send_buffer` is either recycled from `self.protocol.get_status()` or newly
+                    //   allocated from `self.tx_frames`. There's no other references to it.
+                    // * `self.curr` stricly increases for each new allocation until
+                    //   `reset_simple_network()`. Thus there'll be no other references to the buffer
+                    //    until it is either recycled or `reset_simple_network()` is called.
+                    let _ = unsafe {
+                        self.protocol.transmit(
+                            0,
+                            send_buffer.as_mut().unwrap().get_mut(..len).unwrap(),
+                            Default::default(), // Src mac address don't care
+                            Default::default(), // Dest mac address don't care
+                            0,
+                        )
+                    };
+
+                    return result;
+                }
+                Ok(None) => {
+                    // Some UEFI firmware has internal network service that also recycle buffers,
+                    // in which case our buffer may be hijacked and will never be returned from our
+                    // call. If we run into this case, shutdown and restart the network and try
+                    // again. Shutting down network releases all pending send/receive buffers
+                    // internally retained.
+                    efi_println!(
+                        self.efi_entry,
+                        "Timeout recycling TX buffers. Resetting network."
+                    );
+                    // Panics if this fails, as we have effectively lost control over network's
+                    // used of buffers.
+                    reset_simple_network(self.protocol).unwrap();
+                    *self.curr = 0;
+                }
+                _ => {} // `loop_with_timeout` failure. Try again.
+            };
+        }
+    }
+}
+
+/// Returns the current value of timestamp.
+fn timestamp() -> u64 {
+    NETWORK_TIMESTAMP.load(Ordering::Relaxed)
+}
+
+/// Returns a smoltcp time `Instant` value.
+fn time_instant() -> Instant {
+    Instant::from_millis(i64::try_from(timestamp()).unwrap())
+}
+
+/// Find the first available network device.
+fn find_net_device(efi_entry: &EfiEntry) -> Result<DeviceHandle> {
+    // Find the device whose path is the "smallest" lexicographically, this ensures that it's not
+    // any child network device of some other node. e1000 tends to add a child network device for
+    // ipv4 and ipv6 configuration information.
+    efi_entry
+        .system_table()
+        .boot_services()
+        .locate_handle_buffer_by_protocol::<SimpleNetworkProtocol>()?
+        .handles()
+        .iter()
+        .map(|handle| (*handle, get_device_path(efi_entry, *handle)))
+        // Ignore devices that fail to get device path.
+        .filter_map(|(handle, path)| path.ok().map(|v| (handle, v)))
+        // Ignore devices that have NULL path.
+        .filter_map(|(handle, path)| path.text().is_some().then(|| (handle, path)))
+        // Finds the minimum path lexicographically.
+        .min_by(|lhs, rhs| Ord::cmp(lhs.1.text().unwrap(), rhs.1.text().unwrap()))
+        .map(|(h, _)| h)
+        .ok_or(EfiAppError::NotFound.into())
+}
+
+/// Derives a link local ethernet mac address and IPv6 address from `EfiMacAddress`.
+fn ll_mac_ip6_addr_from_efi_mac(mac: EfiMacAddress) -> (EthernetAddress, IpAddress) {
+    let ll_mac_bytes = &mac.addr[..6];
+    let mut ip6_bytes = [0u8; 16];
+    ip6_bytes[0] = 0xfe;
+    ip6_bytes[1] = 0x80;
+    ip6_bytes[8] = ll_mac_bytes[0] ^ 2;
+    ip6_bytes[9] = ll_mac_bytes[1];
+    ip6_bytes[10] = ll_mac_bytes[2];
+    ip6_bytes[11] = 0xff;
+    ip6_bytes[12] = 0xfe;
+    ip6_bytes[13] = ll_mac_bytes[3];
+    ip6_bytes[14] = ll_mac_bytes[4];
+    ip6_bytes[15] = ll_mac_bytes[5];
+
+    (
+        EthernetAddress::from_bytes(ll_mac_bytes),
+        IpAddress::Ipv6(Ipv6Address::from_bytes(&ip6_bytes[..])),
+    )
+}
+
+/// `EfiTcpSocket` groups together necessary components for performing TCP.
+pub struct EfiTcpSocket<'a, 'b> {
+    efi_net_dev: &'b mut EfiNetworkDevice<'a>,
+    interface: &'b mut Interface,
+    sockets: &'b mut SocketSet<'b>,
+    efi_entry: &'a EfiEntry,
+}
+
+impl<'a, 'b> EfiTcpSocket<'a, 'b> {
+    /// Resets the socket and starts listening for new TCP connection.
+    pub fn listen(&mut self, port: u16) -> Result<()> {
+        self.get_socket().abort();
+        self.get_socket().listen(port)?;
+        Ok(())
+    }
+
+    /// Polls network device.
+    pub fn poll(&mut self) {
+        self.interface.poll(time_instant(), self.efi_net_dev, self.sockets);
+    }
+
+    /// Polls network and check if the socket is in an active state.
+    pub fn check_active(&mut self) -> bool {
+        self.poll();
+        self.get_socket().is_active()
+    }
+
+    /// Gets a reference to the smoltcp socket object.
+    pub fn get_socket(&mut self) -> &mut tcp::Socket<'b> {
+        // We only consider single socket use case for now.
+        let handle = self.sockets.iter().next().unwrap().0;
+        self.sockets.get_mut::<tcp::Socket>(handle)
+    }
+
+    /// Checks whether a socket is closed.
+    fn is_closed(&mut self) -> bool {
+        return !self.get_socket().is_open() || self.get_socket().state() == tcp::State::CloseWait;
+    }
+
+    /// Receives exactly `out.len()` number of bytes to `out`.
+    pub fn receive_exact(&mut self, out: &mut [u8], timeout: u64) -> Result<()> {
+        let mut recv_size = 0;
+        loop_with_timeout(self.efi_entry, timeout, || -> core::result::Result<Result<()>, bool> {
+            self.poll();
+            if self.is_closed() {
+                return Ok(Err(EfiAppError::PeerClosed.into()));
+            } else if self.get_socket().can_recv() {
+                let this_recv = match self.get_socket().recv_slice(&mut out[recv_size..]) {
+                    Err(e) => return Ok(Err(e.into())),
+                    Ok(v) => v,
+                };
+                recv_size += this_recv;
+                if recv_size == out.len() {
+                    return Ok(Ok(()));
+                }
+
+                return Err(this_recv > 0);
+            }
+            Err(false)
+        })?
+        .ok_or(EfiAppError::Timeout)?
+    }
+
+    /// Sends exactly `data.len()` number of bytes from `data`.
+    pub fn send_exact(&mut self, data: &[u8], timeout: u64) -> Result<()> {
+        let mut sent_size = 0;
+        let mut last_send_queue = 0usize;
+        loop_with_timeout(self.efi_entry, timeout, || -> core::result::Result<Result<()>, bool> {
+            self.poll();
+            if sent_size == data.len() && self.get_socket().send_queue() == 0 {
+                return Ok(Ok(()));
+            } else if self.is_closed() {
+                return Ok(Err(EfiAppError::PeerClosed.into()));
+            }
+            // As long as some data is sent, reset the timeout.
+            let reset = self.get_socket().send_queue() != last_send_queue;
+            if self.get_socket().can_send() && sent_size < data.len() {
+                sent_size += match self.get_socket().send_slice(&data[sent_size..]) {
+                    Err(e) => return Ok(Err(e.into())),
+                    Ok(v) => v,
+                };
+            }
+            last_send_queue = self.get_socket().send_queue();
+            Err(reset)
+        })?
+        .ok_or(EfiAppError::Timeout)?
+    }
+
+    /// Gets the smoltcp `Interface` for this socket.
+    pub fn interface(&self) -> &Interface {
+        self.interface
+    }
+
+    /// Returns the number of milliseconds elapsed since the `base` timestamp.
+    pub fn timestamp(base: u64) -> u64 {
+        let curr = timestamp();
+        // Assume there can be at most one overflow.
+        match curr < base {
+            true => u64::MAX - (base - curr),
+            false => curr - base,
+        }
+    }
+}
+
+/// Initializes network environment and provides a TCP socket for callers to run a closure. The API
+/// handles clean up automatically after returning.
+pub fn with_efi_network<F, R>(efi_entry: &EfiEntry, mut f: F) -> Result<R>
+where
+    F: FnMut(&mut EfiTcpSocket) -> R,
+{
+    let bs = efi_entry.system_table().boot_services();
+
+    // Creates timestamp update event.
+    let _ = NETWORK_TIMESTAMP.swap(0, Ordering::Relaxed);
+    let mut notify_fn = |_: EfiEvent| {
+        NETWORK_TIMESTAMP.fetch_add(NETWORK_TIMESTAMP_UPDATE_PERIOD, Ordering::Relaxed);
+    };
+    let mut notify = EventNotify::new(Tpl::Callback, &mut notify_fn);
+    let timer = bs.create_event(EventType::TimerNotifySignal, Some(&mut notify))?;
+    bs.set_timer(
+        &timer,
+        EFI_TIMER_DELAY_TIMER_PERIODIC,
+        ms_to_100ns(NETWORK_TIMESTAMP_UPDATE_PERIOD)?,
+    )?;
+
+    // Creates and initializes simple network protocol.
+    let snp_dev = find_net_device(efi_entry).unwrap();
+    let snp = bs.open_protocol::<SimpleNetworkProtocol>(snp_dev).unwrap();
+    reset_simple_network(&snp)?;
+
+    // Gets our MAC address and IPv6 address.
+    // We can also consider getting this from vendor configuration.
+    let (ll_mac, ll_ip6_addr) = ll_mac_ip6_addr_from_efi_mac(snp.mode()?.current_address);
+
+    // Creates an `EfiNetworkDevice`.
+    // Allocates 7(chosen randomly) extra TX frames. Revisits if it is not enough.
+    let mut efi_net_dev = EfiNetworkDevice::new(snp, 7, &efi_entry);
+    // Configures smoltcp network interface.
+    let mut interface =
+        Interface::new(Config::new(ll_mac.into()), &mut efi_net_dev, time_instant());
+    interface.update_ip_addrs(|ip_addrs| ip_addrs.push(IpCidr::new(ll_ip6_addr, 64)).unwrap());
+    // Creates an instance of socket.
+    let mut tx_buffer = vec![0u8; SOCKET_TX_RX_BUFFER];
+    let mut rx_buffer = vec![0u8; SOCKET_TX_RX_BUFFER];
+    let tx_socket_buffer = tcp::SocketBuffer::new(&mut tx_buffer[..]);
+    let rx_socket_buffer = tcp::SocketBuffer::new(&mut rx_buffer[..]);
+    let socket = tcp::Socket::new(rx_socket_buffer, tx_socket_buffer);
+    let mut sockets: [_; 1] = Default::default();
+    let mut sockets = SocketSet::new(&mut sockets[..]);
+    let _ = sockets.add(socket);
+    let mut socket = EfiTcpSocket {
+        efi_net_dev: &mut efi_net_dev,
+        interface: &mut interface,
+        sockets: &mut sockets,
+        efi_entry: efi_entry,
+    };
+
+    Ok(f(&mut socket))
+}
diff --git a/gbl/efi/src/utils.rs b/gbl/efi/src/utils.rs
index c02b7bb..d74989e 100644
--- a/gbl/efi/src/utils.rs
+++ b/gbl/efi/src/utils.rs
@@ -184,7 +184,7 @@ pub fn cstr_bytes_to_str(data: &[u8]) -> Result<&str> {
 }
 
 /// Converts 1 ms to number of 100 nano seconds
-fn ms_to_100ns(ms: u64) -> Result<u64> {
+pub fn ms_to_100ns(ms: u64) -> Result<u64> {
     Ok(ms.checked_mul(1000 * 10).ok_or(EfiAppError::ArithmeticOverflow)?)
 }
 
diff --git a/gbl/libefi/src/protocol.rs b/gbl/libefi/src/protocol.rs
index d8961be..ef7c31a 100644
--- a/gbl/libefi/src/protocol.rs
+++ b/gbl/libefi/src/protocol.rs
@@ -486,19 +486,21 @@ impl<'a> Protocol<'a, SimpleNetworkProtocol> {
     ///
     /// # Safety
     ///
-    /// Caller must ensure that `buf` remains valid until either 1) the buffer address re-appears
-    /// in `recycled_buffer` from `Self::get_status()`, or 2) Self::Shutdown() is called and
-    /// returns either Ok(()) or EFI_STATUS_NOT_STARTED. Otherwise, the driver may still have
-    /// modifiable access to the buffer and causes undefined behavior if the buffer goes out of
-    /// scope earlier.
+    /// * `buf` needs to be a valid buffer.
+    /// * There should not be any existing references to memory pointed by `buf`.
+    /// * Because `buf` is internally retained by the network. `buf` should remain valid and not
+    ///   dereferenced until either 1) the buffer address re-appears in `recycled_buffer` from
+    ///   `Self::get_status()` or 2) Self::Shutdown() is called and returns either Ok(()) or
+    ///   EFI_STATUS_NOT_STARTED.
     pub unsafe fn transmit(
         &self,
         header_size: usize,
-        buf: &mut [u8],
+        buf: *mut [u8],
         mut src: EfiMacAddress,
         mut dest: EfiMacAddress,
         mut protocol: u16,
     ) -> EfiResult<()> {
+        let buf = buf.as_mut().unwrap();
         // SAFETY:
         // See safety reasoning of `start()`.
         // All pointers passed are valid, outlive the call and are not retained by the call.