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 aligned,
148        // and we have checked the space.
149        let mut cmsg = unsafe { self.inner.current_mut(self.buffer.buf_mut_ptr().cast()) }
150            .expect("sufficient space");
151        cmsg.set_level(level);
152        cmsg.set_ty(ty);
153        unsafe {
154            self.buffer.advance(cmsg.encode_data(value)?);
155        }
156
157        unsafe { self.inner.next(self.buffer.buf_mut_ptr().cast()) };
158
159        Ok(())
160    }
161}
162
163/// A fixed-size, stack-allocated buffer for ancillary (control) messages.
164///
165/// Properly aligned for the platform's control message header type
166/// (`cmsghdr` on Unix, `CMSGHDR` on Windows), so it can be passed directly
167/// to [`AncillaryIter`] and [`AncillaryBuilder`].
168pub struct AncillaryBuf<const N: usize> {
169    inner: [u8; N],
170    len: usize,
171    _align: [sys::cmsghdr; 0],
172}
173
174impl<const N: usize> AncillaryBuf<N> {
175    /// Create a new zeroed [`AncillaryBuf`].
176    pub fn new() -> Self {
177        Self {
178            inner: [0u8; N],
179            len: 0,
180            _align: [],
181        }
182    }
183
184    /// Create [`AncillaryBuilder`] with this buffer. The buffer will be zeroed
185    /// on creation.
186    ///
187    /// # Panics
188    ///
189    /// This function will panic if this buffer is too short.
190    pub fn builder(&mut self) -> AncillaryBuilder<'_, Self> {
191        AncillaryBuilder::new(self)
192    }
193}
194
195impl<const N: usize> Default for AncillaryBuf<N> {
196    fn default() -> Self {
197        Self::new()
198    }
199}
200
201impl<const N: usize> IoBuf for AncillaryBuf<N> {
202    fn as_init(&self) -> &[u8] {
203        &self.inner[..self.len]
204    }
205}
206
207impl<const N: usize> SetLen for AncillaryBuf<N> {
208    unsafe fn set_len(&mut self, len: usize) {
209        debug_assert!(len <= N);
210        self.len = len;
211    }
212}
213
214impl<const N: usize> IoBufMut for AncillaryBuf<N> {
215    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
216        self.inner.as_uninit()
217    }
218}
219
220impl<const N: usize> Deref for AncillaryBuf<N> {
221    type Target = [u8];
222
223    fn deref(&self) -> &Self::Target {
224        &self.inner[0..self.len]
225    }
226}
227
228impl<const N: usize> DerefMut for AncillaryBuf<N> {
229    fn deref_mut(&mut self) -> &mut Self::Target {
230        &mut self.inner[0..self.len]
231    }
232}
233
234/// Returns the buffer size required to hold one ancillary message carrying a
235/// value of type `T`.
236///
237/// This is the platform-appropriate equivalent of `CMSG_SPACE(T::SIZE)` on
238/// Unix or `WSA_CMSG_SPACE(T::SIZE)` on Windows, and can be used as a const
239/// generic argument for [`AncillaryBuf`].
240pub const fn ancillary_space<T: AncillaryData>() -> usize {
241    // SAFETY: CMSG_SPACE is always safe
242    #[allow(clippy::unnecessary_cast)]
243    unsafe {
244        sys::CMSG_SPACE(T::SIZE as _) as usize
245    }
246}
247
248/// Error that can occur when encoding or decoding ancillary data.
249#[derive(Debug)]
250pub enum CodecError {
251    /// The provided buffer is too small to hold the encoded data.
252    BufferTooSmall,
253    /// Another error occurred during encoding or decoding.
254    Other(Box<dyn std::error::Error + Send + Sync + 'static>),
255}
256
257impl CodecError {
258    /// Create a new [`CodecError::Other`] from any error type.
259    pub fn other(error: impl Into<Box<dyn std::error::Error + Send + Sync + 'static>>) -> Self {
260        Self::Other(error.into())
261    }
262
263    /// Attempt to downcast the error to a concrete type.
264    ///
265    /// Returns `Some(&T)` if the error is of type `T`, otherwise `None`.
266    pub fn downcast_ref<T: std::error::Error + 'static>(&self) -> Option<&T> {
267        match self {
268            Self::Other(e) => e.downcast_ref(),
269            _ => None,
270        }
271    }
272
273    /// Attempt to downcast the error to a concrete type.
274    ///
275    /// Returns `Some(&mut T)` if the error is of type `T`, otherwise `None`.
276    pub fn downcast_mut<T: std::error::Error + 'static>(&mut self) -> Option<&mut T> {
277        match self {
278            Self::Other(e) => e.downcast_mut(),
279            _ => None,
280        }
281    }
282}
283
284impl std::fmt::Display for CodecError {
285    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
286        match self {
287            Self::BufferTooSmall => write!(f, "buffer too small for encoding/decoding"),
288            Self::Other(e) => write!(f, "codec error: {}", e),
289        }
290    }
291}
292
293impl std::error::Error for CodecError {
294    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
295        match self {
296            Self::Other(e) => Some(e.as_ref()),
297            _ => None,
298        }
299    }
300}
301
302/// Trait for types that can be encoded and decoded as ancillary data payloads.
303///
304/// This trait enables a type to be used as the data payload in control messages
305/// (ancillary data). Types implementing this trait can be passed to
306/// [`AncillaryBuilder::push`] and retrieved via [`AncillaryRef::data`].
307///
308/// # Built-in Implementations
309///
310/// This trait is implemented for the following platform-specific types:
311///
312/// - Unix: `libc::in_addr`, `libc::in_pktinfo`, `libc::in6_pktinfo`
313/// - Windows: `IN_PKTINFO`, `IN6_PKTINFO`
314///
315/// When the `bytemuck` feature is enabled, this trait is also automatically
316/// implemented for types that implement [`bytemuck_ext::BitwiseAncillaryData`]:
317///
318/// - Primitive types: `()`, `u8`, `u16`, `u32`, `u64`, `u128`, `usize`, `i8`,
319///   `i16`, `i32`, `i64`, `i128`, `isize`, `f32`, `f64`
320/// - Fixed-size arrays of the above types (up to size 512)
321///
322/// For custom types with the `bytemuck` feature enabled, you can implement
323/// [`bytemuck_ext::BitwiseAncillaryData`] to automatically get
324/// [`AncillaryData`] (see [`bytemuck_ext`] for details). Otherwise, you must
325/// manually implement this trait with custom encoding/decoding logic.
326///
327/// # Example
328///
329/// ```
330/// use std::mem::MaybeUninit;
331///
332/// use compio_io::ancillary::{AncillaryData, CodecError};
333///
334/// struct MyData {
335///     value: u32,
336/// }
337///
338/// impl AncillaryData for MyData {
339///     const SIZE: usize = std::mem::size_of::<u32>();
340///
341///     fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
342///         if buffer.len() < Self::SIZE {
343///             return Err(CodecError::BufferTooSmall);
344///         }
345///         let bytes = self.value.to_ne_bytes();
346///         for (i, &byte) in bytes.iter().enumerate() {
347///             buffer[i] = MaybeUninit::new(byte);
348///         }
349///         Ok(())
350///     }
351///
352///     fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
353///         if buffer.len() < Self::SIZE {
354///             return Err(CodecError::BufferTooSmall);
355///         }
356///         let mut bytes = [0u8; 4];
357///         bytes.copy_from_slice(&buffer[..4]);
358///         Ok(MyData {
359///             value: u32::from_ne_bytes(bytes),
360///         })
361///     }
362/// }
363/// ```
364pub trait AncillaryData: Sized {
365    /// The size in bytes of the encoded representation.
366    ///
367    /// This defaults to `std::mem::size_of::<Self>()` but can be overridden
368    /// for types with custom encoding.
369    const SIZE: usize = std::mem::size_of::<Self>();
370
371    /// Encode this value into the provided buffer.
372    ///
373    /// # Errors
374    ///
375    /// Returns [`CodecError::BufferTooSmall`] if the buffer is too small to
376    /// hold the encoded data, or [`CodecError::Other`] for other encoding
377    /// errors.
378    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError>;
379
380    /// Decode a value from the provided buffer.
381    ///
382    /// # Errors
383    ///
384    /// Returns [`CodecError::BufferTooSmall`] if the buffer is too small,
385    /// or [`CodecError::Other`] for other decoding errors.
386    fn decode(buffer: &[u8]) -> Result<Self, CodecError>;
387}
388
389unsafe fn copy_to_bytes<T: AncillaryData>(
390    src: &T,
391    dest: &mut [MaybeUninit<u8>],
392) -> Result<(), CodecError> {
393    if dest.len() < T::SIZE {
394        return Err(CodecError::BufferTooSmall);
395    }
396    unsafe {
397        ptr::copy_nonoverlapping::<u8>(src as *const T as _, dest.as_mut_ptr() as _, T::SIZE);
398    }
399    Ok(())
400}
401
402unsafe fn copy_from_bytes<T: AncillaryData>(src: &[u8]) -> Result<T, CodecError> {
403    if src.len() < T::SIZE {
404        return Err(CodecError::BufferTooSmall);
405    }
406    let src_ptr = src.as_ptr() as *const T;
407    unsafe { Ok(ptr::read_unaligned(src_ptr)) }
408}
compio_io/ancillary/mod.rs

compio_io/ancillary/
mod.rs
