compio_io\ancillary/
mod.rs

1//! Ancillary data (control message) support for connected streams.
2//!
3//! Ancillary messages are used to pass out-of-band information such as file
4//! descriptors (Unix domain sockets), credentials, or kTLS record types.
5//!
6//! # Types
7//!
8//! - [`AncillaryBuf`]: A fixed-size, properly aligned stack buffer for
9//!   ancillary messages.
10//! - [`AncillaryBuilder`]: A builder for constructing ancillary messages into a
11//!   [`AncillaryBuf`].
12//! - [`AncillaryIter`]: An iterator over a buffer of ancillary messages.
13//! - [`AncillaryRef`]: A reference to a single ancillary data entry.
14//! - [`AncillaryData`]: Trait for types that can be encoded/decoded as
15//!   ancillary data payloads.
16//! - [`CodecError`]: Error type for encoding/decoding operations.
17//!
18//! # Traits
19//!
20//! - [`AsyncReadAncillary`]: read data together with ancillary data
21//! - [`AsyncWriteAncillary`]: write data together with ancillary data
22//!
23//! # Functions
24//!
25//! - [`ancillary_space`]: Helper function to calculate ancillary message size
26//!   for a type.
27//!
28//! # Modules
29//!
30//! - [`bytemuck_ext`]: Extension module for automatic [`AncillaryData`]
31//!   implementation via bytemuck (requires `bytemuck` feature).
32
33use std::{
34    mem::MaybeUninit,
35    ops::{Deref, DerefMut},
36    ptr,
37};
38
39use compio_buf::{IoBuf, IoBufMut, IoBufMutExt, SetLen, SetLenExt};
40
41mod io;
42
43pub use rustix::net::ReturnFlags;
44
45pub use self::io::*;
46
47#[cfg(feature = "bytemuck")]
48pub mod bytemuck_ext;
49mod sys;
50
51/// Reference to an ancillary (control) message.
52pub struct AncillaryRef<'a>(sys::CMsgRef<'a>);
53
54impl AncillaryRef<'_> {
55    /// Returns the level of the control message.
56    pub fn level(&self) -> i32 {
57        self.0.level()
58    }
59
60    /// Returns the type of the control message.
61    pub fn ty(&self) -> i32 {
62        self.0.ty()
63    }
64
65    /// Returns the length of the control message.
66    #[allow(clippy::len_without_is_empty)]
67    pub fn len(&self) -> usize {
68        self.0.len() as _
69    }
70
71    /// Returns a copy of the data in the control message.
72    pub fn data<T: AncillaryData>(&self) -> Result<T, CodecError> {
73        self.0.decode_data()
74    }
75}
76
77/// An iterator for ancillary (control) messages.
78pub struct AncillaryIter<'a> {
79    inner: sys::CMsgIter,
80    buffer: &'a [u8],
81}
82
83impl<'a> AncillaryIter<'a> {
84    /// Create [`AncillaryIter`] with the given buffer.
85    ///
86    /// # Panics
87    ///
88    /// This function will panic if the buffer is too short or not properly
89    /// aligned.
90    ///
91    /// # Safety
92    ///
93    /// The buffer should contain valid control messages.
94    pub unsafe fn new(buffer: &'a [u8]) -> Self {
95        Self {
96            inner: sys::CMsgIter::new(buffer.as_ptr(), buffer.len()),
97            buffer,
98        }
99    }
100}
101
102impl<'a> Iterator for AncillaryIter<'a> {
103    type Item = AncillaryRef<'a>;
104
105    fn next(&mut self) -> Option<Self::Item> {
106        unsafe {
107            let cmsg = self.inner.current(self.buffer.as_ptr());
108            self.inner.next(self.buffer.as_ptr());
109            cmsg.map(AncillaryRef)
110        }
111    }
112}
113
114/// Helper to construct ancillary (control) messages.
115pub struct AncillaryBuilder<'a, B: ?Sized> {
116    inner: sys::CMsgIter,
117    buffer: &'a mut B,
118}
119
120impl<'a, B: IoBufMut + ?Sized> AncillaryBuilder<'a, B> {
121    /// Create [`AncillaryBuilder`] with the given buffer. The buffer will be
122    /// cleared on creation.
123    ///
124    /// # Panics
125    ///
126    /// This function will panic if the buffer is too short or not properly
127    /// aligned.
128    pub fn new(buffer: &'a mut B) -> Self {
129        // SAFETY: always safe to make it empty.
130        unsafe { buffer.set_len(0) };
131        let slice = buffer.ensure_init();
132        let inner = sys::CMsgIter::new(slice.as_ptr(), slice.len());
133        Self { inner, buffer }
134    }
135
136    /// Append a control message into the buffer.
137    pub fn push<T: AncillaryData>(
138        &mut self,
139        level: i32,
140        ty: i32,
141        value: &T,
142    ) -> Result<(), CodecError> {
143        if !self.inner.is_space_enough(T::SIZE) {
144            return Err(CodecError::BufferTooSmall);
145        }
146
147        // SAFETY: method `new` guarantees the buffer is zeroed and properly
148        // aligned, and we have checked the space.
149        let mut cmsg = unsafe {
150            self.inner
151                .current_mut::<T>(self.buffer.buf_mut_ptr().cast())
152        }
153        .expect("sufficient space");
154        cmsg.set_level(level);
155        cmsg.set_ty(ty);
156        unsafe {
157            self.buffer.advance(cmsg.encode_data(value)?);
158        }
159
160        unsafe { self.inner.next(self.buffer.buf_mut_ptr().cast()) };
161
162        Ok(())
163    }
164}
165
166/// A fixed-size, stack-allocated buffer for ancillary (control) messages.
167///
168/// Properly aligned for the platform's control message header type
169/// (`cmsghdr` on Unix, `CMSGHDR` on Windows), so it can be passed directly
170/// to [`AncillaryIter`] and [`AncillaryBuilder`].
171pub struct AncillaryBuf<const N: usize> {
172    inner: [u8; N],
173    len: usize,
174    _align: [sys::cmsghdr; 0],
175}
176
177impl<const N: usize> AncillaryBuf<N> {
178    /// Create a new zeroed [`AncillaryBuf`].
179    pub fn new() -> Self {
180        Self {
181            inner: [0u8; N],
182            len: 0,
183            _align: [],
184        }
185    }
186
187    /// Create [`AncillaryBuilder`] with this buffer. The buffer will be zeroed
188    /// on creation.
189    ///
190    /// # Panics
191    ///
192    /// This function will panic if this buffer is too short.
193    pub fn builder(&mut self) -> AncillaryBuilder<'_, Self> {
194        AncillaryBuilder::new(self)
195    }
196}
197
198impl<const N: usize> Default for AncillaryBuf<N> {
199    fn default() -> Self {
200        Self::new()
201    }
202}
203
204impl<const N: usize> IoBuf for AncillaryBuf<N> {
205    fn as_init(&self) -> &[u8] {
206        &self.inner[..self.len]
207    }
208}
209
210impl<const N: usize> SetLen for AncillaryBuf<N> {
211    unsafe fn set_len(&mut self, len: usize) {
212        debug_assert!(len <= N);
213        self.len = len;
214    }
215}
216
217impl<const N: usize> IoBufMut for AncillaryBuf<N> {
218    unsafe fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
219        unsafe { std::slice::from_raw_parts_mut(self.inner.as_mut_ptr().cast(), N) }
220    }
221}
222
223impl<const N: usize> Deref for AncillaryBuf<N> {
224    type Target = [u8];
225
226    fn deref(&self) -> &Self::Target {
227        &self.inner[0..self.len]
228    }
229}
230
231impl<const N: usize> DerefMut for AncillaryBuf<N> {
232    fn deref_mut(&mut self) -> &mut Self::Target {
233        &mut self.inner[0..self.len]
234    }
235}
236
237/// Returns the buffer size required to hold one ancillary message carrying a
238/// value of type `T`.
239///
240/// This is the platform-appropriate equivalent of `CMSG_SPACE(T::SIZE)` on
241/// Unix or `WSA_CMSG_SPACE(T::SIZE)` on Windows, and can be used as a const
242/// generic argument for [`AncillaryBuf`].
243pub const fn ancillary_space<T: AncillaryData>() -> usize {
244    // SAFETY: CMSG_SPACE is always safe
245    #[allow(clippy::unnecessary_cast)]
246    unsafe {
247        sys::CMSG_SPACE(T::SIZE as _) as usize
248    }
249}
250
251/// Error that can occur when encoding or decoding ancillary data.
252#[derive(Debug)]
253pub enum CodecError {
254    /// The provided buffer is too small to hold the encoded data.
255    BufferTooSmall,
256    /// Another error occurred during encoding or decoding.
257    Other(Box<dyn std::error::Error + Send + Sync + 'static>),
258}
259
260impl CodecError {
261    /// Create a new [`CodecError::Other`] from any error type.
262    pub fn other(error: impl Into<Box<dyn std::error::Error + Send + Sync + 'static>>) -> Self {
263        Self::Other(error.into())
264    }
265
266    /// Attempt to downcast the error to a concrete type.
267    ///
268    /// Returns `Some(&T)` if the error is of type `T`, otherwise `None`.
269    pub fn downcast_ref<T: std::error::Error + 'static>(&self) -> Option<&T> {
270        match self {
271            Self::Other(e) => e.downcast_ref(),
272            _ => None,
273        }
274    }
275
276    /// Attempt to downcast the error to a concrete type.
277    ///
278    /// Returns `Some(&mut T)` if the error is of type `T`, otherwise `None`.
279    pub fn downcast_mut<T: std::error::Error + 'static>(&mut self) -> Option<&mut T> {
280        match self {
281            Self::Other(e) => e.downcast_mut(),
282            _ => None,
283        }
284    }
285}
286
287impl std::fmt::Display for CodecError {
288    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
289        match self {
290            Self::BufferTooSmall => write!(f, "buffer too small for encoding/decoding"),
291            Self::Other(e) => write!(f, "codec error: {}", e),
292        }
293    }
294}
295
296impl std::error::Error for CodecError {
297    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
298        match self {
299            Self::Other(e) => Some(e.as_ref()),
300            _ => None,
301        }
302    }
303}
304
305/// Trait for types that can be encoded and decoded as ancillary data payloads.
306///
307/// This trait enables a type to be used as the data payload in control messages
308/// (ancillary data). Types implementing this trait can be passed to
309/// [`AncillaryBuilder::push`] and retrieved via [`AncillaryRef::data`].
310///
311/// # Built-in Implementations
312///
313/// This trait is implemented for the following platform-specific types:
314///
315/// - Unix: `libc::in_addr`, `libc::in_pktinfo`, `libc::in6_pktinfo`
316/// - Windows: `IN_PKTINFO`, `IN6_PKTINFO`
317///
318/// When the `bytemuck` feature is enabled, this trait is also automatically
319/// implemented for types that implement [`bytemuck_ext::BitwiseAncillaryData`]:
320///
321/// - Primitive types: `()`, `u8`, `u16`, `u32`, `u64`, `u128`, `usize`, `i8`,
322///   `i16`, `i32`, `i64`, `i128`, `isize`, `f32`, `f64`
323/// - Fixed-size arrays of the above types (up to size 512)
324///
325/// For custom types with the `bytemuck` feature enabled, you can implement
326/// [`bytemuck_ext::BitwiseAncillaryData`] to automatically get
327/// [`AncillaryData`] (see [`bytemuck_ext`] for details). Otherwise, you must
328/// manually implement this trait with custom encoding/decoding logic.
329///
330/// # Example
331///
332/// ```
333/// use std::mem::MaybeUninit;
334///
335/// use compio_io::ancillary::{AncillaryData, CodecError};
336///
337/// struct MyData {
338///     value: u32,
339/// }
340///
341/// impl AncillaryData for MyData {
342///     const SIZE: usize = std::mem::size_of::<u32>();
343///
344///     fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
345///         if buffer.len() < Self::SIZE {
346///             return Err(CodecError::BufferTooSmall);
347///         }
348///         let bytes = self.value.to_ne_bytes();
349///         for (i, &byte) in bytes.iter().enumerate() {
350///             buffer[i] = MaybeUninit::new(byte);
351///         }
352///         Ok(())
353///     }
354///
355///     fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
356///         if buffer.len() < Self::SIZE {
357///             return Err(CodecError::BufferTooSmall);
358///         }
359///         let mut bytes = [0u8; 4];
360///         bytes.copy_from_slice(&buffer[..4]);
361///         Ok(MyData {
362///             value: u32::from_ne_bytes(bytes),
363///         })
364///     }
365/// }
366/// ```
367pub trait AncillaryData: Sized {
368    /// The size in bytes of the encoded representation.
369    ///
370    /// This defaults to `std::mem::size_of::<Self>()` but can be overridden
371    /// for types with custom encoding.
372    const SIZE: usize = std::mem::size_of::<Self>();
373
374    /// Encode this value into the provided buffer.
375    ///
376    /// # Errors
377    ///
378    /// Returns [`CodecError::BufferTooSmall`] if the buffer is too small to
379    /// hold the encoded data, or [`CodecError::Other`] for other encoding
380    /// errors.
381    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError>;
382
383    /// Decode a value from the provided buffer.
384    ///
385    /// # Errors
386    ///
387    /// Returns [`CodecError::BufferTooSmall`] if the buffer is too small,
388    /// or [`CodecError::Other`] for other decoding errors.
389    fn decode(buffer: &[u8]) -> Result<Self, CodecError>;
390}
391
392unsafe fn copy_to_bytes<T: AncillaryData>(
393    src: &T,
394    dest: &mut [MaybeUninit<u8>],
395) -> Result<(), CodecError> {
396    if dest.len() < T::SIZE {
397        return Err(CodecError::BufferTooSmall);
398    }
399    unsafe {
400        ptr::copy_nonoverlapping::<u8>(src as *const T as _, dest.as_mut_ptr() as _, T::SIZE);
401    }
402    Ok(())
403}
404
405unsafe fn copy_from_bytes<T: AncillaryData>(src: &[u8]) -> Result<T, CodecError> {
406    if src.len() < T::SIZE {
407        return Err(CodecError::BufferTooSmall);
408    }
409    let src_ptr = src.as_ptr() as *const T;
410    unsafe { Ok(ptr::read_unaligned(src_ptr)) }
411}
compio_io\ancillary/mod.rs

compio_io\ancillary/
mod.rs
