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

netlink_socket2/
chained.rs
