[STALE] Model TCP protocol and implement a client

uefi-raw/CHANGELOG.md

@@ -3,6 +3,9 @@
 ## Added
 - Added `TimestampProtocol`.
 - Added `DevicePathToTextProtocol` and `DevicePathFromTextProtocol`.
+- Added minor utility methods to `Ipv4Address`: `new(u8, u8, u8, u8) -> Self` and `zero() -> Self`.
+- Added `Ip4ModeData`, `Ip4ConfigData`, and `Ip4IcmpType`.
+- Added `TCPv4Option`, `TCPv4ConnectionState`, and `TCPv4ServiceBindingProtocol`.
 
 # uefi-raw - 0.5.1 (2024-03-17)
 

uefi-raw/src/lib.rs

@@ -65,6 +65,16 @@ pub type VirtualAddress = u64;
 #[repr(transparent)]
 pub struct Ipv4Address(pub [u8; 4]);
 
+impl Ipv4Address {
+    pub fn new(b1: u8, b2: u8, b3: u8, b4: u8) -> Self {
+        Self([b1, b2, b3, b4])
+    }
+
+    pub fn zero() -> Self {
+        Self([0, 0, 0, 0])
+    }
+}
+
 /// An IPv6 internet protocol address.
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash)]
 #[repr(transparent)]

uefi-raw/src/protocol/network/ip4.rs

