unix_udp_sock/
cmsg.rs

1use std::{
2    mem,
3    net::{IpAddr, Ipv6Addr, SocketAddr},
4    ptr,
5};
6
7use bytes::{Bytes, BytesMut};
8
9#[derive(Copy, Clone)]
10#[repr(align(8))] // Conservative bound for align_of<cmsghdr>
11pub struct Aligned<T>(pub T);
12
13/// Helper to encode a series of control messages ("cmsgs") to a buffer for use in `sendmsg`.
14///
15/// The operation must be "finished" for the msghdr to be usable, either by calling `finish`
16/// explicitly or by dropping the `Encoder`.
17pub struct Encoder<'a> {
18    hdr: &'a mut libc::msghdr,
19    cmsg: Option<&'a mut libc::cmsghdr>,
20    len: usize,
21}
22
23impl<'a> Encoder<'a> {
24    /// # Safety
25    /// - `hdr.msg_control` must be a suitably aligned pointer to `hdr.msg_controllen` bytes that
26    ///   can be safely written
27    /// - The `Encoder` must be dropped before `hdr` is passed to a system call, and must not be leaked.
28    pub unsafe fn new(hdr: &'a mut libc::msghdr) -> Self {
29        Self {
30            cmsg: libc::CMSG_FIRSTHDR(hdr).as_mut(),
31            hdr,
32            len: 0,
33        }
34    }
35
36    /// Append a control message to the buffer.
37    ///
38    /// # Panics
39    /// - If insufficient buffer space remains.
40    /// - If `T` has stricter alignment requirements than `cmsghdr`
41    #[allow(clippy::unnecessary_cast)]
42    pub fn push<T: Copy + ?Sized>(&mut self, level: libc::c_int, ty: libc::c_int, value: T) {
43        assert!(mem::align_of::<T>() <= mem::align_of::<libc::cmsghdr>());
44        let space = unsafe { libc::CMSG_SPACE(mem::size_of_val(&value) as _) as usize };
45        assert!(
46            self.hdr.msg_controllen as usize >= self.len + space,
47            "control message buffer too small. Required: {}, Available: {}",
48            self.len + space,
49            self.hdr.msg_controllen
50        );
51        let cmsg = self.cmsg.take().expect("no control buffer space remaining");
52        cmsg.cmsg_level = level;
53        cmsg.cmsg_type = ty;
54        cmsg.cmsg_len = unsafe { libc::CMSG_LEN(mem::size_of_val(&value) as _) } as _;
55        unsafe {
56            ptr::write(libc::CMSG_DATA(cmsg) as *const T as *mut T, value);
57        }
58        self.len += space;
59        self.cmsg = unsafe { libc::CMSG_NXTHDR(self.hdr, cmsg).as_mut() };
60    }
61
62    /// Finishes appending control messages to the buffer
63    pub fn finish(self) {
64        // Delegates to the `Drop` impl
65    }
66}
67
68// Statically guarantees that the encoding operation is "finished" before the control buffer is read
69// by `sendmsg`.
70impl<'a> Drop for Encoder<'a> {
71    fn drop(&mut self) {
72        self.hdr.msg_controllen = self.len as _;
73    }
74}
75
76/// # Safety
77///
78/// `cmsg` must refer to a cmsg containing a payload of type `T`
79#[allow(clippy::unnecessary_cast)]
80pub unsafe fn decode<T: Copy>(cmsg: &libc::cmsghdr) -> T {
81    assert!(mem::align_of::<T>() <= mem::align_of::<libc::cmsghdr>());
82    debug_assert_eq!(
83        cmsg.cmsg_len as usize,
84        libc::CMSG_LEN(mem::size_of::<T>() as _) as usize
85    );
86    ptr::read(libc::CMSG_DATA(cmsg) as *const T)
87}
88
89pub struct Iter<'a> {
90    hdr: &'a libc::msghdr,
91    cmsg: Option<&'a libc::cmsghdr>,
92}
93
94impl<'a> Iter<'a> {
95    /// # Safety
96    ///
97    /// `hdr.msg_control` must point to memory outliving `'a` which can be soundly read for the
98    /// lifetime of the constructed `Iter` and contains a buffer of cmsgs, i.e. is aligned for
99    /// `cmsghdr`, is fully initialized, and has correct internal links.
100    pub unsafe fn new(hdr: &'a libc::msghdr) -> Self {
101        Self {
102            hdr,
103            cmsg: libc::CMSG_FIRSTHDR(hdr).as_ref(),
104        }
105    }
106}
107
108impl<'a> Iterator for Iter<'a> {
109    type Item = &'a libc::cmsghdr;
110    fn next(&mut self) -> Option<&'a libc::cmsghdr> {
111        let current = self.cmsg.take()?;
112        self.cmsg = unsafe { libc::CMSG_NXTHDR(self.hdr, current).as_ref() };
113        Some(current)
114    }
115}
116
117/// Explicit congestion notification codepoint
118#[repr(u8)]
119#[derive(Debug, Copy, Clone, Eq, PartialEq)]
120pub enum EcnCodepoint {
121    #[doc(hidden)]
122    Ect0 = 0b10,
123    #[doc(hidden)]
124    Ect1 = 0b01,
125    #[doc(hidden)]
126    Ce = 0b11,
127}
128
129impl EcnCodepoint {
130    /// Create new object from the given bits
131    pub fn from_bits(x: u8) -> Option<Self> {
132        use self::EcnCodepoint::*;
133        Some(match x & 0b11 {
134            0b10 => Ect0,
135            0b01 => Ect1,
136            0b11 => Ce,
137            _ => {
138                return None;
139            }
140        })
141    }
142}
143
144/// An outgoing packet
145#[derive(Debug)]
146pub struct Transmit<B> {
147    /// The socket this datagram should be sent to
148    pub dst: SocketAddr,
149    /// Explicit congestion notification bits to set on the packet
150    pub ecn: Option<EcnCodepoint>,
151    /// Contents of the datagram
152    pub contents: B,
153    /// The segment size if this transmission contains multiple datagrams.
154    /// This is `None` if the transmit only contains a single datagram
155    pub segment_size: Option<usize>,
156    /// Optional source IP address for the datagram
157    pub src: Option<Source>,
158}
159
160impl<B: AsPtr<u8>> Transmit<B> {
161    pub fn new(dst: SocketAddr, contents: B) -> Self {
162        Self {
163            dst,
164            contents,
165            ecn: None,
166            segment_size: None,
167            src: None,
168        }
169    }
170    pub fn src_ip(self, src_ip: Source) -> Self {
171        Transmit {
172            src: Some(src_ip),
173            ..self
174        }
175    }
176    pub fn segment_size(self, size: usize) -> Self {
177        Transmit {
178            segment_size: Some(size),
179            ..self
180        }
181    }
182    pub fn ecn(self, ecn: EcnCodepoint) -> Self {
183        Transmit {
184            ecn: Some(ecn),
185            ..self
186        }
187    }
188}
189
190/// Select how to set the source IP - using either interface id or the IP itself
191#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord)]
192pub enum Source {
193    /// Set source IP explicitly by IP
194    Ip(IpAddr),
195    /// Set via interface index, ipv4 only
196    Interface(u32),
197    /// Set v6 interface index and IP
198    InterfaceV6(u32, Ipv6Addr),
199}
200
201/// A buffer that can be turned into a raw ptr and has a len.
202/// Is used to be generic over Vec<u8>, [u8], Bytes, BytesMut
203pub trait AsPtr<T> {
204    fn as_ptr(&self) -> *const T;
205    fn len(&self) -> usize;
206    fn is_empty(&self) -> bool {
207        self.len() == 0
208    }
209}
210
211impl<T, const N: usize> AsPtr<T> for &[T; N] {
212    fn as_ptr(&self) -> *const T {
213        self.as_slice().as_ptr()
214    }
215
216    fn len(&self) -> usize {
217        N
218    }
219}
220
221impl<T, const N: usize> AsPtr<T> for [T; N] {
222    fn as_ptr(&self) -> *const T {
223        self.as_slice().as_ptr()
224    }
225
226    fn len(&self) -> usize {
227        N
228    }
229}
230
231impl<T> AsPtr<T> for Vec<T> {
232    fn as_ptr(&self) -> *const T {
233        <Vec<T>>::as_ptr(self)
234    }
235    fn len(&self) -> usize {
236        <Vec<T>>::len(self)
237    }
238}
239
240impl<T> AsPtr<T> for [T] {
241    fn as_ptr(&self) -> *const T {
242        <[T]>::as_ptr(self)
243    }
244    fn len(&self) -> usize {
245        <[T]>::len(self)
246    }
247}
248
249impl AsPtr<u8> for BytesMut {
250    fn as_ptr(&self) -> *const u8 {
251        <[u8]>::as_ptr(self.as_ref())
252    }
253    fn len(&self) -> usize {
254        self.len()
255    }
256}
257
258impl AsPtr<u8> for Bytes {
259    fn as_ptr(&self) -> *const u8 {
260        <[u8]>::as_ptr(self.as_ref())
261    }
262    fn len(&self) -> usize {
263        self.len()
264    }
265}
unix_udp_sock/cmsg.rs

unix_udp_sock/
cmsg.rs
