platform: add GSO support

quic/s2n-quic-core/src/io/tx.rs

@@ -85,6 +85,9 @@ pub trait Message {
     /// Returns the IPv6 flow label for the message
     fn ipv6_flow_label(&mut self) -> u32;
 
+    /// Returns true if the packet can be used in a GSO packet
+    fn can_gso(&self) -> bool;
+
     /// Writes the payload of the message to an output buffer
     fn write_payload(&mut self, buffer: &mut [u8]) -> usize;
 }
@@ -106,6 +109,10 @@ impl<Payload: AsRef<[u8]>> Message for (SocketAddress, Payload) {
         0
     }
 
+    fn can_gso(&self) -> bool {
+        true
+    }
+
     fn write_payload(&mut self, buffer: &mut [u8]) -> usize {
         let payload = self.1.as_ref();
         let len = payload.len();

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

@@ -1,16 +1,14 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-mod segment;
+use core::ops::{Deref, DerefMut};
+
 mod vec;
 // TODO support mmap buffers
 
-pub use segment::*;
 pub use vec::*;
 
-use core::ops::{Index, IndexMut};
-
-pub trait Buffer: Index<usize, Output = [u8]> + IndexMut<usize> {
+pub trait Buffer: Deref<Target = [u8]> + DerefMut {
     fn len(&self) -> usize;
 
     fn is_empty(&self) -> bool {

quic/s2n-quic-platform/src/buffer/vec.rs

@@ -1,24 +1,28 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::buffer::{Buffer, SegmentBuffer};
+use crate::buffer::Buffer;
 use core::{
     fmt,
-    ops::{Index, IndexMut},
+    ops::{Deref, DerefMut},
 };
 use s2n_quic_core::path::DEFAULT_MAX_MTU;
 
 // TODO decide on better defaults
-const DEFAULT_MESSAGE_COUNT: usize = 4096;
+const DEFAULT_MESSAGE_COUNT: usize = 1024;
 
-pub struct VecBuffer(SegmentBuffer<alloc::vec::Vec<u8>>);
+pub struct VecBuffer {
+    region: alloc::vec::Vec<u8>,
+    mtu: usize,
+}
 
 impl VecBuffer {
     /// Create a contiguous buffer with the specified number of messages
     pub fn new(message_count: usize, mtu: usize) -> Self {
         let len = message_count * mtu;
-        let vec = alloc::vec![0; len];
-        Self(SegmentBuffer::new(vec, mtu))
+        let region = alloc::vec![0; len];
+
+        Self { region, mtu }
     }
 }
 
@@ -45,24 +49,24 @@ impl fmt::Debug for VecBuffer {
 
 impl Buffer for VecBuffer {
     fn len(&self) -> usize {
-        self.0.len()
+        self.region.len()
     }
 
     fn mtu(&self) -> usize {
-        self.0.mtu()
+        self.mtu
     }
 }
 
-impl Index<usize> for VecBuffer {
-    type Output = [u8];
+impl Deref for VecBuffer {
+    type Target = [u8];
 
-    fn index(&self, index: usize) -> &Self::Output {
-        self.0.index(index)
+    fn deref(&self) -> &[u8] {
+        self.region.as_ref()
     }
 }
 
-impl IndexMut<usize> for VecBuffer {
-    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
-        self.0.index_mut(index)
+impl DerefMut for VecBuffer {
+    fn deref_mut(&mut self) -> &mut [u8] {
+        self.region.as_mut()
     }
 }

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

@@ -0,0 +1,20 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use lazy_static::lazy_static;
+
+mod gso;
+pub use gso::Gso;
+
+lazy_static! {
+    static ref FEATURES: Features = Features::default();
+}
+
+pub fn get() -> &'static Features {
+    &*FEATURES
+}
+
+#[derive(Debug, Default)]
+pub struct Features {
+    pub gso: Gso,
+}

quic/s2n-quic-platform/src/features/gso.rs

@@ -0,0 +1,83 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use core::num::NonZeroUsize;
+
+#[derive(Debug)]
+pub struct Gso {
+    max_segments: NonZeroUsize,
+}
+
+impl Default for Gso {
+    fn default() -> Self {
+        let max_segments = if cfg!(target_os = "linux") {
+            // https://github.com/torvalds/linux/blob/e9f1cbc0c4114880090c7a578117d3b9cf184ad4/tools/testing/selftests/net/udpgso.c#L37
+            // ```
+            // #define UDP_MAX_SEGMENTS	(1 << 6UL)
+            // ```
+            1 << 6
+        } else {
+            1
+        };
+
+        let max_segments = NonZeroUsize::new(max_segments).unwrap();
+
+        Self { max_segments }
+    }
+}
+
+#[cfg(target_os = "linux")]
+impl Gso {
+    pub fn set(cmsg: &mut [u8], segment_size: usize) -> usize {
+        use core::mem::size_of;
+        type SegmentType = u16;
+
+        let len = unsafe { libc::CMSG_SPACE(size_of::<SegmentType>() as _) } as usize;
+        debug_assert_ne!(len, 0);
+        assert!(
+            cmsg.len() >= len,
+            "out of space in cmsg: needed {}, got {}",
+            len,
+            cmsg.len()
+        );
+
+        let cmsg = unsafe { &mut *(&mut cmsg[0] as *mut u8 as *mut libc::cmsghdr) };
+        cmsg.cmsg_level = libc::SOL_UDP;
+        cmsg.cmsg_type = libc::UDP_SEGMENT;
+        cmsg.cmsg_len = unsafe { libc::CMSG_LEN(size_of::<SegmentType>() as _) } as _;
+        unsafe {
+            core::ptr::write(
+                libc::CMSG_DATA(cmsg) as *const _ as *mut _,
+                segment_size as SegmentType,
+            );
+        }
+
+        len
+    }
+
+    pub fn max_segments(&self) -> usize {
+        self.max_segments.get()
+    }
+
+    pub fn default_max_segments(&self) -> usize {
+        // TODO profile a good default
+        const DEFAULT_MAX_SEGMENTS: usize = 16;
+
+        self.max_segments().min(DEFAULT_MAX_SEGMENTS)
+    }
+}
+
+#[cfg(not(target_os = "linux"))]
+impl Gso {
+    pub fn set(_cmsg: &mut [u8], _segment_size: usize) -> usize {
+        panic!("cannot use GSO on the current platform")
+    }
+
+    pub fn max_segments(&self) -> usize {
+        1
+    }
+
+    pub fn default_max_segments(&self) -> usize {
+        1
+    }
+}

quic/s2n-quic-platform/src/io/tokio.rs

@@ -664,8 +664,11 @@ mod tests {
             let len = entries.len();
             for entry in entries {
                 let payload: &[u8] = entry.payload_mut();
-                let payload = payload.try_into().unwrap();
-                let id = u32::from_be_bytes(payload);
+                if payload.len() != 4 {
+                    panic!("invalid payload {:?}", payload);
+                }
+                let id = payload.try_into().unwrap();
+                let id = u32::from_be_bytes(id);
                 self.messages.remove(&id);
             }
             queue.finish(len);

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

@@ -12,6 +12,7 @@ extern crate alloc;
 mod macros;
 
 pub mod buffer;
+pub mod features;
 pub mod io;
 pub mod message;
 pub mod socket;

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

@@ -18,6 +18,8 @@ use s2n_quic_core::inet::{ExplicitCongestionNotification, SocketAddress};
 
 /// An abstract message that can be sent and received on a network
 pub trait Message {
+    const SUPPORTS_GSO: bool;
+
     /// Returns the ECN values for the message
     fn ecn(&self) -> ExplicitCongestionNotification;
 
@@ -30,9 +32,6 @@ pub trait Message {
     /// Sets the `SocketAddress` for the message
     fn set_remote_address(&mut self, remote_address: &SocketAddress);
 
-    /// Resets the `SocketAddress` for the message
-    fn reset_remote_address(&mut self);
-
     /// Returns the length of the payload
     fn payload_len(&self) -> usize;
 
@@ -65,6 +64,15 @@ pub trait Message {
         unsafe { core::slice::from_raw_parts_mut(self.payload_ptr_mut(), self.payload_len()) }
     }
 
+    /// Sets the segment size for the message payload
+    fn set_segment_size(&mut self, size: usize);
+
+    /// Resets the message for future use
+    ///
+    /// # Safety
+    /// This method should only set the MTU to the original value
+    unsafe fn reset(&mut self, mtu: usize);
+
     /// Returns a pointer to the Message
     fn as_ptr(&self) -> *const c_void {
         self as *const _ as *const _
@@ -96,6 +104,15 @@ pub trait Ring {
     /// Returns the maximum transmission unit for the ring
     fn mtu(&self) -> usize;
 
+    /// Returns the maximum number of GSO segments that can be used
+    fn max_gso(&self) -> usize;
+
+    /// Disables the ability for the ring to send GSO messages
+    ///
+    /// This will be called in case the runtime encounters an IO error and will
+    /// try again with GSO disabled.
+    fn disable_gso(&mut self);
+
     /// Returns all of the messages in the ring
     ///
     /// The first half of the slice should be duplicated into the second half

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

@@ -4,8 +4,10 @@
 #![allow(unused_macros)]
 
 macro_rules! impl_message_delegate {
-    ($name:ident, $field:tt) => {
+    ($name:ident, $field:tt, $field_ty:ty) => {
         impl $crate::message::Message for $name {
+            const SUPPORTS_GSO: bool = <$field_ty as $crate::message::Message>::SUPPORTS_GSO;
+
             fn ecn(&self) -> ExplicitCongestionNotification {
                 $crate::message::Message::ecn(&self.$field)
             }
@@ -22,10 +24,6 @@ macro_rules! impl_message_delegate {
                 $crate::message::Message::set_remote_address(&mut self.$field, remote_address)
             }
 
-            fn reset_remote_address(&mut self) {
-                $crate::message::Message::reset_remote_address(&mut self.$field)
-            }
-
             fn payload_len(&self) -> usize {
                 $crate::message::Message::payload_len(&self.$field)
             }
@@ -34,6 +32,14 @@ macro_rules! impl_message_delegate {
                 $crate::message::Message::set_payload_len(&mut self.$field, payload_len)
             }
 
+            fn set_segment_size(&mut self, size: usize) {
+                $crate::message::Message::set_segment_size(&mut self.$field, size)
+            }
+
+            unsafe fn reset(&mut self, mtu: usize) {
+                $crate::message::Message::reset(&mut self.$field, mtu)
+            }
+
             fn replicate_fields_from(&mut self, other: &Self) {
                 $crate::message::Message::replicate_fields_from(&mut self.$field, &other.$field)
             }