1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

use std::io;

use crate::{ring::XskRingCons, umem::frame::FrameDesc};

use super::{fd::Fd, Socket};

/// The receiving side of an AF_XDP [`Socket`].
///
/// More details can be found in the
/// [docs](https://www.kernel.org/doc/html/latest/networking/af_xdp.html#rx-ring).
#[derive(Debug)]
pub struct RxQueue {
    ring: XskRingCons,
    socket: Socket,
}

impl RxQueue {
    pub(super) fn new(ring: XskRingCons, socket: Socket) -> Self {
        Self { ring, socket }
    }

    /// Update `descs` with information on which [`Umem`] frames have
    /// received packets. Returns the number of elements of `descs`
    /// which have been updated.
    ///
    /// The number of entries updated will be less than or equal to
    /// the length of `descs`. Entries will be updated sequentially
    /// from the start of `descs` until the end.
    ///
    /// Once the contents of the consumed frames have been dealt with
    /// and are no longer required, the frames should eventually be
    /// added back on to either the [`FillQueue`] or the [`TxQueue`].
    ///
    /// # Safety
    ///
    /// The frames passed to this queue must belong to the same
    /// [`Umem`] that this `RxQueue` instance is tied to.
    ///
    /// [`Umem`]: crate::Umem
    /// [`FillQueue`]: crate::FillQueue
    /// [`TxQueue`]: crate::TxQueue
    #[inline]
    pub unsafe fn consume(&mut self, descs: &mut [FrameDesc]) -> usize {
        let nb = descs.len() as u64;

        if nb == 0 {
            return 0;
        }

        let mut idx = 0;

        let cnt = unsafe { libbpf_sys::_xsk_ring_cons__peek(self.ring.as_mut(), nb, &mut idx) };

        if cnt > 0 {
            for desc in descs.iter_mut().take(cnt as usize) {
                let recv_pkt_desc =
                    unsafe { libbpf_sys::_xsk_ring_cons__rx_desc(self.ring.as_ref(), idx) };

                unsafe {
                    desc.addr = (*recv_pkt_desc).addr as usize;
                    desc.lengths.data = (*recv_pkt_desc).len as usize;
                    desc.lengths.headroom = 0;
                    desc.options = (*recv_pkt_desc).options;
                }

                idx += 1;
            }

            unsafe { libbpf_sys::_xsk_ring_cons__release(self.ring.as_mut(), cnt) };
        }

        cnt as usize
    }

    /// Same as [`consume`] but for a single frame descriptor.
    ///
    /// # Safety
    ///
    /// See [`consume`].
    ///
    /// [`consume`]: Self::consume
    #[inline]
    pub unsafe fn consume_one(&mut self, desc: &mut FrameDesc) -> usize {
        let mut idx = 0;

        let cnt = unsafe { libbpf_sys::_xsk_ring_cons__peek(self.ring.as_mut(), 1, &mut idx) };

        if cnt > 0 {
            let recv_pkt_desc =
                unsafe { libbpf_sys::_xsk_ring_cons__rx_desc(self.ring.as_ref(), idx) };

            unsafe {
                desc.addr = (*recv_pkt_desc).addr as usize;
                desc.lengths.data = (*recv_pkt_desc).len as usize;
                desc.lengths.headroom = 0;
                desc.options = (*recv_pkt_desc).options;
            }

            unsafe { libbpf_sys::_xsk_ring_cons__release(self.ring.as_mut(), cnt) };
        }

        cnt as usize
    }

    /// Same as [`consume`] but poll first to check if there is
    /// anything to read beforehand.
    ///
    /// # Safety
    ///
    /// See [`consume`].
    ///
    /// [`consume`]: RxQueue::consume
    #[inline]
    pub unsafe fn poll_and_consume(
        &mut self,
        descs: &mut [FrameDesc],
        poll_timeout: i32,
    ) -> io::Result<usize> {
        match self.poll(poll_timeout)? {
            true => Ok(unsafe { self.consume(descs) }),
            false => Ok(0),
        }
    }

    /// Same as [`poll_and_consume`] but for a single frame descriptor.
    ///
    /// # Safety
    ///
    /// See [`consume`].
    ///
    /// [`poll_and_consume`]: Self::poll_and_consume
    /// [`consume`]: Self::consume
    #[inline]
    pub unsafe fn poll_and_consume_one(
        &mut self,
        desc: &mut FrameDesc,
        poll_timeout: i32,
    ) -> io::Result<usize> {
        match self.poll(poll_timeout)? {
            true => Ok(unsafe { self.consume_one(desc) }),
            false => Ok(0),
        }
    }

    /// Polls the socket, returning `true` if there is data to read.
    #[inline]
    pub fn poll(&mut self, poll_timeout: i32) -> io::Result<bool> {
        self.socket.fd.poll_read(poll_timeout)
    }

    /// A reference to the underlying [`Socket`]'s file descriptor.
    #[inline]
    pub fn fd(&self) -> &Fd {
        &self.socket.fd
    }

    /// A mutable reference to the underlying [`Socket`]'s file descriptor.
    #[inline]
    pub fn fd_mut(&mut self) -> &mut Fd {
        &mut self.socket.fd
    }
}