From c57cedd2d46a73fcc4e4253b76e9c1ba55593982 Mon Sep 17 00:00:00 2001
From: Kenjiro Nakayama <nakayamakenjiro@gmail.com>
Date: Sat, 28 May 2022 14:45:29 +0900
Subject: [PATCH] Add missing BPF_MAP_TYPE_STACK in BPF code

Currently BPF_MAP_TYPE_STACK is supported by user code (`aya/src/maps/stack.rs`),
but it does not support the BPF_MAP_TYPE_STACK BPF code.

This patch fixes it.
---
 bpf/aya-bpf/src/maps/mod.rs   |  2 ++
 bpf/aya-bpf/src/maps/stack.rs | 67 +++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 bpf/aya-bpf/src/maps/stack.rs

diff --git a/bpf/aya-bpf/src/maps/mod.rs b/bpf/aya-bpf/src/maps/mod.rs
index fc4f1732..bc94e5f7 100644
--- a/bpf/aya-bpf/src/maps/mod.rs
+++ b/bpf/aya-bpf/src/maps/mod.rs
@@ -14,6 +14,7 @@ pub mod program_array;
 pub mod queue;
 pub mod sock_hash;
 pub mod sock_map;
+pub mod stack;
 pub mod stack_trace;
 
 pub use array::Array;
@@ -25,4 +26,5 @@ pub use program_array::ProgramArray;
 pub use queue::Queue;
 pub use sock_hash::SockHash;
 pub use sock_map::SockMap;
+pub use stack::Stack;
 pub use stack_trace::StackTrace;
diff --git a/bpf/aya-bpf/src/maps/stack.rs b/bpf/aya-bpf/src/maps/stack.rs
new file mode 100644
index 00000000..d78df353
--- /dev/null
+++ b/bpf/aya-bpf/src/maps/stack.rs
@@ -0,0 +1,67 @@
+use core::{marker::PhantomData, mem};
+
+use crate::{
+    bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_STACK},
+    helpers::{bpf_map_pop_elem, bpf_map_push_elem},
+    maps::PinningType,
+};
+
+#[repr(transparent)]
+pub struct Stack<T> {
+    def: bpf_map_def,
+    _t: PhantomData<T>,
+}
+
+impl<T> Stack<T> {
+    pub const fn with_max_entries(max_entries: u32, flags: u32) -> Stack<T> {
+        Stack {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_STACK,
+                key_size: 0,
+                value_size: mem::size_of::<T>() as u32,
+                max_entries,
+                map_flags: flags,
+                id: 0,
+                pinning: PinningType::None as u32,
+            },
+            _t: PhantomData,
+        }
+    }
+
+    pub const fn pinned(max_entries: u32, flags: u32) -> Stack<T> {
+        Stack {
+            def: bpf_map_def {
+                type_: BPF_MAP_TYPE_STACK,
+                key_size: 0,
+                value_size: mem::size_of::<T>() as u32,
+                max_entries,
+                map_flags: flags,
+                id: 0,
+                pinning: PinningType::ByName as u32,
+            },
+            _t: PhantomData,
+        }
+    }
+
+    pub fn push(&mut self, value: &T, flags: u64) -> Result<(), i64> {
+        let ret = unsafe {
+            bpf_map_push_elem(
+                &mut self.def as *mut _ as *mut _,
+                value as *const _ as *const _,
+                flags,
+            )
+        };
+        (ret == 0).then(|| ()).ok_or(ret)
+    }
+
+    pub fn pop(&mut self) -> Option<T> {
+        unsafe {
+            let mut value = mem::MaybeUninit::uninit();
+            let ret = bpf_map_pop_elem(
+                &mut self.def as *mut _ as *mut _,
+                value.as_mut_ptr() as *mut _,
+            );
+            (ret == 0).then(|| value.assume_init())
+        }
+    }
+}