kbpf_basic/map/
queue.rs

1use alloc::vec::Vec;
2use core::fmt::Debug;
3
4use super::{BpfMapCommonOps, BpfMapMeta, BpfMapUpdateElemFlags};
5use crate::{BpfError, Result};
6
7type BpfQueueValue = Vec<u8>;
8
9/// BPF_MAP_TYPE_QUEUE provides FIFO storage and BPF_MAP_TYPE_STACK provides LIFO storage for BPF programs.
10/// These maps support peek, pop and push operations that are exposed to BPF programs through the respective helpers.
11/// These operations are exposed to userspace applications using the existing bpf syscall in the following way:
12/// - `BPF_MAP_LOOKUP_ELEM` -> `peek`
13/// - `BPF_MAP_UPDATE_ELEM` -> `push`
14/// - `BPF_MAP_LOOKUP_AND_DELETE_ELEM ` -> `pop`
15///
16/// See <https://docs.kernel.org/bpf/map_queue_stack.html>
17pub trait SpecialMap: Debug + Send + Sync + 'static {
18    /// Returns the number of elements the queue can hold.
19    fn push(&mut self, value: BpfQueueValue, flags: BpfMapUpdateElemFlags) -> Result<()>;
20    /// Removes the first element and returns it.
21    fn pop(&mut self) -> Option<BpfQueueValue>;
22    /// Returns the first element without removing it.
23    fn peek(&self) -> Option<&BpfQueueValue>;
24    /// Returns the fixed value size of each entry.
25    fn value_size(&self) -> usize;
26    /// Get the memory usage of the map.
27    fn mem_usage(&self) -> Result<usize>;
28}
29
30/// The queue map type is a generic map type, resembling a FIFO (First-In First-Out) queue.
31///
32/// This map type has no keys, only values. The size and type of the values can be specified by the user
33/// to fit a large variety of use cases. The typical use-case for this map type is to keep track of
34/// a pool of elements such as available network ports when implementing NAT (network address translation).
35///
36/// As apposed to most map types, this map type uses a custom set of helpers to pop, peek and push elements.
37///
38/// See <https://ebpf-docs.dylanreimerink.nl/linux/map-type/BPF_MAP_TYPE_QUEUE/>
39#[derive(Debug)]
40pub struct QueueMap {
41    max_entries: u32,
42    value_size: u32,
43    data: Vec<BpfQueueValue>,
44}
45
46impl QueueMap {
47    pub fn new(map_meta: &BpfMapMeta) -> Result<Self> {
48        if map_meta.value_size == 0 || map_meta.max_entries == 0 || map_meta.key_size != 0 {
49            return Err(BpfError::EINVAL);
50        }
51        let data = Vec::with_capacity(map_meta.max_entries as usize);
52        Ok(Self {
53            max_entries: map_meta.max_entries,
54            value_size: map_meta.value_size,
55            data,
56        })
57    }
58}
59
60impl SpecialMap for QueueMap {
61    fn push(&mut self, value: BpfQueueValue, flags: BpfMapUpdateElemFlags) -> Result<()> {
62        if flags != BpfMapUpdateElemFlags::empty() {
63            return Err(BpfError::EINVAL);
64        }
65        if self.data.len() == self.max_entries as usize {
66            self.data.remove(0);
67        }
68        self.data.push(value);
69        Ok(())
70    }
71
72    fn pop(&mut self) -> Option<BpfQueueValue> {
73        if self.data.is_empty() {
74            return None;
75        }
76        Some(self.data.remove(0))
77    }
78
79    fn peek(&self) -> Option<&BpfQueueValue> {
80        self.data.first()
81    }
82
83    fn value_size(&self) -> usize {
84        self.value_size as usize
85    }
86
87    fn mem_usage(&self) -> Result<usize> {
88        let mut total = 0;
89        for v in &self.data {
90            total += v.len();
91        }
92        Ok(total)
93    }
94}
95/// The stack map type is a generic map type, resembling a stack data structure.
96///
97/// See <https://ebpf-docs.dylanreimerink.nl/linux/map-type/BPF_MAP_TYPE_STACK/>
98#[derive(Debug)]
99pub struct StackMap(QueueMap);
100
101impl StackMap {
102    /// Create a new [StackMap] with the given key size, value size, and maximum number of entries.
103    pub fn new(map_meta: &BpfMapMeta) -> Result<Self> {
104        QueueMap::new(map_meta).map(StackMap)
105    }
106}
107
108impl SpecialMap for StackMap {
109    fn push(&mut self, value: BpfQueueValue, flags: BpfMapUpdateElemFlags) -> Result<()> {
110        if self.0.data.len() == self.0.max_entries as usize {
111            if flags.contains(BpfMapUpdateElemFlags::BPF_EXISTS) {
112                // remove the last element
113                self.0.data.pop();
114            } else {
115                return Err(BpfError::ENOMEM);
116            }
117        }
118        self.0.data.push(value);
119        Ok(())
120    }
121
122    fn pop(&mut self) -> Option<BpfQueueValue> {
123        self.0.data.pop()
124    }
125
126    fn peek(&self) -> Option<&BpfQueueValue> {
127        self.0.data.last()
128    }
129
130    fn value_size(&self) -> usize {
131        self.0.value_size()
132    }
133
134    fn mem_usage(&self) -> Result<usize> {
135        self.0.mem_usage()
136    }
137}
138
139impl<T: SpecialMap> BpfMapCommonOps for T {
140    /// Equal to QueueMap::peek
141    fn lookup_elem(&mut self, key: &[u8]) -> Result<Option<&[u8]>> {
142        if !key.is_empty() {
143            return Err(BpfError::EINVAL);
144        }
145        Ok(self.peek().map(|v| v.as_slice()))
146    }
147
148    /// Equal to QueueMap::push
149    fn update_elem(&mut self, key: &[u8], value: &[u8], flags: u64) -> Result<()> {
150        if !key.is_empty() || value.len() != self.value_size() {
151            return Err(BpfError::EINVAL);
152        }
153        let flag = BpfMapUpdateElemFlags::from_bits(flags).ok_or(BpfError::EINVAL)?;
154        if flag.contains(BpfMapUpdateElemFlags::BPF_F_LOCK)
155            || (flag.contains(BpfMapUpdateElemFlags::BPF_NOEXIST)
156                && flag.contains(BpfMapUpdateElemFlags::BPF_EXISTS))
157        {
158            return Err(BpfError::EINVAL);
159        }
160        self.push(value.to_vec(), flag)
161    }
162
163    /// Equal to QueueMap::pop
164    fn lookup_and_delete_elem(&mut self, key: &[u8], value: &mut [u8]) -> Result<()> {
165        if !key.is_empty() || value.len() != self.value_size() {
166            return Err(BpfError::EINVAL);
167        }
168        if let Some(v) = self.pop() {
169            value.copy_from_slice(&v);
170            Ok(())
171        } else {
172            Err(BpfError::ENOENT)
173        }
174    }
175
176    fn push_elem(&mut self, value: &[u8], flags: u64) -> Result<()> {
177        self.update_elem(&[], value, flags)
178    }
179
180    fn pop_elem(&mut self, value: &mut [u8]) -> Result<()> {
181        self.lookup_and_delete_elem(&[], value)
182    }
183
184    fn peek_elem(&self, value: &mut [u8]) -> Result<()> {
185        if value.len() != self.value_size() {
186            return Err(BpfError::EINVAL);
187        }
188        self.peek()
189            .map(|v| value.copy_from_slice(v))
190            .ok_or(BpfError::ENOENT)
191    }
192
193    fn map_mem_usage(&self) -> Result<usize> {
194        self.mem_usage()
195    }
196
197    fn as_any(&self) -> &dyn core::any::Any {
198        self
199    }
200
201    fn as_any_mut(&mut self) -> &mut dyn core::any::Any {
202        self
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use super::QueueMap;
209    use crate::{
210        BpfError,
211        map::{BpfMapCommonOps, BpfMapMeta},
212    };
213
214    #[test]
215    fn test_queue_validation() {
216        let mut meta = BpfMapMeta::default();
217        meta.key_size = 0;
218        meta.value_size = 4;
219        meta.max_entries = 2;
220
221        let mut queue = QueueMap::new(&meta).unwrap();
222
223        assert_eq!(queue.update_elem(b"x", b"abcd", 0), Err(BpfError::EINVAL));
224        assert_eq!(queue.update_elem(&[], b"abc", 0), Err(BpfError::EINVAL));
225        assert_eq!(queue.update_elem(&[], b"abcd", 8), Err(BpfError::EINVAL));
226        assert_eq!(queue.update_elem(&[], b"abcd", 0), Ok(()));
227
228        let mut value = [0; 3];
229        assert_eq!(queue.peek_elem(&mut value), Err(BpfError::EINVAL));
230    }
231}
kbpf_basic/map/queue.rs

kbpf_basic/map/
queue.rs
