Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 75 additions & 1 deletion compio-io/src/ancillary/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)] {
Expand Down Expand Up @@ -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]
}
}
62 changes: 5 additions & 57 deletions compio-quic/src/socket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);

Expand Down