@@ -7,3 +7,44 @@ pub struct Ip4RouteTable {
     pub subnet_mask: Ipv4Address,
     pub gateway_addr: Ipv4Address,
 }
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct Ip4ModeData<'a> {
+    is_started: bool,
+    max_packet_size: u32,
+    config_data: Ip4ConfigData,
+    is_configured: bool,
+    group_count: bool,
+    group_table: &'a [Ipv4Address; 0],
+    route_count: u32,
+    ip4_route_table: &'a [Ip4RouteTable; 0],
+    icmp_type_count: u32,
+    icmp_type_list: &'a [Ip4IcmpType; 0],
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct Ip4ConfigData {
+    default_protocol: u8,
+    accept_any_protocol: bool,
+    accept_icmp_errors: bool,
+    accept_broadcast: bool,
+    accept_promiscuous: bool,
+    use_default_address: bool,
+    station_address: Ipv4Address,
+    subnet_mask: Ipv4Address,
+    type_of_service: u8,
+    time_to_live: u8,
+    do_not_fragment: bool,
+    raw_data: bool,
+    receive_timeout: u32,
+    transmit_timeout: u32,
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct Ip4IcmpType {
+    _type: u8,
+    code: u8,
+}

uefi-raw/src/protocol/network/mod.rs

@@ -3,3 +3,4 @@ pub mod http;
 pub mod ip4;
 pub mod ip4_config2;
 pub mod tls;
+pub mod tcpv4;

uefi-raw/src/protocol/network/tcpv4.rs

@@ -0,0 +1,53 @@
+use crate::{Handle, Status};
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4Option {
+    receive_buffer_size: u32,
+    send_buffer_size: u32,
+    max_syn_back_log: u32,
+    connection_timeout: u32,
+    data_retries: u32,
+    fin_timeout: u32,
+    time_wait_timeout: u32,
+    keep_alive_probes: u32,
+    keep_alive_time: u32,
+    keep_alive_interval: u32,
+    enable_nagle: bool,
+    enable_time_stamp: bool,
+    enable_window_scaling: bool,
+    enable_selective_ack: bool,
+    enable_path_mtu_discovery: bool,
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub enum TCPv4ConnectionState {
+    Closed = 0,
+    Listen = 1,
+    SynSent = 2,
+    SynReceived = 3,
+    Established = 4,
+    FinWait1 = 5,
+    FinWait2 = 6,
+    Closing = 7,
+    TimeWait = 8,
+    CloseWait = 9,
+    LastAck = 10,
+}
+
+#[derive(Debug)]
+#[repr(C)]
+#[unsafe_protocol("00720665-67EB-4a99-BAF7-D3C33A1C7CC9")]
+pub struct TCPv4ServiceBindingProtocol {
+    pub(crate) create_child: extern "efiapi" fn(
+        this: &Self,
+        out_child_handle: &mut Handle,
+    ) -> Status,
+
+    destroy_child: extern "efiapi" fn(
+        this: &Self,
+        child_handle: Handle,
+    ) -> Status,
+}
+

uefi/CHANGELOG.md

@@ -4,6 +4,8 @@
 - Added `Timestamp` protocol.
 - Added `UnalignedSlice::as_ptr`.
 - Added common derives for `Event` and `Handle`.
+- Added `TCPv4Protocol` and implemented a client that can be used to 
+  connect, transmit, and receive data over TCP.
 
 # uefi - 0.27.0 (2024-03-17)
 

uefi/src/proto/network/mod.rs

@@ -4,6 +4,7 @@
 
 pub mod pxe;
 pub mod snp;
+pub mod tcpv4;
 
 /// Represents an IPv4/v6 address.
 ///

uefi/src/proto/network/tcpv4/definitions.rs

@@ -0,0 +1,215 @@
+use alloc::format;
+use core::alloc::Layout;
+use core::ffi::c_void;
+use core::fmt::{Debug, Formatter};
+use core::ptr::copy_nonoverlapping;
+use uefi::{Event, Status};
+use uefi_raw::Ipv4Address;
+use uefi_raw::protocol::network::tcpv4::TCPv4Option;
+
+use crate::proto::network::tcpv4::managed_event::ManagedEvent;
+use crate::proto::network::tcpv4::receive_data::TCPv4ReceiveData;
+use crate::proto::network::tcpv4::transmit_data::TCPv4TransmitData;
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct UnmodelledPointer(pub *mut c_void);
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4AccessPoint {
+    use_default_address: bool,
+    station_address: Ipv4Address,
+    subnet_mask: Ipv4Address,
+    station_port: u16,
+    remote_address: Ipv4Address,
+    remote_port: u16,
+    active_flag: bool,
+}
+
+impl TCPv4AccessPoint {
+    fn new(connection_mode: TCPv4ConnectionMode) -> Self {
+        let (remote_ip, remote_port, is_client) = match connection_mode {
+            TCPv4ConnectionMode::Client(params) => {
+                (params.remote_ip, params.remote_port, true)
+            }
+            TCPv4ConnectionMode::Server => {
+                (Ipv4Address::zero(), 0, false)
+            }
+        };
+        Self {
+            use_default_address: true,
+            // These two fields are meaningless because we set use_default_address above
+            station_address: Ipv4Address::zero(),
+            subnet_mask: Ipv4Address::zero(),
+            // Chosen on-demand
+            station_port: 0,
+            remote_address: remote_ip,
+            remote_port,
+            active_flag: is_client,
+
+        }
+    }
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4ConfigData<'a> {
+    type_of_service: u8,
+    time_to_live: u8,
+    access_point: TCPv4AccessPoint,
+    option: Option<&'a TCPv4Option>,
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4IoToken<'a> {
+    pub completion_token: TCPv4CompletionToken,
+    packet: TCPv4Packet<'a>,
+}
+
+impl<'a> TCPv4IoToken<'a> {
+    pub fn new(
+        event: &ManagedEvent,
+        tx: Option<&'a TCPv4TransmitData>,
+        rx: Option<&'a TCPv4ReceiveData>,
+    ) -> Self {
+        let packet = {
+            if tx.is_some() {
+                TCPv4Packet { tx_data: tx }
+            }
+            else {
+                let rx_ref = rx.as_ref();
+                rx_ref.expect("Either RX or TX data handles must be provided");
+                TCPv4Packet { rx_data: rx }
+            }
+        };
+        Self {
+            completion_token: TCPv4CompletionToken::new(event),
+            packet,
+        }
+    }
+}
+
+impl Drop for TCPv4IoToken<'_> {
+    fn drop(&mut self) {
+        // TODO(PT): I'm unsure offhand whether this empty Drop implementation is important,
+        // or if it can just be... dropped.
+    }
+}
+
+#[derive(Debug)]
+pub struct TCPv4ClientConnectionModeParams {
+    remote_ip: Ipv4Address,
+    remote_port: u16,
+}
+
+impl TCPv4ClientConnectionModeParams {
+    pub fn new(
+        remote_ip: Ipv4Address,
+        remote_port: u16,
+    ) -> Self {
+        Self {
+            remote_ip,
+            remote_port,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum TCPv4ConnectionMode {
+    Client(TCPv4ClientConnectionModeParams),
+    // TODO(PT): There may be parameters we need to model when operating as a server
+    Server,
+}
+
+impl<'a> TCPv4ConfigData<'a> {
+    pub(crate) fn new(
+        connection_mode: TCPv4ConnectionMode,
+        options: Option<&'a TCPv4Option>,
+    ) -> Self {
+        Self {
+            type_of_service: 0,
+            time_to_live: 255,
+            access_point: TCPv4AccessPoint::new(connection_mode),
+            option: options,
+        }
+    }
+}
+
+#[repr(C)]
+union TCPv4Packet<'a> {
+    rx_data: Option<&'a TCPv4ReceiveData>,
+    tx_data: Option<&'a TCPv4TransmitData>,
+}
+
+impl Debug for TCPv4Packet<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
+        unsafe {
+            let rx_data = self.rx_data;
+            let tx_data = self.tx_data;
+            f.write_str(&format!("<TCPv4Packet {rx_data:?} {tx_data:?}"))?;
+        }
+        Ok(())
+    }
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4CompletionToken {
+    pub event: Event,
+    status: Status,
+}
+
+impl TCPv4CompletionToken {
+    pub fn new(event: &ManagedEvent) -> Self {
+        // Safety: The lifetime of this token is bound by the lifetime of the ManagedEvent.
+        let event_clone = unsafe { event.event.unsafe_clone() };
+        Self {
+            event: event_clone,
+            status: Status::SUCCESS,
+        }
+    }
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4FragmentData {
+    pub(crate) fragment_length: u32,
+    pub(crate) fragment_buf: *const c_void,
+}
+
+impl TCPv4FragmentData {
+    pub fn with_buffer_len(len: usize) -> Self {
+        unsafe {
+            let layout = Layout::array::<u8>(len).unwrap();
+            let buffer = alloc::alloc::alloc(layout);
+            Self {
+                fragment_length: len as u32,
+                fragment_buf: buffer as *const c_void,
+            }
+        }
+    }
+    pub fn with_data(data: &[u8]) -> Self {
+        unsafe {
+            let data_len = data.len();
+            let _self = Self::with_buffer_len(data_len);
+            let buffer = _self.fragment_buf as *mut u8;
+            copy_nonoverlapping(
+                data.as_ptr(),
+                buffer,
+                data_len,
+            );
+            _self
+        }
+    }
+}
+
+impl Drop for TCPv4FragmentData {
+    fn drop(&mut self) {
+        unsafe {
+            let layout = Layout::array::<u8>(self.fragment_length as usize).unwrap();
+            alloc::alloc::dealloc(self.fragment_buf as *mut u8, layout);
+        }
+    }
+}

uefi/src/proto/network/tcpv4/managed_event.rs

@@ -0,0 +1,89 @@
+use alloc::boxed::Box;
+use alloc::vec::Vec;
+use core::ffi::c_void;
+use uefi::Event;
+use uefi::prelude::BootServices;
+use uefi::table::boot::{EventType, Tpl};
+use core::ptr::NonNull;
+
+pub struct ManagedEvent<'a> {
+    pub event: Event,
+    boxed_closure: *mut (dyn FnMut(Event) + 'static),
+    boot_services: &'a BootServices,
+}
+
+/// Higher level modelling on top of the thin wrapper that uefi-rs provides.
+/// The wrapper as-is can't be used because the wrapper can be cheaply cloned and passed around,
+/// whereas we need there to be a single instance per event (so the destructor only runs once).
+impl<'a> ManagedEvent<'a> {
+    pub fn new<F>(
+        bs: &'static BootServices,
+        event_type: EventType,
+        callback: F,
+    ) -> Self
+    where
+        F: FnMut(Event) + 'static {
+        let boxed_closure = Box::into_raw(Box::new(callback));
+        unsafe {
+            let event = bs.create_event(
+                event_type,
+                Tpl::CALLBACK,
+                Some(call_closure::<F>),
+                Some(NonNull::new(boxed_closure as *mut _ as *mut c_void).unwrap()),
+            ).expect("Failed to create event");
+            Self {
+                event,
+                boxed_closure,
+                boot_services: bs,
+            }
+        }
+    }
+
+    pub fn wait(&self) {
+        // Safety: The event clone is discarded after being passed to the UEFI function.
+        unsafe {
+            self.boot_services.wait_for_event(
+                &mut [self.event.unsafe_clone()]
+            ).expect("Failed to wait for transmit to complete");
+        }
+    }
+
+    pub fn wait_for_events(bs: &BootServices, events: &[&Self]) -> usize {
+        // Safety: The event clone is discarded after being passed to the UEFI function.
+        unsafe {
+            bs.wait_for_event(
+                &mut events.iter().map(|e| e.event.unsafe_clone()).collect::<Vec<Event>>()
+            ).expect("Failed to wait for transmit to complete")
+        }
+    }
+}
+
+impl Drop for ManagedEvent<'_> {
+    fn drop(&mut self) {
+        unsafe {
+            // Close the UEFI handle
+            // Safety: We're dropping the event here and don't use the handle again after
+            // passing it to the UEFI function.
+            self.boot_services.close_event(self.event.unsafe_clone()).expect("Failed to close event");
+            // *Drop the box* that carries the closure.
+            let _ = Box::from_raw(self.boxed_closure);
+        }
+    }
+}
+
+unsafe extern "efiapi" fn call_closure<F>(
+    event: Event,
+    raw_context: Option<NonNull<c_void>>,
+) where F: FnMut(Event) + 'static {
+    let unwrapped_context = cast_ctx(raw_context);
+    let callback_ptr = unwrapped_context as *mut F;
+    let callback = &mut *callback_ptr;
+    callback(event);
+    // Safety: *Don't drop the box* that carries the closure yet, because
+    // the closure might be invoked again.
+}
+
+unsafe fn cast_ctx<T>(raw_val: Option<NonNull<c_void>>) -> &'static mut T {
+    let val_ptr = raw_val.unwrap().as_ptr() as *mut c_void as *mut T;
+    &mut *val_ptr
+}

