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        let mut done = Bits::with_len(request.chain_len());
33        for i in 0..request.chain_len() {
34            let Some(has_ack) = request.supports_ack(i) else {
35                break;
36            };
37            if !has_ack {
38                done.set(i);
39            }
40        }
41
42        Ok(NetlinkReplyChained {
43            sock,
44            buf: &mut self.buf,
45            request,
46            inner: NetlinkReplyInner {
47                buf_offset: 0,
48                buf_read: 0,
49            },
50            done,
51        })
52    }
53}
54
55pub struct NetlinkReplyChained<'sock> {
56    inner: NetlinkReplyInner,
57    request: &'sock (dyn NetlinkChained + Send + Sync),
58    sock: &'sock mut Socket,
59    buf: &'sock mut Arc<[u8; RECV_BUF_SIZE]>,
60    done: Bits,
61}
62
63impl NetlinkReplyChained<'_> {
64    #[cfg_attr(not(feature = "async"), maybe_async::maybe_async)]
65    pub async fn recv_all(&mut self) -> Result<(), ReplyError> {
66        while let Some(res) = self.recv().await {
67            res?;
68        }
69        Ok(())
70    }
71
72    #[cfg_attr(not(feature = "async"), maybe_async::maybe_async)]
73    pub async fn recv(&mut self) -> Option<Result<(), ReplyError>> {
74        if self.done.is_all() {
75            return None;
76        }
77
78        let buf = Arc::make_mut(self.buf);
79
80        loop {
81            match self.inner.recv(self.sock, buf).await {
82                Err(io_err) => {
83                    self.done.set_all();
84                    return Some(Err(io_err.into()));
85                }
86                Ok((seq, _type, res)) => {
87                    let Some(index) = self.request.get_index(seq) else {
88                        continue;
89                    };
90                    match res {
91                        Ok(_) => return Some(Ok(())),
92                        Err(mut err) => {
93                            if err.code.raw_os_error().unwrap() == 0 {
94                                self.done.set(index);
95                                return Some(Ok(()));
96                            } else {
97                                self.done.set_all();
98                                err.chained_name = Some(self.request.name(index));
99                                if err.has_context() {
100                                    err.lookup = self.request.lookup(index);
101                                    err.reply_buf = Some(self.buf.clone());
102                                }
103                                return Some(Err(err));
104                            };
105                        }
106                    }
107                }
108            };
109        }
110    }
111}
112
113#[derive(Clone)]
114enum Bits {
115    Inline(u64),
116    Vec(Vec<u64>),
117}
118
119impl fmt::Debug for Bits {
120    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121        let n = self.count_zeros();
122        write!(f, "{n} replies pending")
123    }
124}
125
126impl Bits {
127    fn with_len(len: usize) -> Self {
128        if len < 64 {
129            Self::Inline(u64::MAX << (len % 64))
130        } else {
131            let mut vec = vec![0; len.div_ceil(64)];
132            *vec.last_mut().unwrap() |= u64::MAX << (len % 64);
133            Self::Vec(vec)
134        }
135    }
136
137    fn set(&mut self, index: usize) {
138        match self {
139            Self::Inline(w) => *w |= 1u64 << index,
140            Self::Vec(bits) => bits[index / 64] |= 1u64 << (index % 64),
141        }
142    }
143
144    fn is_all(&self) -> bool {
145        match self {
146            Self::Inline(w) => *w == u64::MAX,
147            Self::Vec(bits) => bits.iter().all(|w| *w == u64::MAX),
148        }
149    }
150
151    fn set_all(&mut self) {
152        match self {
153            Self::Inline(w) => *w = u64::MAX,
154            Self::Vec(bits) => bits.iter_mut().for_each(|w| *w = u64::MAX),
155        }
156    }
157
158    fn count_zeros(&self) -> usize {
159        match self {
160            Self::Inline(w) => w.count_zeros() as usize,
161            Self::Vec(bits) => bits.iter().map(|s| s.count_zeros() as usize).sum(),
162        }
163    }
164}
165
166#[cfg(test)]
167mod test {
168    use super::*;
169
170    #[allow(unused)]
171    trait SpawnCompatible: Send {}
172    impl<'a> SpawnCompatible for NetlinkReplyChained<'a> {}
173}
netlink_socket2/chained.rs

netlink_socket2/
chained.rs
