nuclei/buf/
vectored.rs

1//! Owned vectored IO: several of the caller's buffers one call writes from
2//! in order ([`IoVecBuf`]) or reads into in order ([`IoVecBufMut`]), given
3//! back together with its result.
4
5use core::mem::MaybeUninit;
6
7use super::owned::{IoBuf, IoBufMut, OwnedBuf};
8
9mod sealed {
10    use super::OwnedBuf;
11
12    /// The buffers of a vectored call as the runtime holds them.
13    pub struct Bufs(pub(crate) Vec<OwnedBuf>);
14
15    /// How the runtime holds the buffers of a vectored write, and gives
16    /// them back.
17    pub trait Written: Sized {
18        fn into_owned(self) -> Bufs;
19        fn from_owned(bufs: Bufs) -> Self;
20    }
21
22    /// As [`Written`], for a vectored read.
23    pub trait Filled: Sized {
24        fn into_owned(self) -> Bufs;
25        fn from_owned(bufs: Bufs) -> Self;
26    }
27}
28
29use sealed::Bufs;
30
31/// The buffers of an owned vectored write: each one's bytes, in order, go
32/// out as one write. Implemented for `Vec<B>` and `[B; N]` of any [`IoBuf`],
33/// `Vec<Vec<u8>>` among them.
34pub trait IoVecBuf: sealed::Written + Send + 'static {}
35
36/// The buffers of an owned vectored read: the bytes read fill each one's
37/// spare capacity in order. Implemented for `Vec<B>` and `[B; N]` of any
38/// [`IoBufMut`].
39pub trait IoVecBufMut: sealed::Filled + Send + 'static {}
40
41impl<B> sealed::Written for Vec<B>
42where
43    B: IoBuf,
44{
45    fn into_owned(self) -> Bufs {
46        Bufs(self.into_iter().map(OwnedBuf::written).collect())
47    }
48
49    fn from_owned(bufs: Bufs) -> Self {
50        bufs.0.into_iter().map(OwnedBuf::into_written).collect()
51    }
52}
53
54impl<B> IoVecBuf for Vec<B> where B: IoBuf {}
55
56impl<B> sealed::Filled for Vec<B>
57where
58    B: IoBufMut,
59{
60    fn into_owned(self) -> Bufs {
61        Bufs(self.into_iter().map(OwnedBuf::filled).collect())
62    }
63
64    fn from_owned(bufs: Bufs) -> Self {
65        bufs.0.into_iter().map(OwnedBuf::into_filled).collect()
66    }
67}
68
69impl<B> IoVecBufMut for Vec<B> where B: IoBufMut {}
70
71impl<B, const N: usize> sealed::Written for [B; N]
72where
73    B: IoBuf,
74{
75    fn into_owned(self) -> Bufs {
76        Bufs(self.into_iter().map(OwnedBuf::written).collect())
77    }
78
79    fn from_owned(bufs: Bufs) -> Self {
80        let mut bufs = bufs.0.into_iter();
81        core::array::from_fn(|_| bufs.next().map_or_else(B::default, OwnedBuf::into_written))
82    }
83}
84
85impl<B, const N: usize> IoVecBuf for [B; N] where B: IoBuf {}
86
87impl<B, const N: usize> sealed::Filled for [B; N]
88where
89    B: IoBufMut,
90{
91    fn into_owned(self) -> Bufs {
92        Bufs(self.into_iter().map(OwnedBuf::filled).collect())
93    }
94
95    fn from_owned(bufs: Bufs) -> Self {
96        let mut bufs = bufs.0.into_iter();
97        core::array::from_fn(|_| bufs.next().map_or_else(B::default, OwnedBuf::into_filled))
98    }
99}
100
101impl<B, const N: usize> IoVecBufMut for [B; N] where B: IoBufMut {}
102
103/// The buffers of a vectored write, as the runtime holds them.
104pub(crate) fn written<V>(bufs: V) -> Vec<OwnedBuf>
105where
106    V: IoVecBuf,
107{
108    sealed::Written::into_owned(bufs).0
109}
110
111/// The caller's buffers of a vectored write back.
112pub(crate) fn into_written<V>(bufs: Vec<OwnedBuf>) -> V
113where
114    V: IoVecBuf,
115{
116    <V as sealed::Written>::from_owned(Bufs(bufs))
117}
118
119/// The buffers of a vectored read, as the runtime holds them.
120pub(crate) fn filled<V>(bufs: V) -> Vec<OwnedBuf>
121where
122    V: IoVecBufMut,
123{
124    sealed::Filled::into_owned(bufs).0
125}
126
127/// The caller's buffers of a vectored read back, holding what was read.
128pub(crate) fn into_filled<V>(bufs: Vec<OwnedBuf>) -> V
129where
130    V: IoVecBufMut,
131{
132    <V as sealed::Filled>::from_owned(Bufs(bufs))
133}
134
135/// The descriptor of one buffer's bytes or room that the kernel is handed:
136/// an `iovec` on unix, a `WSABUF` on Windows.
137#[cfg(unix)]
138pub(crate) type IoVec = libc::iovec;
139#[cfg(windows)]
140pub(crate) type IoVec = crate::syscore::windows::WSABUF;
141
142/// The most buffers one vectored call hands the kernel: `IOV_MAX` on the
143/// unix systems nuclei runs on. A call over more is a short one, which the
144/// write goes on from and the read answers as it is.
145pub(crate) const IOV_MAX: usize = 1024;
146
147fn iovec(ptr: *const u8, len: usize) -> IoVec {
148    #[cfg(unix)]
149    return libc::iovec {
150        iov_base: ptr.cast_mut().cast(),
151        iov_len: len,
152    };
153    #[cfg(windows)]
154    return IoVec {
155        len: u32::try_from(len).unwrap_or(u32::MAX),
156        buf: ptr.cast_mut(),
157    };
158}
159
160/// The buffers of a vectored call while the runtime holds them: written
161/// from in order, or read into in order, each buffer's room filled before
162/// the next one's.
163pub(crate) struct Vectored {
164    bufs: Vec<OwnedBuf>,
165    /// A read's: how many bytes each buffer held when the call took it; the
166    /// bytes read land in the room after them.
167    held: Option<Vec<usize>>,
168    /// The descriptors of the call out now, rebuilt for each call and kept
169    /// here, so they stay in place while the kernel reads them.
170    iov: Vec<IoVec>,
171    /// The message a ring's vectored send names.
172    #[cfg(ring)]
173    msg: libc::msghdr,
174}
175
176// SAFETY: the descriptors and the message point only into the buffers the
177// value owns, and are read only by the call that is handed them.
178unsafe impl Send for Vectored {}
179// SAFETY: as for `Send`: shared, the value is only read.
180unsafe impl Sync for Vectored {}
181
182impl Vectored {
183    fn new(bufs: Vec<OwnedBuf>, held: Option<Vec<usize>>) -> Self {
184        Self {
185            bufs,
186            held,
187            iov: Vec::new(),
188            // SAFETY: an all-zero `msghdr` is valid (no name, no control data).
189            #[cfg(ring)]
190            msg: unsafe { core::mem::zeroed() },
191        }
192    }
193
194    /// The buffers of a write.
195    pub(crate) fn written(bufs: Vec<OwnedBuf>) -> Self {
196        Self::new(bufs, None)
197    }
198
199    /// The buffers of a read: each one's bytes stay, and its room is filled.
200    pub(crate) fn filled(bufs: Vec<OwnedBuf>) -> Self {
201        let held = bufs.iter().map(OwnedBuf::len).collect();
202        Self::new(bufs, Some(held))
203    }
204
205    pub(crate) fn into_bufs(self) -> Vec<OwnedBuf> {
206        self.bufs
207    }
208
209    /// How many bytes the buffers hold.
210    pub(crate) fn len(&self) -> usize {
211        self.bufs.iter().map(OwnedBuf::len).sum()
212    }
213
214    /// How many bytes the buffers' memory holds: for a write, the bytes they
215    /// hold (nothing is read into it).
216    pub(crate) fn capacity(&self) -> usize {
217        match self.held {
218            Some(_) => self.bufs.iter().map(OwnedBuf::capacity).sum(),
219            None => self.len(),
220        }
221    }
222
223    /// The bytes from the `from`th one they hold to the end of the buffer it
224    /// is in.
225    pub(crate) fn rest_from(&self, from: usize) -> &[u8] {
226        let mut from = from;
227        for buf in &self.bufs {
228            if from < buf.len() {
229                return buf.rest_from(from);
230            }
231            from -= buf.len();
232        }
233        &[]
234    }
235
236    /// The descriptors of the bytes from the `from`th on, in order, for one
237    /// vectored write (at most [`IOV_MAX`] of them); empty past the end.
238    pub(crate) fn write_iovecs(&mut self, from: usize) -> &[IoVec] {
239        self.iov.clear();
240        let mut skip = from;
241        for buf in &self.bufs {
242            if self.iov.len() == IOV_MAX {
243                break;
244            }
245            let bytes = buf.rest_from(skip);
246            skip = skip.saturating_sub(buf.len());
247            if !bytes.is_empty() {
248                self.iov.push(iovec(bytes.as_ptr(), bytes.len()));
249            }
250        }
251        &self.iov
252    }
253
254    /// The descriptors of the buffers' room, in order, for one vectored read
255    /// (at most [`IOV_MAX`] of them): every buffer before the first one with
256    /// room is full.
257    pub(crate) fn read_iovecs(&mut self) -> &mut [IoVec] {
258        self.iov.clear();
259        if self.held.is_some() {
260            for buf in &mut self.bufs {
261                if self.iov.len() == IOV_MAX {
262                    break;
263                }
264                let room = buf.spare_capacity_mut();
265                if !room.is_empty() {
266                    self.iov.push(iovec(room.as_mut_ptr().cast(), room.len()));
267                }
268            }
269        }
270        &mut self.iov
271    }
272
273    /// The message of a ring's vectored send of the bytes from the `from`th
274    /// on: the descriptors of [`write_iovecs`](Self::write_iovecs).
275    #[cfg(ring)]
276    pub(crate) fn send_msg(&mut self, from: usize) -> *const libc::msghdr {
277        let count = self.write_iovecs(from).len();
278        self.msg.msg_iov = self.iov.as_mut_ptr();
279        self.msg.msg_iovlen = count as _;
280        &raw const self.msg
281    }
282
283    /// The message of a ring's vectored receive into the rooms of
284    /// [`read_iovecs`](Self::read_iovecs).
285    #[cfg(ring)]
286    pub(crate) fn recv_msg(&mut self) -> *mut libc::msghdr {
287        let count = self.read_iovecs().len();
288        self.msg.msg_iov = self.iov.as_mut_ptr();
289        self.msg.msg_iovlen = count as _;
290        &raw mut self.msg
291    }
292
293    /// The descriptors the last [`write_iovecs`](Self::write_iovecs) or
294    /// [`read_iovecs`](Self::read_iovecs) built.
295    #[cfg(windows)]
296    pub(crate) fn built(&mut self) -> &mut [IoVec] {
297        &mut self.iov
298    }
299
300    /// The room of the first buffer that has any.
301    pub(crate) fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
302        if self.held.is_none() {
303            return &mut [];
304        }
305        match self.bufs.iter_mut().find(|b| b.len() < b.capacity()) {
306            Some(buf) => buf.spare_capacity_mut(),
307            None => &mut [],
308        }
309    }
310
311    /// Sets how many bytes the buffers hold to `len`: the bytes past what
312    /// they held when the call took them fill each buffer's room in order.
313    ///
314    /// # Safety
315    ///
316    /// `len <= capacity()`, and those bytes were written into the rooms in
317    /// that order.
318    pub(crate) unsafe fn set_len(&mut self, len: usize) {
319        let Some(held) = &self.held else {
320            return;
321        };
322        let mut left = len.saturating_sub(held.iter().sum());
323        for (buf, &held) in self.bufs.iter_mut().zip(held) {
324            let take = left.min(buf.capacity() - held);
325            left -= take;
326            // SAFETY: the caller's: the first `held + take` bytes of this
327            // buffer are initialized, and within its capacity.
328            unsafe { buf.set_len(held + take) };
329        }
330    }
331
332    /// Copies as much of `bytes` as the rooms hold into them in order;
333    /// returns how many.
334    pub(crate) fn append(&mut self, bytes: &[u8]) -> usize {
335        if self.held.is_none() {
336            return 0;
337        }
338        let mut n = 0;
339        for buf in &mut self.bufs {
340            if n == bytes.len() {
341                break;
342            }
343            n += buf.append(&bytes[n..]);
344        }
345        n
346    }
347
348    /// The buffers of a read whose task went away, each with what it held
349    /// before the read: only the buffers that got bytes.
350    pub(crate) fn read_parts(self) -> impl Iterator<Item = (OwnedBuf, usize)> {
351        let held = self.held.unwrap_or_default();
352        self.bufs
353            .into_iter()
354            .zip(held)
355            .filter(|(buf, held)| buf.len() > *held)
356    }
357}
nuclei/buf/vectored.rs

nuclei/buf/
vectored.rs