uefi/src/proto/network/tcpv4/mod.rs

@@ -0,0 +1,5 @@
+mod definitions;
+mod managed_event;
+mod proto;
+mod receive_data;
+mod transmit_data;

uefi/src/proto/network/tcpv4/proto.rs

@@ -0,0 +1,181 @@
+use alloc::string::{String, ToString};
+use uefi::{Status, StatusExt};
+use uefi::prelude::BootServices;
+use uefi::proto::unsafe_protocol;
+use uefi::Error;
+use uefi::table::boot::EventType;
+use uefi_raw::Ipv4Address;
+use uefi_raw::protocol::network::ip4::Ip4ModeData;
+use uefi_raw::protocol::network::tcpv4::TCPv4ConnectionState;
+use crate::proto::network::tcpv4::definitions::{TCPv4CompletionToken, TCPv4ConfigData, TCPv4ConnectionMode, TCPv4IoToken, UnmodelledPointer};
+use crate::proto::network::tcpv4::managed_event::ManagedEvent;
+use crate::proto::network::tcpv4::transmit_data::TCPv4TransmitDataHandle;
+
+#[derive(Debug)]
+#[repr(C)]
+#[unsafe_protocol("65530BC7-A359-410F-B010-5AADC7EC2B62")]
+pub struct TCPv4Protocol {
+    get_mode_data_fn: extern "efiapi" fn(
+        this: &Self,
+        out_connection_state: Option<&mut TCPv4ConnectionState>,
+        out_config_data: Option<&mut UnmodelledPointer>,
+        out_ip4_mode_data: Option<&mut Ip4ModeData>,
+        out_managed_network_config_data: Option<&mut UnmodelledPointer>,
+        out_simple_network_mode: Option<&mut UnmodelledPointer>,
+    ) -> Status,
+
+    configure_fn: extern "efiapi" fn(
+        this: &Self,
+        config_data: Option<&TCPv4ConfigData>,
+    ) -> Status,
+
+    routes_fn: extern "efiapi" fn(
+        this: &Self,
+        delete_route: bool,
+        subnet_address: &Ipv4Address,
+        subnet_mask: &Ipv4Address,
+        gateway_address: &Ipv4Address,
+    ) -> Status,
+
+    connect_fn: extern "efiapi" fn(
+        this: &Self,
+        connection_token: &TCPv4CompletionToken,
+    ) -> Status,
+
+    accept_fn: extern "efiapi" fn(
+        this: &Self,
+        listen_token: &UnmodelledPointer,
+    ) -> Status,
+
+    pub(crate) transmit_fn: extern "efiapi" fn(
+        this: &Self,
+        token: &TCPv4IoToken,
+    ) -> Status,
+
+    pub receive_fn: extern "efiapi" fn(
+        this: &Self,
+        token: &TCPv4IoToken,
+    ) -> Status,
+
+    close_fn: extern "efiapi" fn(
+        this: &Self,
+        close_token: &UnmodelledPointer,
+    ) -> Status,
+
+    cancel_fn: extern "efiapi" fn(
+        this: &Self,
+        completion_token: &UnmodelledPointer,
+    ) -> Status,
+
+    poll_fn: extern "efiapi" fn(this: &Self) -> Status,
+}
+
+impl TCPv4Protocol {
+    pub fn reset_stack(&self) {
+        // The UEFI specification states that configuring with NULL options "brutally resets" the TCP stack
+        (self.configure_fn)(
+            self,
+            None,
+        ).to_result().expect("Failed to reset TCP stack")
+    }
+
+    pub fn configure(
+        &self,
+        bt: &BootServices,
+        connection_mode: TCPv4ConnectionMode,
+    ) -> uefi::Result<(), String> {
+        let configuration = TCPv4ConfigData::new(connection_mode, None);
+        // Maximum timeout of 10 seconds
+        for _ in 0..10 {
+            let result = (self.configure_fn)(
+                self,
+                Some(&configuration),
+            );
+            if result == Status::SUCCESS {
+                // Configured connection
+                return Ok(())
+            }
+            else if result == Status::NO_MAPPING {
+                // DHCP is still running, wait...
+                bt.stall(1_000_000);
+            }
+            else {
+                // Error, spin and try again
+                bt.stall(1_000_000);
+            }
+        }
+        Err(Error::new(Status::PROTOCOL_ERROR, "Timeout before configuring the connection succeeded.".to_string()))
+    }
+
+    pub fn get_tcp_connection_state(&self) -> TCPv4ConnectionState {
+        let mut connection_state = core::mem::MaybeUninit::<TCPv4ConnectionState>::uninit();
+        let connection_state_ptr = connection_state.as_mut_ptr();
+        unsafe {
+            (self.get_mode_data_fn)(
+                self,
+                Some(&mut *connection_state_ptr),
+                None,
+                None,
+                None,
+                None,
+            ).to_result().expect("Failed to read connection state");
+            connection_state.assume_init()
+        }
+    }
+
+    pub fn get_ipv4_mode_data(&self) -> Ip4ModeData {
+        let mut mode_data = core::mem::MaybeUninit::<Ip4ModeData>::uninit();
+        let mode_data_ptr = mode_data.as_mut_ptr();
+        unsafe {
+            (self.get_mode_data_fn)(
+                self,
+                None,
+                None,
+                Some(&mut *mode_data_ptr),
+                None,
+                None,
+            ).to_result().expect("Failed to read mode data");
+            mode_data.assume_init()
+        }
+    }
+
+    pub fn connect(
+        &mut self,
+        bs: &'static BootServices,
+    ) {
+        let event = ManagedEvent::new(
+            bs,
+            EventType::NOTIFY_WAIT,
+            |_| {},
+        );
+        let completion_token = TCPv4CompletionToken::new(&event);
+        (self.connect_fn)(
+            &self,
+            &completion_token,
+        ).to_result().expect("Failed to call Connect()");
+        event.wait();
+    }
+
+    pub fn transmit(
+        &mut self,
+        bs: &'static BootServices,
+        data: &[u8],
+    ) {
+        let event = ManagedEvent::new(
+            bs,
+            EventType::NOTIFY_WAIT,
+            move |_| {
+                // TODO(PT): Accept a user-provided closure?
+            },
+        );
+
+        let tx_data_handle = TCPv4TransmitDataHandle::new(data);
+        let tx_data = tx_data_handle.get_data_ref();
+        let io_token = TCPv4IoToken::new(&event, Some(&tx_data), None);
+        (self.transmit_fn)(
+            &self,
+            &io_token,
+        ).to_result().expect("Failed to transmit");
+        event.wait();
+    }
+}

