netlink_socket2/
chained.rs

1use std::{
2    fmt,
3    io::{self, IoSlice},
4    marker::{Send, Sync},
5    sync::Arc,
6};
7
8use netlink_bindings::traits::NetlinkChained;
9
10use crate::{NetlinkReplyInner, NetlinkSocket, ReplyError, Socket, RECV_BUF_SIZE};
11
12impl NetlinkSocket {
13    /// Execute a chained request (experimental)
14    ///
15    /// Some subsystems have special requirements for related requests,
16    /// expecting certain types of messages to be sent within a single write
17    /// operation. For example transactions in nftables subsystem.
18    ///
19    /// Chained requests currently don't support replies carrying data.
20    #[cfg_attr(not(feature = "async"), maybe_async::maybe_async)]
21    pub async fn request_chained<'a, Chained>(
22        &'a mut self,
23        request: &'a Chained,
24    ) -> io::Result<NetlinkReplyChained<'a>>
25    where
26        Chained: NetlinkChained + Send + Sync,
27    {
28        let sock = Self::get_socket_cached(&mut self.sock, request.protonum())?;
29
30        Self::write_buf(sock, &[IoSlice::new(request.payload())]).await?;
31
32        Ok(NetlinkReplyChained {
33            sock,
34            buf: &mut self.buf,
35            request,
36            inner: NetlinkReplyInner {
37                buf_offset: 0,
38                buf_read: 0,
39            },
40            done: Bits::with_len(request.chain_len()),
41        })
42    }
43}
44
45pub struct NetlinkReplyChained<'sock> {
46    inner: NetlinkReplyInner,
47    request: &'sock (dyn NetlinkChained + Send + Sync),
48    sock: &'sock mut Socket,
49    buf: &'sock mut Arc<[u8; RECV_BUF_SIZE]>,
50    done: Bits,
51}
52
53impl NetlinkReplyChained<'_> {
54    #[cfg_attr(not(feature = "async"), maybe_async::maybe_async)]
55    pub async fn recv_all(&mut self) -> Result<(), ReplyError> {
56        while let Some(res) = self.recv().await {
57            res?;
58        }
59        Ok(())
60    }
61
62    #[cfg_attr(not(feature = "async"), maybe_async::maybe_async)]
63    pub async fn recv(&mut self) -> Option<Result<(), ReplyError>> {
64        if self.done.is_all() {
65            return None;
66        }
67
68        let buf = Arc::make_mut(self.buf);
69
70        loop {
71            match self.inner.recv(self.sock, buf).await {
72                Err(io_err) => {
73                    self.done.set_all();
74                    return Some(Err(io_err.into()));
75                }
76                Ok((seq, res)) => {
77                    let Some(index) = self.request.get_index(seq) else {
78                        continue;
79                    };
80                    match res {
81                        Ok(_) => return Some(Ok(())),
82                        Err(mut err) => {
83                            if err.code.raw_os_error().unwrap() == 0 {
84                                self.done.set(index);
85                                return Some(Ok(()));
86                            } else {
87                                self.done.set_all();
88                                err.chained_name = Some(self.request.name(index));
89                                if err.has_context() {
90                                    err.lookup = self.request.lookup(index);
91                                    err.reply_buf = Some(self.buf.clone());
92                                }
93                                return Some(Err(err));
94                            };
95                        }
96                    }
97                }
98            };
99        }
100    }
101}
102
103#[derive(Clone)]
104enum Bits {
105    Inline(u64),
106    Vec(Vec<u64>),
107}
108
109impl fmt::Debug for Bits {
110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111        let n = self.count_zeros();
112        write!(f, "{n} replies pending")
113    }
114}
115
116impl Bits {
117    fn with_len(len: usize) -> Self {
118        if len < 64 {
119            Self::Inline(u64::MAX << (len % 64))
120        } else {
121            let mut vec = vec![0; len.div_ceil(64)];
122            *vec.last_mut().unwrap() |= u64::MAX << (len % 64);
123            Self::Vec(vec)
124        }
125    }
126
127    fn set(&mut self, index: usize) {
128        match self {
129            Self::Inline(w) => *w |= 1u64 << index,
130            Self::Vec(bits) => bits[index / 64] |= 1u64 << (index % 64),
131        }
132    }
133
134    fn is_all(&self) -> bool {
135        match self {
136            Self::Inline(w) => *w == u64::MAX,
137            Self::Vec(bits) => bits.iter().all(|w| *w == u64::MAX),
138        }
139    }
140
141    fn set_all(&mut self) {
142        match self {
143            Self::Inline(w) => *w = u64::MAX,
144            Self::Vec(bits) => bits.iter_mut().for_each(|w| *w = u64::MAX),
145        }
146    }
147
148    fn count_zeros(&self) -> usize {
149        match self {
150            Self::Inline(w) => w.count_zeros() as usize,
151            Self::Vec(bits) => bits.iter().map(|s| s.count_zeros() as usize).sum(),
152        }
153    }
154}
155
156#[cfg(test)]
157mod test {
158    use super::*;
159
160    #[allow(unused)]
161    trait SpawnCompatible: Send {}
162    impl<'a> SpawnCompatible for NetlinkReplyChained<'a> {}
163}
netlink_socket2/chained.rs

netlink_socket2/
chained.rs
