Skip to main content

compio_io/ancillary/
sys.rs

1use std::{mem::MaybeUninit, slice};
2
3#[cfg(unix)]
4use libc::{CMSG_DATA, CMSG_FIRSTHDR, CMSG_LEN, CMSG_NXTHDR, msghdr};
5#[cfg(unix)]
6pub use libc::{CMSG_SPACE, cmsghdr};
7#[cfg(windows)]
8pub use windows_sys::Win32::Networking::WinSock::CMSGHDR as cmsghdr;
9#[cfg(windows)]
10use windows_sys::Win32::Networking::WinSock::{self, IN_PKTINFO, IN6_PKTINFO};
11
12use super::{AncillaryData, CodecError, copy_from_bytes, copy_to_bytes};
13
14#[cfg(windows)]
15#[allow(non_snake_case)]
16mod windows_macros {
17    use std::ptr::null_mut;
18
19    use windows_sys::Win32::Networking::WinSock::{CMSGHDR, WSABUF, WSAMSG};
20
21    const fn CMSG_ALIGN(length: usize) -> usize {
22        (length + align_of::<CMSGHDR>() - 1) & !(align_of::<CMSGHDR>() - 1)
23    }
24
25    const WSA_CMSGDATA_OFFSET: usize = CMSG_ALIGN(size_of::<CMSGHDR>());
26
27    pub unsafe fn CMSG_DATA(cmsg: *const CMSGHDR) -> *mut u8 {
28        unsafe { cmsg.offset(1) as *mut u8 }
29    }
30
31    pub const unsafe fn CMSG_SPACE(length: usize) -> usize {
32        WSA_CMSGDATA_OFFSET + CMSG_ALIGN(length)
33    }
34
35    pub const unsafe fn CMSG_LEN(length: usize) -> usize {
36        WSA_CMSGDATA_OFFSET + length
37    }
38
39    pub unsafe fn CMSG_FIRSTHDR(msg: *const WSAMSG) -> *mut CMSGHDR {
40        unsafe {
41            if (*msg).Control.len as usize >= size_of::<CMSGHDR>() {
42                (*msg).Control.buf as _
43            } else {
44                null_mut()
45            }
46        }
47    }
48
49    pub unsafe fn CMSG_NXTHDR(msg: *const WSAMSG, cmsg: *const CMSGHDR) -> *mut CMSGHDR {
50        unsafe {
51            if cmsg.is_null() {
52                CMSG_FIRSTHDR(msg)
53            } else {
54                let next = cmsg
55                    .cast::<u8>()
56                    .wrapping_add(CMSG_ALIGN((*cmsg).cmsg_len))
57                    .cast::<CMSGHDR>();
58                let end = (*msg).Control.buf.addr() + (*msg).Control.len as usize;
59                if next.addr() + size_of::<CMSGHDR>() > end {
60                    null_mut()
61                } else {
62                    next.cast_mut()
63                }
64            }
65        }
66    }
67
68    pub fn msghdr_from_raw(ptr: *const u8, len: usize) -> WSAMSG {
69        WSAMSG {
70            Control: WSABUF {
71                len: len as _,
72                buf: ptr as _,
73            },
74            ..unsafe { std::mem::zeroed() }
75        }
76    }
77}
78
79#[cfg(windows)]
80pub use windows_macros::CMSG_SPACE;
81#[cfg(windows)]
82use windows_macros::{CMSG_DATA, CMSG_FIRSTHDR, CMSG_LEN, CMSG_NXTHDR, msghdr_from_raw};
83
84#[cfg(unix)]
85fn msghdr_from_raw(ptr: *const u8, len: usize) -> msghdr {
86    let mut msg: msghdr = unsafe { std::mem::zeroed() };
87    msg.msg_control = ptr as _;
88    msg.msg_controllen = len as _;
89    msg
90}
91
92pub(crate) struct CMsgRef<'a>(&'a cmsghdr, &'a [u8]);
93
94impl CMsgRef<'_> {
95    pub(crate) fn level(&self) -> i32 {
96        self.0.cmsg_level as _
97    }
98
99    pub(crate) fn ty(&self) -> i32 {
100        self.0.cmsg_type as _
101    }
102
103    pub(crate) fn len(&self) -> usize {
104        self.0.cmsg_len as _
105    }
106
107    pub(crate) fn decode_data<T: AncillaryData>(&self) -> Result<T, CodecError> {
108        T::decode(self.1)
109    }
110}
111
112pub(crate) struct CMsgMut<'a>(&'a mut cmsghdr, &'a mut [MaybeUninit<u8>]);
113
114impl CMsgMut<'_> {
115    pub(crate) fn set_level(&mut self, level: i32) {
116        self.0.cmsg_level = level as _;
117    }
118
119    pub(crate) fn set_ty(&mut self, ty: i32) {
120        self.0.cmsg_type = ty as _;
121    }
122
123    pub(crate) fn encode_data<T: AncillaryData>(&mut self, value: &T) -> Result<usize, CodecError> {
124        debug_assert_eq!(self.1.len(), T::SIZE);
125        self.0.cmsg_len = unsafe { CMSG_LEN(T::SIZE as _) } as _;
126        value.encode(self.1)?;
127        Ok(unsafe { CMSG_SPACE(T::SIZE as _) } as _)
128    }
129}
130
131pub(crate) struct CMsgIter {
132    len: usize,
133    offset: Option<usize>,
134}
135
136impl CMsgIter {
137    pub(crate) fn new(ptr: *const u8, len: usize) -> Self {
138        assert!(len >= unsafe { CMSG_SPACE(0) as _ }, "buffer too short");
139        assert!(ptr.cast::<cmsghdr>().is_aligned(), "misaligned buffer");
140
141        let msg = msghdr_from_raw(ptr.cast_mut(), len);
142        let first_cmsg = unsafe { CMSG_FIRSTHDR(&msg) };
143
144        let offset = if first_cmsg.is_null() {
145            None
146        } else {
147            Some(first_cmsg.addr() - ptr.addr())
148        };
149        Self { len, offset }
150    }
151
152    pub(crate) unsafe fn current<'a>(&self, ptr: *const u8) -> Option<CMsgRef<'a>> {
153        let offset = self.offset?;
154        let header = unsafe { ptr.add(offset).cast::<cmsghdr>().as_ref() }?;
155        let data_offset = unsafe { CMSG_DATA(header) }.addr() - ptr.addr();
156        let len: usize = header.cmsg_len as _;
157        let data_len = len.saturating_sub(data_offset - offset);
158        let buffer = unsafe { slice::from_raw_parts(ptr.add(data_offset), data_len) };
159        Some(CMsgRef(header, buffer))
160    }
161
162    pub(crate) unsafe fn next(&mut self, ptr: *const u8) {
163        if let Some(offset) = self.offset {
164            let msg = msghdr_from_raw(ptr, self.len);
165            let next_cmsg = unsafe { CMSG_NXTHDR(&msg, ptr.add(offset).cast()) };
166            if next_cmsg.is_null() {
167                self.offset = None;
168            } else {
169                self.offset = Some(next_cmsg.addr() - ptr.addr());
170            }
171        }
172    }
173
174    pub(crate) unsafe fn current_mut<'a, T: AncillaryData>(
175        &self,
176        ptr: *mut u8,
177    ) -> Option<CMsgMut<'a>> {
178        let offset = self.offset?;
179        let header = unsafe { ptr.add(offset).cast::<cmsghdr>().as_mut() }?;
180        let data_ptr = unsafe { CMSG_DATA(header) };
181        // Retrieve the offset of the data pointer again to make miri happy.
182        let data_offset = unsafe { data_ptr.offset_from(ptr) };
183        let buffer = unsafe { slice::from_raw_parts_mut(ptr.offset(data_offset).cast(), T::SIZE) };
184        Some(CMsgMut(header, buffer))
185    }
186
187    pub(crate) fn is_space_enough(&self, space: usize) -> bool {
188        if let Some(offset) = self.offset {
189            #[allow(clippy::unnecessary_cast)]
190            let space = unsafe { CMSG_SPACE(space as _) } as usize;
191            offset + space <= self.len
192        } else {
193            false
194        }
195    }
196}
197
198#[cfg(unix)]
199impl AncillaryData for libc::in_addr {
200    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
201        unsafe { copy_to_bytes(self, buffer) }
202    }
203
204    fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
205        unsafe { copy_from_bytes(buffer) }
206    }
207}
208
209#[cfg(any(target_os = "linux", target_os = "android"))]
210impl AncillaryData for libc::in_pktinfo {
211    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
212        let mut pktinfo: libc::in_pktinfo = unsafe { std::mem::zeroed() };
213        pktinfo.ipi_ifindex = self.ipi_ifindex;
214        pktinfo.ipi_spec_dst.s_addr = self.ipi_spec_dst.s_addr;
215        pktinfo.ipi_addr.s_addr = self.ipi_addr.s_addr;
216        unsafe { copy_to_bytes(&pktinfo, buffer) }
217    }
218
219    fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
220        let pktinfo: libc::in_pktinfo = unsafe { copy_from_bytes(buffer) }?;
221        Ok(libc::in_pktinfo {
222            ipi_ifindex: pktinfo.ipi_ifindex,
223            ipi_spec_dst: libc::in_addr {
224                s_addr: pktinfo.ipi_spec_dst.s_addr,
225            },
226            ipi_addr: libc::in_addr {
227                s_addr: pktinfo.ipi_addr.s_addr,
228            },
229        })
230    }
231}
232
233#[cfg(unix)]
234impl AncillaryData for libc::in6_pktinfo {
235    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
236        let mut pktinfo: libc::in6_pktinfo = unsafe { std::mem::zeroed() };
237        pktinfo.ipi6_ifindex = self.ipi6_ifindex;
238        pktinfo.ipi6_addr.s6_addr = self.ipi6_addr.s6_addr;
239        unsafe { copy_to_bytes(&pktinfo, buffer) }
240    }
241
242    fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
243        let pktinfo: libc::in6_pktinfo = unsafe { copy_from_bytes(buffer) }?;
244        Ok(libc::in6_pktinfo {
245            ipi6_ifindex: pktinfo.ipi6_ifindex,
246            ipi6_addr: libc::in6_addr {
247                s6_addr: pktinfo.ipi6_addr.s6_addr,
248            },
249        })
250    }
251}
252
253#[cfg(windows)]
254impl AncillaryData for IN_PKTINFO {
255    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
256        let mut pktinfo: IN_PKTINFO = unsafe { std::mem::zeroed() };
257        unsafe {
258            pktinfo.ipi_addr.S_un.S_addr = self.ipi_addr.S_un.S_addr;
259        }
260        pktinfo.ipi_ifindex = self.ipi_ifindex;
261        unsafe { copy_to_bytes(&pktinfo, buffer) }
262    }
263
264    fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
265        let pktinfo: IN_PKTINFO = unsafe { copy_from_bytes(buffer) }?;
266        Ok(IN_PKTINFO {
267            ipi_addr: WinSock::IN_ADDR {
268                S_un: WinSock::IN_ADDR_0 {
269                    S_addr: unsafe { pktinfo.ipi_addr.S_un.S_addr },
270                },
271            },
272            ipi_ifindex: pktinfo.ipi_ifindex,
273        })
274    }
275}
276
277#[cfg(windows)]
278impl AncillaryData for IN6_PKTINFO {
279    fn encode(&self, buffer: &mut [MaybeUninit<u8>]) -> Result<(), CodecError> {
280        let mut pktinfo: IN6_PKTINFO = unsafe { std::mem::zeroed() };
281        unsafe {
282            pktinfo.ipi6_addr.u.Byte = self.ipi6_addr.u.Byte;
283        }
284        pktinfo.ipi6_ifindex = self.ipi6_ifindex;
285        unsafe { copy_to_bytes(&pktinfo, buffer) }
286    }
287
288    fn decode(buffer: &[u8]) -> Result<Self, CodecError> {
289        let pktinfo: IN6_PKTINFO = unsafe { copy_from_bytes(buffer) }?;
290        Ok(IN6_PKTINFO {
291            ipi6_addr: WinSock::IN6_ADDR {
292                u: WinSock::IN6_ADDR_0 {
293                    Byte: unsafe { pktinfo.ipi6_addr.u.Byte },
294                },
295            },
296            ipi6_ifindex: pktinfo.ipi6_ifindex,
297        })
298    }
299}