uefi/src/proto/network/tcpv4/receive_data.rs

@@ -0,0 +1,112 @@
+use alloc::vec;
+use alloc::vec::Vec;
+use core::alloc::Layout;
+use core::marker::PhantomData;
+use core::mem;
+use core::mem::ManuallyDrop;
+use core::ptr::copy_nonoverlapping;
+use crate::proto::network::tcpv4::definitions::TCPv4FragmentData;
+use crate::tcpv4::TCPv4FragmentData;
+
+/// This type is necessary because the underlying structure has a flexible array member.
+/// Due to this, the memory for the instance needs to be carefully managed.
+/// A Box cannot be used because the Box doesn't have the full knowledge of the layout.
+/// A wide pointer also cannot be used because the layout needs to be precisely controlled for FFI.
+/// Therefore, we use a wrapper 'handle' to manage the lifecycle of the allocation manually.
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4ReceiveDataHandle<'a> {
+    ptr: *const TCPv4ReceiveData,
+    layout: Layout,
+    phantom: PhantomData<&'a ()>,
+}
+
+impl<'a> TCPv4ReceiveDataHandle<'a> {
+    fn total_layout_size(fragment_count: usize) -> usize {
+        let size_of_fragments = mem::size_of::<ManuallyDrop<TCPv4FragmentData>>() * fragment_count;
+        mem::size_of::<Self>() + size_of_fragments
+    }
+
+    pub(crate) fn new() -> Self {
+        let buffer_len = 2048*16;
+        let fragment = ManuallyDrop::new(TCPv4FragmentData::with_buffer_len(buffer_len));
+        let layout = Layout::from_size_align(
+            Self::total_layout_size(1),
+            mem::align_of::<Self>(),
+        ).unwrap();
+        unsafe {
+            let ptr = alloc::alloc::alloc(layout) as *mut TCPv4ReceiveData;
+            (*ptr).urgent = false;
+            (*ptr).data_length = buffer_len as _;
+
+            let fragment_count = 1;
+            (*ptr).fragment_count = fragment_count as _;
+            copy_nonoverlapping(
+                &fragment as *const _,
+                (*ptr).fragment_table.as_mut_ptr(),
+                fragment_count,
+            );
+
+            Self {
+                ptr: ptr as _,
+                layout,
+                phantom: PhantomData,
+            }
+        }
+    }
+
+    pub(crate) fn get_data_ref(&self) -> &'a TCPv4ReceiveData {
+        // Safety: The reference is strictly tied to the lifetime of this handle
+        unsafe { &*self.ptr }
+    }
+}
+
+impl Drop for TCPv4ReceiveDataHandle<'_> {
+    fn drop(&mut self) {
+        let ptr = self.ptr as *mut TCPv4ReceiveData;
+        unsafe {
+            // First, drop all the fragments
+            let fragment_table: *mut ManuallyDrop<TCPv4FragmentData> = (*ptr).fragment_table.as_mut_ptr();
+            for i in 0..((*ptr).fragment_count as usize) {
+                let fragment_ptr = fragment_table.add(i as _);
+                ManuallyDrop::drop(&mut *fragment_ptr);
+            }
+
+            // Lastly, drop the allocation itself
+            alloc::alloc::dealloc(ptr as *mut u8, self.layout);
+        }
+    }
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4ReceiveData {
+    urgent: bool,
+    data_length: u32,
+    fragment_count: u32,
+    fragment_table: [ManuallyDrop<TCPv4FragmentData>; 0],
+}
+
+impl TCPv4ReceiveData {
+    pub fn read_buffers(&self) -> Vec<u8> {
+        let mut out = vec![];
+        unsafe {
+            let ptr = self as *const Self;
+            let fragment_table: *const ManuallyDrop<TCPv4FragmentData> = (*ptr).fragment_table.as_ptr();
+            for i in 0..(self.fragment_count as usize) {
+                let fragment_ptr = fragment_table.add(i as _);
+                let fragment = &*fragment_ptr;
+                let fragment_buf = fragment.fragment_buf as *const u8;
+                let fragment_slice = core::slice::from_raw_parts(fragment_buf, self.data_length as _);
+                out.extend_from_slice(fragment_slice);
+            }
+        }
+        out
+    }
+}
+
+impl Drop for TCPv4ReceiveData {
+    fn drop(&mut self) {
+        panic!("Should be manually dropped by TCPv4ReceiveDataHandle")
+    }
+}

