From 971cc68cbcf3e4015447470dd7a8ef6fd356633f Mon Sep 17 00:00:00 2001
From: Asakura Mizu <asakuramizu111@gmail.com>
Date: Thu, 1 Aug 2024 02:42:18 +0800
Subject: [PATCH] refactor(io,quic): move quic Ancillary to io

Co-authored-by: Fantix King <fantix.king@gmail.com>
---
 compio-io/src/ancillary/mod.rs | 76 +++++++++++++++++++++++++++++++++-
 compio-quic/src/socket.rs      | 62 +++------------------------
 2 files changed, 80 insertions(+), 58 deletions(-)

diff --git a/compio-io/src/ancillary/mod.rs b/compio-io/src/ancillary/mod.rs
index a6ff839d7..95a281158 100644
--- a/compio-io/src/ancillary/mod.rs
+++ b/compio-io/src/ancillary/mod.rs
@@ -9,8 +9,18 @@
 //! - [`AncillaryIter`]: An iterator over a buffer of ancillary messages.
 //! - [`AncillaryBuilder`]: A builder for constructing ancillary messages into a
 //!   caller-supplied send buffer.
+//! - [`AncillaryBuf`]: A fixed-size, properly aligned stack buffer for
+//!   ancillary data
 
-use std::{marker::PhantomData, mem::MaybeUninit};
+use std::{
+    marker::PhantomData,
+    mem::MaybeUninit,
+    ops::{Deref, DerefMut},
+};
+
+use compio_buf::{IoBuf, IoBufMut, SetLen};
+#[cfg(windows)]
+use windows_sys::Win32::Networking::WinSock;
 
 cfg_if::cfg_if! {
     if #[cfg(windows)] {
@@ -141,3 +151,67 @@ impl<'a> AncillaryBuilder<'a> {
         Some(())
     }
 }
+
+/// A fixed-size, stack-allocated buffer for ancillary (control) messages.
+///
+/// Properly aligned for the platform's control message header type
+/// (`cmsghdr` on Unix, `CMSGHDR` on Windows), so it can be passed directly
+/// to [`AncillaryIter`] and [`AncillaryBuilder`].
+pub struct AncillaryBuf<const N: usize> {
+    inner: [u8; N],
+    len: usize,
+    #[cfg(unix)]
+    _align: [libc::cmsghdr; 0],
+    #[cfg(windows)]
+    _align: [WinSock::CMSGHDR; 0],
+}
+
+impl<const N: usize> AncillaryBuf<N> {
+    /// Create a new zeroed [`AncillaryBuf`].
+    pub fn new() -> Self {
+        Self {
+            inner: [0u8; N],
+            len: 0,
+            _align: [],
+        }
+    }
+}
+
+impl<const N: usize> Default for AncillaryBuf<N> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<const N: usize> IoBuf for AncillaryBuf<N> {
+    fn as_init(&self) -> &[u8] {
+        &self.inner[..self.len]
+    }
+}
+
+impl<const N: usize> SetLen for AncillaryBuf<N> {
+    unsafe fn set_len(&mut self, len: usize) {
+        debug_assert!(len <= N);
+        self.len = len;
+    }
+}
+
+impl<const N: usize> IoBufMut for AncillaryBuf<N> {
+    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
+        self.inner.as_uninit()
+    }
+}
+
+impl<const N: usize> Deref for AncillaryBuf<N> {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner[0..self.len]
+    }
+}
+
+impl<const N: usize> DerefMut for AncillaryBuf<N> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner[0..self.len]
+    }
+}
diff --git a/compio-quic/src/socket.rs b/compio-quic/src/socket.rs
index 810e08f79..34f01ed57 100644
--- a/compio-quic/src/socket.rs
+++ b/compio-quic/src/socket.rs
@@ -12,12 +12,11 @@ use std::{
     future::Future,
     io,
     net::{IpAddr, SocketAddr},
-    ops::{Deref, DerefMut},
     sync::atomic::Ordering,
 };
 
 use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, SetLen, buf_try};
-use compio_io::ancillary::{AncillaryBuilder, AncillaryIter};
+use compio_io::ancillary::{AncillaryBuf, AncillaryBuilder, AncillaryIter};
 use compio_net::UdpSocket;
 use quinn_proto::{EcnCodepoint, Transmit};
 #[cfg(windows)]
@@ -58,58 +57,6 @@ pub(crate) struct RecvMeta {
 
 const CMSG_LEN: usize = 128;
 
-struct Ancillary<const N: usize> {
-    inner: [u8; N],
-    len: usize,
-    #[cfg(unix)]
-    _align: [libc::cmsghdr; 0],
-    #[cfg(windows)]
-    _align: [WinSock::CMSGHDR; 0],
-}
-
-impl<const N: usize> Ancillary<N> {
-    fn new() -> Self {
-        Self {
-            inner: [0u8; N],
-            len: 0,
-            _align: [],
-        }
-    }
-}
-
-impl<const N: usize> IoBuf for Ancillary<N> {
-    fn as_init(&self) -> &[u8] {
-        &self.inner[..self.len]
-    }
-}
-
-impl<const N: usize> SetLen for Ancillary<N> {
-    unsafe fn set_len(&mut self, len: usize) {
-        debug_assert!(len <= N);
-        self.len = len;
-    }
-}
-
-impl<const N: usize> IoBufMut for Ancillary<N> {
-    fn as_uninit(&mut self) -> &mut [std::mem::MaybeUninit<u8>] {
-        self.inner.as_uninit()
-    }
-}
-
-impl<const N: usize> Deref for Ancillary<N> {
-    type Target = [u8];
-
-    fn deref(&self) -> &Self::Target {
-        &self.inner[0..self.len]
-    }
-}
-
-impl<const N: usize> DerefMut for Ancillary<N> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.inner[0..self.len]
-    }
-}
-
 #[cfg(linux_all)]
 #[inline]
 fn max_gso_segments(socket: &UdpSocket) -> io::Result<usize> {
@@ -333,7 +280,7 @@ impl Socket {
     }
 
     pub async fn recv<T: IoBufMut>(&self, buffer: T) -> BufResult<RecvMeta, T> {
-        let control = Ancillary::<CMSG_LEN>::new();
+        let control = AncillaryBuf::<CMSG_LEN>::new();
 
         let BufResult(res, (buffer, control)) = self.inner.recv_msg(buffer, control).await;
         let ((len, _, remote), buffer) = buf_try!(res, buffer);
@@ -420,7 +367,7 @@ impl Socket {
         let is_ipv4 = transmit.destination.ip().to_canonical().is_ipv4();
         let ecn = transmit.ecn.map_or(0, |x| x as u8);
 
-        let mut control = Ancillary::<CMSG_LEN>::new();
+        let mut control = AncillaryBuf::<CMSG_LEN>::new();
         let mut builder = AncillaryBuilder::new(control.as_uninit());
 
         // ECN
@@ -516,7 +463,8 @@ impl Socket {
         }
 
         let len = builder.finish();
-        control.len = len;
+        // SAFETY: AncillaryBuilder ensures the buffer is initialized within len
+        unsafe { control.set_len(len) };
 
         let mut buffer = buffer.slice(0..transmit.size);