From 0ebee15a12803ea3986e287319d182dca7a424ed Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Mon, 11 Oct 2021 15:54:28 +0100
Subject: [PATCH] bpf: Implement Array/Hash of Maps

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf/src/maps/array_of_maps.rs | 57 +++++++++++++++++++++++++
 bpf/aya-bpf/src/maps/hash_of_maps.rs  | 60 +++++++++++++++++++++++++++
 bpf/aya-bpf/src/maps/mod.rs           |  4 ++
 3 files changed, 121 insertions(+)
 create mode 100644 bpf/aya-bpf/src/maps/array_of_maps.rs
 create mode 100644 bpf/aya-bpf/src/maps/hash_of_maps.rs

diff --git a/bpf/aya-bpf/src/maps/array_of_maps.rs b/bpf/aya-bpf/src/maps/array_of_maps.rs
new file mode 100644
index 00000000..3b381191
--- /dev/null
+++ b/bpf/aya-bpf/src/maps/array_of_maps.rs
@@ -0,0 +1,57 @@
+use core::mem;
+
+use aya_bpf_cty::c_void;
+
+use crate::{
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_ARRAY_OF_MAPS},
+    helpers::bpf_map_lookup_elem,
+    maps::PinningType,
+};
+
+#[repr(transparent)]
+pub struct ArrayOfMaps {
+    def: bpf_map_def,
+}
+
+impl ArrayOfMaps {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> ArrayOfMaps {
+        ArrayOfMaps {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_ARRAY_OF_MAPS,
+                key_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<u32>() as u32,
+                max_entries,
+                map_flags: flags,
+                id: 0,
+                pinning: PinningType::None as u32,
+            },
+        }
+    }
+
+    pub const fn pinned(max_entries: u32, flags: u32) -> ArrayOfMaps {
+        ArrayOfMaps {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_ARRAY_OF_MAPS,
+                key_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<u32>() as u32,
+                max_entries,
+                map_flags: flags,
+                id: 0,
+                pinning: PinningType::ByName as u32,
+            },
+        }
+    }
+
+    pub unsafe fn get(&mut self, index: u32) -> Option<&u32> {
+        let value = bpf_map_lookup_elem(
+            &mut self.def as *mut _ as *mut _,
+            &index as *const _ as *const c_void,
+        );
+        if value.is_null() {
+            None
+        } else {
+            // FIXME: alignment
+            Some(&*(value as *const u32))
+        }
+    }
+}
diff --git a/bpf/aya-bpf/src/maps/hash_of_maps.rs b/bpf/aya-bpf/src/maps/hash_of_maps.rs
new file mode 100644
index 00000000..5299ca72
--- /dev/null
+++ b/bpf/aya-bpf/src/maps/hash_of_maps.rs
@@ -0,0 +1,60 @@
+use core::{marker::PhantomData, mem};
+
+use aya_bpf_cty::c_void;
+
+use crate::{
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_HASH_OF_MAPS},
+    helpers::bpf_map_lookup_elem,
+    maps::PinningType,
+};
+
+#[repr(transparent)]
+pub struct HashOfMaps<K> {
+    def: bpf_map_def,
+    _k: PhantomData<K>,
+}
+
+impl<K> HashOfMaps<K> {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> HashOfMaps<K> {
+        HashOfMaps {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_HASH_OF_MAPS,
+                key_size: mem::size_of::<K>() as u32,
+                value_size: mem::size_of::<u32>() as u32,
+                max_entries,
+                map_flags: flags,
+                id: 0,
+                pinning: PinningType::None as u32,
+            },
+            _k: PhantomData,
+        }
+    }
+
+    pub const fn pinned(max_entries: u32, flags: u32) -> HashOfMaps<K> {
+        HashOfMaps {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_HASH_OF_MAPS,
+                key_size: mem::size_of::<K>() as u32,
+                value_size: mem::size_of::<u32>() as u32,
+                max_entries,
+                map_flags: flags,
+                id: 0,
+                pinning: PinningType::ByName as u32,
+            },
+            _k: PhantomData,
+        }
+    }
+
+    pub unsafe fn get(&mut self, key: &K) -> Option<&u32> {
+        let value = bpf_map_lookup_elem(
+            &mut self.def as *mut _ as *mut _,
+            key as *const _ as *const c_void,
+        );
+        if value.is_null() {
+            None
+        } else {
+            // FIXME: alignment
+            Some(&*(value as *const u32))
+        }
+    }
+}
diff --git a/bpf/aya-bpf/src/maps/mod.rs b/bpf/aya-bpf/src/maps/mod.rs
index 0ba88308..93bb99ae 100644
--- a/bpf/aya-bpf/src/maps/mod.rs
+++ b/bpf/aya-bpf/src/maps/mod.rs
@@ -6,7 +6,9 @@ pub(crate) enum PinningType {
 }
 
 pub mod array;
+pub mod array_of_maps;
 pub mod hash_map;
+pub mod hash_of_maps;
 pub mod per_cpu_array;
 pub mod perf;
 pub mod queue;
@@ -15,7 +17,9 @@ pub mod sock_map;
 pub mod stack_trace;
 
 pub use array::Array;
+pub use array_of_maps::ArrayOfMaps;
 pub use hash_map::HashMap;
+pub use hash_of_maps::HashOfMaps;
 pub use per_cpu_array::PerCpuArray;
 pub use perf::{PerfEventArray, PerfEventByteArray};
 pub use queue::Queue;