refactor(s2n-quic-platform): simplify Message trait

quic/s2n-quic-core/src/inet/datagram.rs

@@ -32,4 +32,6 @@ pub struct AncillaryData {
     /// Correctly threading this value through to connections ensures packets end up on the same
     /// network interfaces and thereby have consistent MAC addresses.
     pub local_interface: Option<u32>,
+    /// Set when the packet buffer is an aggregate of multiple received packets
+    pub segment_size: u16,
 }

quic/s2n-quic-platform/src/message.rs

@@ -1,49 +1,36 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-#[macro_use]
-mod macros;
+use core::ffi::c_void;
+use s2n_quic_core::{inet::datagram, io::tx, path};
 
+#[cfg(any(s2n_quic_platform_socket_msg, s2n_quic_platform_socket_mmsg))]
+pub mod cmsg;
 #[cfg(s2n_quic_platform_socket_mmsg)]
 pub mod mmsg;
-
 #[cfg(s2n_quic_platform_socket_msg)]
 pub mod msg;
-
-#[cfg(any(s2n_quic_platform_socket_msg, s2n_quic_platform_socket_mmsg))]
-pub mod cmsg;
-
 pub mod queue;
 pub mod simple;
 
-use core::ffi::c_void;
-use s2n_quic_core::{
-    inet::{datagram, ExplicitCongestionNotification, SocketAddress},
-    io::tx,
-    path,
-};
+pub mod default {
+    cfg_if::cfg_if! {
+        if #[cfg(s2n_quic_platform_socket_mmsg)] {
+            pub use super::mmsg::*;
+        } else if #[cfg(s2n_quic_platform_socket_msg)] {
+            pub use super::msg::*;
+        } else {
+            pub use super::simple::*;
+        }
+    }
+}
 
 /// An abstract message that can be sent and received on a network
 pub trait Message {
     type Handle: path::Handle;
 
     const SUPPORTS_GSO: bool;
 
-    /// Returns the ECN values for the message
-    fn ecn(&self) -> ExplicitCongestionNotification;
-
-    /// Sets the ECN values for the message
-    fn set_ecn(&mut self, ecn: ExplicitCongestionNotification, remote_address: &SocketAddress);
-
-    /// Returns the `SocketAddress` for the message
-    fn remote_address(&self) -> Option<SocketAddress>;
-
-    /// Sets the `SocketAddress` for the message
-    fn set_remote_address(&mut self, remote_address: &SocketAddress);
-
-    /// Returns the path handle for the message
-    fn path_handle(&self) -> Option<Self::Handle>;
-
     /// Returns the length of the payload
     fn payload_len(&self) -> usize;
 
@@ -63,18 +50,11 @@ pub trait Message {
     /// This should used in scenarios where the data pointers are the same.
     fn replicate_fields_from(&mut self, other: &Self);
 
-    /// Returns a pointer for the message payload
-    fn payload_ptr(&self) -> *const u8;
-
     /// Returns a mutable pointer for the message payload
     fn payload_ptr_mut(&mut self) -> *mut u8;
 
-    /// Returns a slice for the message payload
-    fn payload(&self) -> &[u8] {
-        unsafe { core::slice::from_raw_parts(self.payload_ptr(), self.payload_len()) }
-    }
-
     /// Returns a mutable slice for the message payload
+    #[inline]
     fn payload_mut(&mut self) -> &mut [u8] {
         unsafe { core::slice::from_raw_parts_mut(self.payload_ptr_mut(), self.payload_len()) }
     }
@@ -101,10 +81,7 @@ pub trait Message {
     }
 
     /// Reads the message as an RX packet
-    fn rx_read(
-        &mut self,
-        local_address: &path::LocalAddress,
-    ) -> Option<(datagram::Header<Self::Handle>, &mut [u8])>;
+    fn rx_read(&mut self, local_address: &path::LocalAddress) -> Option<RxMessage<Self::Handle>>;
 
     /// Writes the message into the TX packet
     fn tx_write<M: tx::Message<Handle = Self::Handle>>(
@@ -113,6 +90,26 @@ pub trait Message {
     ) -> Result<usize, tx::Error>;
 }
 
+pub struct RxMessage<'a, Handle: Copy> {
+    /// The received header for the message
+    pub header: datagram::Header<Handle>,
+    /// The number of segments inside the message
+    pub segment_size: usize,
+    /// The full payload of the message
+    pub payload: &'a mut [u8],
+}
+
+impl<'a, Handle: Copy> RxMessage<'a, Handle> {
+    #[inline]
+    pub fn for_each<F: FnMut(datagram::Header<Handle>, &mut [u8])>(self, mut on_packet: F) {
+        debug_assert_ne!(self.segment_size, 0);
+
+        for segment in self.payload.chunks_mut(self.segment_size) {
+            on_packet(self.header, segment);
+        }
+    }
+}
+
 /// A message ring used to back a queue
 pub trait Ring {
     /// The type of message that is stored in the ring

quic/s2n-quic-platform/src/message/mmsg.rs

@@ -6,71 +6,18 @@ use crate::message::{
     Message as MessageTrait,
 };
 use alloc::vec::Vec;
-use core::{fmt, mem::zeroed};
+use core::mem::zeroed;
 use libc::mmsghdr;
-use s2n_quic_core::{
-    inet::{datagram, ExplicitCongestionNotification, SocketAddress},
-    io::tx,
-    path,
-};
-
-#[repr(transparent)]
-pub struct Message(pub(crate) mmsghdr);
+use s2n_quic_core::{io::tx, path};
 
+pub use libc::mmsghdr as Message;
 pub type Handle = msg::Handle;
 
-impl_message_delegate!(Message, 0, mmsghdr);
-
-impl fmt::Debug for Message {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        let alt = f.alternate();
-        let mut s = f.debug_struct("mmsghdr");
-
-        s.field("remote_address", &self.remote_address()).field(
-            "ancillary_data",
-            &crate::message::cmsg::decode(&self.0.msg_hdr),
-        );
-
-        if alt {
-            s.field("payload", &self.payload());
-        } else {
-            s.field("payload_len", &self.payload_len());
-        }
-
-        s.finish()
-    }
-}
-
 impl MessageTrait for mmsghdr {
     type Handle = Handle;
 
     const SUPPORTS_GSO: bool = libc::msghdr::SUPPORTS_GSO;
 
-    #[inline]
-    fn ecn(&self) -> ExplicitCongestionNotification {
-        self.msg_hdr.ecn()
-    }
-
-    #[inline]
-    fn set_ecn(&mut self, ecn: ExplicitCongestionNotification, remote_address: &SocketAddress) {
-        self.msg_hdr.set_ecn(ecn, remote_address)
-    }
-
-    #[inline]
-    fn remote_address(&self) -> Option<SocketAddress> {
-        self.msg_hdr.remote_address()
-    }
-
-    #[inline]
-    fn set_remote_address(&mut self, remote_address: &SocketAddress) {
-        self.msg_hdr.set_remote_address(remote_address)
-    }
-
-    #[inline]
-    fn path_handle(&self) -> Option<Self::Handle> {
-        self.msg_hdr.path_handle()
-    }
-
     #[inline]
     fn payload_len(&self) -> usize {
         self.msg_len as usize
@@ -99,11 +46,6 @@ impl MessageTrait for mmsghdr {
         self.msg_hdr.reset(mtu)
     }
 
-    #[inline]
-    fn payload_ptr(&self) -> *const u8 {
-        self.msg_hdr.payload_ptr()
-    }
-
     #[inline]
     fn payload_ptr_mut(&mut self) -> *mut u8 {
         self.msg_hdr.payload_ptr_mut()
@@ -119,7 +61,7 @@ impl MessageTrait for mmsghdr {
     fn rx_read(
         &mut self,
         local_address: &path::LocalAddress,
-    ) -> Option<(datagram::Header<Self::Handle>, &mut [u8])> {
+    ) -> Option<super::RxMessage<Self::Handle>> {
         unsafe {
             // We need to replicate the `msg_len` field to the inner type before delegating
             // Safety: The `msg_len` is associated with the same buffer as the `msg_hdr`
@@ -172,9 +114,9 @@ impl<Payloads: crate::buffer::Buffer> Ring<Payloads> {
             .map(|msg_hdr| unsafe {
                 let mut mmsghdr = zeroed::<mmsghdr>();
                 let payload_len = msg_hdr.payload_len();
-                mmsghdr.msg_hdr = msg_hdr.0;
+                mmsghdr.msg_hdr = msg_hdr;
                 mmsghdr.set_payload_len(payload_len);
-                Message(mmsghdr)
+                mmsghdr
             })
             .collect();