From d0b6daa0914348110db63be3ef7350ffb31a99a6 Mon Sep 17 00:00:00 2001
From: Nimrod Shneor <nimrodshn@gmail.com>
Date: Thu, 10 Feb 2022 12:02:15 +0000
Subject: [PATCH] Add LPMTrie Map to aya-bpf for BPF programs

---
 bpf/aya-bpf/src/maps/lpm_trie.rs | 101 +++++++++++++++++++++++++++++++
 bpf/aya-bpf/src/maps/mod.rs      |   2 +
 2 files changed, 103 insertions(+)
 create mode 100644 bpf/aya-bpf/src/maps/lpm_trie.rs

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<K, V> {
+    def: bpf_map_def,
+    _k: PhantomData<K>,
+    _v: PhantomData<V>,
+}
+
+#[repr(packed)]
+pub struct Key<K> {
+    /// Represents the number of bytes matched against.
+    pub prefix_len: u32,
+    /// Represents arbitrary data stored in the LpmTrie.
+    pub data: K,
+}
+
+impl<K> Key<K> {
+    pub fn new(prefix_len: u32, data: K) -> Self {
+        Self { prefix_len, data }
+    }
+}
+
+impl<K, V> LpmTrie<K, V> {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> LpmTrie<K, V> {
+        LpmTrie {
+            def: build_def::<K, V>(BPF_MAP_TYPE_LPM_TRIE, max_entries, flags, PinningType::None),
+            _k: PhantomData,
+            _v: PhantomData,
+        }
+    }
+
+    pub const fn pinned(max_entries: u32, flags: u32) -> LpmTrie<K, V> {
+        LpmTrie {
+            def: build_def::<K, V>(
+                BPF_MAP_TYPE_LPM_TRIE,
+                max_entries,
+                flags,
+                PinningType::ByName,
+            ),
+            _k: PhantomData,
+            _v: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn get(&mut self, key: &Key<K>) -> 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<K>, 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<K>) -> 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<K, V>(ty: u32, max_entries: u32, flags: u32, pin: PinningType) -> bpf_map_def {
+    bpf_map_def {
+        type_: ty,
+        key_size: mem::size_of::<Key<K>>() as u32,
+        value_size: mem::size_of::<V>() 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;