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

compio_io/ancillary/
mod.rs