uefi/src/proto/network/tcpv4/transmit_data.rs

@@ -0,0 +1,84 @@
+use core::alloc::Layout;
+use core::mem;
+use core::mem::ManuallyDrop;
+use core::ptr::copy_nonoverlapping;
+use crate::proto::network::tcpv4::definitions::TCPv4FragmentData;
+
+/// This type is necessary because the underlying structure has a flexible array member.
+/// Due to this, the memory for the instance needs to be carefully managed.
+/// A Box cannot be used because the Box doesn't have the full knowledge of the layout.
+/// A wide pointer also cannot be used because the layout needs to be precisely controlled for FFI.
+/// Therefore, we use a wrapper 'handle' to manage the lifecycle of the allocation manually.
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4TransmitDataHandle {
+    ptr: *const TCPv4TransmitData,
+    layout: Layout,
+}
+
+impl TCPv4TransmitDataHandle {
+    fn total_layout_size(fragment_count: usize) -> usize {
+        let size_of_fragments = mem::size_of::<ManuallyDrop<TCPv4FragmentData>>() * fragment_count;
+        mem::size_of::<Self>() + size_of_fragments
+    }
+
+    pub(crate) fn new(data: &[u8]) -> Self {
+        let fragment = ManuallyDrop::new(TCPv4FragmentData::with_data(data));
+        let layout = Layout::from_size_align(
+            Self::total_layout_size(1),
+            mem::align_of::<Self>(),
+        ).unwrap();
+        unsafe {
+            let ptr = alloc::alloc::alloc(layout) as *mut TCPv4TransmitData;
+            (*ptr).push = true;
+            (*ptr).urgent = false;
+            (*ptr).data_length = data.len() as _;
+
+            let fragment_count = 1;
+            (*ptr).fragment_count = fragment_count as _;
+            copy_nonoverlapping(
+                &fragment as *const _,
+                (*ptr).fragment_table.as_mut_ptr(),
+                fragment_count,
+            );
+
+            Self {
+                ptr: ptr as _,
+                layout,
+            }
+        }
+    }
+
+    pub(crate) fn get_data_ref(&self) -> &TCPv4TransmitData {
+        // Safety: The reference is strictly tied to the lifetime of this handle
+        unsafe { &*self.ptr }
+    }
+}
+
+impl Drop for TCPv4TransmitDataHandle {
+    fn drop(&mut self) {
+        unsafe {
+            let ptr = self.ptr as *mut TCPv4TransmitData;
+
+            // First, drop all the fragments
+            let fragment_table: *mut ManuallyDrop<TCPv4FragmentData> = (*ptr).fragment_table.as_mut_ptr();
+            for i in 0..((*ptr).fragment_count as usize) {
+                let fragment_ptr = fragment_table.add(i as _);
+                ManuallyDrop::drop(&mut *fragment_ptr);
+            }
+
+            // Lastly, drop the allocation itself
+            alloc::alloc::dealloc(ptr as *mut u8, self.layout);
+        }
+    }
+}
+
+#[derive(Debug)]
+#[repr(C)]
+pub struct TCPv4TransmitData {
+    push: bool,
+    urgent: bool,
+    data_length: u32,
+    fragment_count: u32,
+    fragment_table: [ManuallyDrop<TCPv4FragmentData>; 0],
+}