xref: /DragonOS/kernel/src/bpf/map/queue.rs (revision 7b0ef10895108a0de5ff5ef3d2f93f40cf2e33a5)
1 use super::{BpfMapCommonOps, Result};
2 use crate::bpf::map::util::{BpfMapMeta, BpfMapUpdateElemFlags};
3 use alloc::vec::Vec;
4 use core::fmt::Debug;
5 use core::ops::Deref;
6 use core::ops::DerefMut;
7 use system_error::SystemError;
8 
9 type BpfQueueValue = Vec<u8>;
10 /// BPF_MAP_TYPE_QUEUE provides FIFO storage and BPF_MAP_TYPE_STACK provides LIFO storage for BPF programs.
11 /// These maps support peek, pop and push operations that are exposed to BPF programs through the respective helpers.
12 /// These operations are exposed to userspace applications using the existing bpf syscall in the following way:
13 /// - `BPF_MAP_LOOKUP_ELEM` -> `peek`
14 /// - `BPF_MAP_UPDATE_ELEM` -> `push`
15 /// - `BPF_MAP_LOOKUP_AND_DELETE_ELEM ` -> `pop`
16 ///
17 /// See https://docs.kernel.org/bpf/map_queue_stack.html
18 pub trait SpecialMap: Debug + Send + Sync + 'static {
19     /// Returns the number of elements the queue can hold.
20     fn push(&mut self, value: BpfQueueValue, flags: BpfMapUpdateElemFlags) -> Result<()>;
21     /// Removes the first element and returns it.
22     fn pop(&mut self) -> Option<BpfQueueValue>;
23     /// Returns the first element without removing it.
24     fn peek(&self) -> Option<&BpfQueueValue>;
25 }
26 
27 /// The queue map type is a generic map type, resembling a FIFO (First-In First-Out) queue.
28 ///
29 /// This map type has no keys, only values. The size and type of the values can be specified by the user
30 /// to fit a large variety of use cases. The typical use-case for this map type is to keep track of
31 /// a pool of elements such as available network ports when implementing NAT (network address translation).
32 ///
33 /// As apposed to most map types, this map type uses a custom set of helpers to pop, peek and push elements.
34 ///
35 /// See https://ebpf-docs.dylanreimerink.nl/linux/map-type/BPF_MAP_TYPE_QUEUE/
36 #[derive(Debug)]
37 pub struct QueueMap {
38     max_entries: u32,
39     data: Vec<BpfQueueValue>,
40 }
41 
42 impl QueueMap {
43     pub fn new(attr: &BpfMapMeta) -> Result<Self> {
44         if attr.value_size == 0 || attr.max_entries == 0 || attr.key_size != 0 {
45             return Err(SystemError::EINVAL);
46         }
47         let data = Vec::with_capacity(attr.max_entries as usize);
48         Ok(Self {
49             max_entries: attr.max_entries,
50             data,
51         })
52     }
53 }
54 
55 impl SpecialMap for QueueMap {
56     fn push(&mut self, value: BpfQueueValue, flags: BpfMapUpdateElemFlags) -> Result<()> {
57         if self.data.len() == self.max_entries as usize {
58             if flags.contains(BpfMapUpdateElemFlags::BPF_EXIST) {
59                 // remove the first element
60                 self.data.remove(0);
61             } else {
62                 return Err(SystemError::ENOSPC);
63             }
64         }
65         self.data.push(value);
66         Ok(())
67     }
68     fn pop(&mut self) -> Option<BpfQueueValue> {
69         if self.data.is_empty() {
70             return None;
71         }
72         Some(self.data.remove(0))
73     }
74     fn peek(&self) -> Option<&BpfQueueValue> {
75         self.data.first()
76     }
77 }
78 /// The stack map type is a generic map type, resembling a stack data structure.
79 ///
80 /// See https://ebpf-docs.dylanreimerink.nl/linux/map-type/BPF_MAP_TYPE_STACK/
81 #[derive(Debug)]
82 pub struct StackMap(QueueMap);
83 
84 impl StackMap {
85     pub fn new(attr: &BpfMapMeta) -> Result<Self> {
86         QueueMap::new(attr).map(StackMap)
87     }
88 }
89 
90 impl Deref for StackMap {
91     type Target = QueueMap;
92     fn deref(&self) -> &Self::Target {
93         &self.0
94     }
95 }
96 
97 impl DerefMut for StackMap {
98     fn deref_mut(&mut self) -> &mut Self::Target {
99         &mut self.0
100     }
101 }
102 
103 impl SpecialMap for StackMap {
104     fn push(&mut self, value: BpfQueueValue, flags: BpfMapUpdateElemFlags) -> Result<()> {
105         if self.data.len() == self.max_entries as usize {
106             if flags.contains(BpfMapUpdateElemFlags::BPF_EXIST) {
107                 // remove the last element
108                 self.data.pop();
109             } else {
110                 return Err(SystemError::ENOSPC);
111             }
112         }
113         self.data.push(value);
114         Ok(())
115     }
116     fn pop(&mut self) -> Option<BpfQueueValue> {
117         self.data.pop()
118     }
119     fn peek(&self) -> Option<&BpfQueueValue> {
120         self.data.last()
121     }
122 }
123 
124 impl<T: SpecialMap> BpfMapCommonOps for T {
125     /// Equal to [QueueMap::peek]
126     fn lookup_elem(&mut self, _key: &[u8]) -> Result<Option<&[u8]>> {
127         Ok(self.peek().map(|v| v.as_slice()))
128     }
129     /// Equal to [QueueMap::push]
130     fn update_elem(&mut self, _key: &[u8], value: &[u8], flags: u64) -> Result<()> {
131         let flag = BpfMapUpdateElemFlags::from_bits_truncate(flags);
132         self.push(value.to_vec(), flag)
133     }
134     /// Equal to [QueueMap::pop]
135     fn lookup_and_delete_elem(&mut self, _key: &[u8], value: &mut [u8]) -> Result<()> {
136         if let Some(v) = self.pop() {
137             value.copy_from_slice(&v);
138             Ok(())
139         } else {
140             Err(SystemError::ENOENT)
141         }
142     }
143     fn push_elem(&mut self, value: &[u8], flags: u64) -> Result<()> {
144         self.update_elem(&[], value, flags)
145     }
146     fn pop_elem(&mut self, value: &mut [u8]) -> Result<()> {
147         self.lookup_and_delete_elem(&[], value)
148     }
149     fn peek_elem(&self, value: &mut [u8]) -> Result<()> {
150         self.peek()
151             .map(|v| value.copy_from_slice(v))
152             .ok_or(SystemError::ENOENT)
153     }
154 }
155