From 8cbffc0e8e01b9aee337b746e41018ec65d14d8f Mon Sep 17 00:00:00 2001 From: Phillip Tennen Date: Sun, 14 Apr 2024 16:35:47 +0100 Subject: [PATCH 1/4] uefi-raw: Add utility methods to Ipv4Address --- uefi-raw/CHANGELOG.md | 1 + uefi-raw/src/lib.rs | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/uefi-raw/CHANGELOG.md b/uefi-raw/CHANGELOG.md index d0315e37e..960a737a0 100644 --- a/uefi-raw/CHANGELOG.md +++ b/uefi-raw/CHANGELOG.md @@ -3,6 +3,7 @@ ## Added - Added `TimestampProtocol`. - Added `DevicePathToTextProtocol` and `DevicePathFromTextProtocol`. +- Added minor utility methods to `Ipv4Address`: `new(u8, u8, u8, u8) -> Self` and `zero() -> Self`. # uefi-raw - 0.5.1 (2024-03-17) diff --git a/uefi-raw/src/lib.rs b/uefi-raw/src/lib.rs index 003dd9584..2d5f061dd 100644 --- a/uefi-raw/src/lib.rs +++ b/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)] From d18c954281650f3b109db87844ee37cfdfdb1821 Mon Sep 17 00:00:00 2001 From: Phillip Tennen Date: Sun, 14 Apr 2024 16:36:03 +0100 Subject: [PATCH 2/4] uefi-raw: Model more IPv4 types --- uefi-raw/CHANGELOG.md | 1 + uefi-raw/src/protocol/network/ip4.rs | 41 ++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/uefi-raw/CHANGELOG.md b/uefi-raw/CHANGELOG.md index 960a737a0..87ffb4a0e 100644 --- a/uefi-raw/CHANGELOG.md +++ b/uefi-raw/CHANGELOG.md @@ -4,6 +4,7 @@ - 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`. # uefi-raw - 0.5.1 (2024-03-17) diff --git a/uefi-raw/src/protocol/network/ip4.rs b/uefi-raw/src/protocol/network/ip4.rs index 30c3f8644..4121f046a 100644 --- a/uefi-raw/src/protocol/network/ip4.rs +++ b/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, +} From 7de5e79dda98a8463f9ed6fe02b770d3d16a70d1 Mon Sep 17 00:00:00 2001 From: Phillip Tennen Date: Sun, 14 Apr 2024 16:36:17 +0100 Subject: [PATCH 3/4] uefi-raw: Model some TCP protocol types --- uefi-raw/CHANGELOG.md | 1 + uefi-raw/src/protocol/network/mod.rs | 1 + uefi-raw/src/protocol/network/tcpv4.rs | 53 ++++++++++++++++++++++++++ 3 files changed, 55 insertions(+) create mode 100644 uefi-raw/src/protocol/network/tcpv4.rs diff --git a/uefi-raw/CHANGELOG.md b/uefi-raw/CHANGELOG.md index 87ffb4a0e..d038136db 100644 --- a/uefi-raw/CHANGELOG.md +++ b/uefi-raw/CHANGELOG.md @@ -5,6 +5,7 @@ - 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) diff --git a/uefi-raw/src/protocol/network/mod.rs b/uefi-raw/src/protocol/network/mod.rs index 71d577d4e..709ff8493 100644 --- a/uefi-raw/src/protocol/network/mod.rs +++ b/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; diff --git a/uefi-raw/src/protocol/network/tcpv4.rs b/uefi-raw/src/protocol/network/tcpv4.rs new file mode 100644 index 000000000..7b5897832 --- /dev/null +++ b/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, +} + From 4e2fcdcf4c92b24e35f2722130a898c5ee0706ce Mon Sep 17 00:00:00 2001 From: Phillip Tennen Date: Sun, 14 Apr 2024 16:36:48 +0100 Subject: [PATCH 4/4] uefi: Port TCP protocol client implementation --- uefi/CHANGELOG.md | 2 + uefi/src/proto/network/mod.rs | 1 + uefi/src/proto/network/tcpv4/definitions.rs | 215 ++++++++++++++++++ uefi/src/proto/network/tcpv4/managed_event.rs | 89 ++++++++ uefi/src/proto/network/tcpv4/mod.rs | 5 + uefi/src/proto/network/tcpv4/proto.rs | 181 +++++++++++++++ uefi/src/proto/network/tcpv4/receive_data.rs | 112 +++++++++ uefi/src/proto/network/tcpv4/transmit_data.rs | 84 +++++++ 8 files changed, 689 insertions(+) create mode 100644 uefi/src/proto/network/tcpv4/definitions.rs create mode 100644 uefi/src/proto/network/tcpv4/managed_event.rs create mode 100644 uefi/src/proto/network/tcpv4/mod.rs create mode 100644 uefi/src/proto/network/tcpv4/proto.rs create mode 100644 uefi/src/proto/network/tcpv4/receive_data.rs create mode 100644 uefi/src/proto/network/tcpv4/transmit_data.rs diff --git a/uefi/CHANGELOG.md b/uefi/CHANGELOG.md index 2e68cd714..6329698ea 100644 --- a/uefi/CHANGELOG.md +++ b/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) diff --git a/uefi/src/proto/network/mod.rs b/uefi/src/proto/network/mod.rs index 0ef7f04bd..0570ac266 100644 --- a/uefi/src/proto/network/mod.rs +++ b/uefi/src/proto/network/mod.rs @@ -4,6 +4,7 @@ pub mod pxe; pub mod snp; +pub mod tcpv4; /// Represents an IPv4/v6 address. /// diff --git a/uefi/src/proto/network/tcpv4/definitions.rs b/uefi/src/proto/network/tcpv4/definitions.rs new file mode 100644 index 000000000..5c79f4cdb --- /dev/null +++ b/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!(" 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::(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::(self.fragment_length as usize).unwrap(); + alloc::alloc::dealloc(self.fragment_buf as *mut u8, layout); + } + } +} diff --git a/uefi/src/proto/network/tcpv4/managed_event.rs b/uefi/src/proto/network/tcpv4/managed_event.rs new file mode 100644 index 000000000..c8d8116c0 --- /dev/null +++ b/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( + 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::), + 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::>() + ).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( + event: Event, + raw_context: Option>, +) 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(raw_val: Option>) -> &'static mut T { + let val_ptr = raw_val.unwrap().as_ptr() as *mut c_void as *mut T; + &mut *val_ptr +} diff --git a/uefi/src/proto/network/tcpv4/mod.rs b/uefi/src/proto/network/tcpv4/mod.rs new file mode 100644 index 000000000..c11898b53 --- /dev/null +++ b/uefi/src/proto/network/tcpv4/mod.rs @@ -0,0 +1,5 @@ +mod definitions; +mod managed_event; +mod proto; +mod receive_data; +mod transmit_data; diff --git a/uefi/src/proto/network/tcpv4/proto.rs b/uefi/src/proto/network/tcpv4/proto.rs new file mode 100644 index 000000000..5067e38b2 --- /dev/null +++ b/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::::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::::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(); + } +} diff --git a/uefi/src/proto/network/tcpv4/receive_data.rs b/uefi/src/proto/network/tcpv4/receive_data.rs new file mode 100644 index 000000000..0fe2ad295 --- /dev/null +++ b/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::>() * fragment_count; + mem::size_of::() + 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::(), + ).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 = (*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; 0], +} + +impl TCPv4ReceiveData { + pub fn read_buffers(&self) -> Vec { + let mut out = vec![]; + unsafe { + let ptr = self as *const Self; + let fragment_table: *const ManuallyDrop = (*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") + } +} diff --git a/uefi/src/proto/network/tcpv4/transmit_data.rs b/uefi/src/proto/network/tcpv4/transmit_data.rs new file mode 100644 index 000000000..94cafbd09 --- /dev/null +++ b/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::>() * fragment_count; + mem::size_of::() + 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::(), + ).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 = (*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; 0], +}