diff --git a/bpf/aya-bpf/src/maps/lpm_trie.rs b/bpf/aya-bpf/src/maps/lpm_trie.rs new file mode 100644 index 00000000..41888f77 --- /dev/null +++ b/bpf/aya-bpf/src/maps/lpm_trie.rs @@ -0,0 +1,101 @@ +use core::{marker::PhantomData, mem, ptr::NonNull}; + +use aya_bpf_cty::{c_long, c_void}; + +use crate::{ + bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_LPM_TRIE}, + helpers::{bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem}, + maps::PinningType, +}; + +#[repr(transparent)] +pub struct LpmTrie { + def: bpf_map_def, + _k: PhantomData, + _v: PhantomData, +} + +#[repr(packed)] +pub struct Key { + /// Represents the number of bytes matched against. + pub prefix_len: u32, + /// Represents arbitrary data stored in the LpmTrie. + pub data: K, +} + +impl Key { + pub fn new(prefix_len: u32, data: K) -> Self { + Self { prefix_len, data } + } +} + +impl LpmTrie { + pub const fn with_max_entries(max_entries: u32, flags: u32) -> LpmTrie { + LpmTrie { + def: build_def::(BPF_MAP_TYPE_LPM_TRIE, max_entries, flags, PinningType::None), + _k: PhantomData, + _v: PhantomData, + } + } + + pub const fn pinned(max_entries: u32, flags: u32) -> LpmTrie { + LpmTrie { + def: build_def::( + BPF_MAP_TYPE_LPM_TRIE, + max_entries, + flags, + PinningType::ByName, + ), + _k: PhantomData, + _v: PhantomData, + } + } + + #[inline] + pub fn get(&mut self, key: &Key) -> Option<&V> { + unsafe { + let value = bpf_map_lookup_elem( + &mut self.def as *mut _ as *mut _, + key as *const _ as *const c_void, + ); + // FIXME: alignment + NonNull::new(value as *mut V).map(|p| p.as_ref()) + } + } + + #[inline] + pub fn insert(&mut self, key: &Key, value: &V, flags: u64) -> Result<(), c_long> { + let ret = unsafe { + bpf_map_update_elem( + &mut self.def as *mut _ as *mut _, + key as *const _ as *const _, + value as *const _ as *const _, + flags, + ) + }; + (ret >= 0).then(|| ()).ok_or(ret) + } + + #[inline] + pub fn remove(&mut self, key: &Key) -> Result<(), c_long> { + let ret = unsafe { + bpf_map_delete_elem( + &mut self.def as *mut _ as *mut _, + key as *const _ as *const c_void, + ) + }; + (ret >= 0).then(|| ()).ok_or(ret) + } +} + +const fn build_def(ty: u32, max_entries: u32, flags: u32, pin: PinningType) -> bpf_map_def { + bpf_map_def { + type_: ty, + key_size: mem::size_of::>() as u32, + value_size: mem::size_of::() as u32, + max_entries, + map_flags: flags, + id: 0, + pinning: pin as u32, + } +} diff --git a/bpf/aya-bpf/src/maps/mod.rs b/bpf/aya-bpf/src/maps/mod.rs index c3ae05ed..2d8ba4c9 100644 --- a/bpf/aya-bpf/src/maps/mod.rs +++ b/bpf/aya-bpf/src/maps/mod.rs @@ -7,6 +7,7 @@ pub(crate) enum PinningType { pub mod array; pub mod hash_map; +pub mod lpm_trie; pub mod per_cpu_array; pub mod perf; pub mod program_array; @@ -17,6 +18,7 @@ pub mod stack_trace; pub use array::Array; pub use hash_map::{HashMap, LruHashMap, LruPerCpuHashMap, PerCpuHashMap}; +pub use lpm_trie::LpmTrie; pub use per_cpu_array::PerCpuArray; pub use perf::{PerfEventArray, PerfEventByteArray}; pub use program_array::ProgramArray;