add stack test

Signed-off-by: Xiaobo Liu <cppcoffee@gmail.com>

ebpf/aya-ebpf/src/maps/stack.rs

@@ -43,10 +43,10 @@ impl<T> Stack<T> {
         }
     }
 
-    pub fn push(&mut self, value: &T, flags: u64) -> Result<(), i64> {
+    pub fn push(&self, value: &T, flags: u64) -> Result<(), i64> {
         let ret = unsafe {
             bpf_map_push_elem(
-                &mut self.def as *mut _ as *mut _,
+                &self.def as *const _ as *mut _,
                 value as *const _ as *const _,
                 flags,
             )
@@ -54,11 +54,11 @@ impl<T> Stack<T> {
         (ret == 0).then_some(()).ok_or(ret)
     }
 
-    pub fn pop(&mut self) -> Option<T> {
+    pub fn pop(&self) -> Option<T> {
         unsafe {
             let mut value = mem::MaybeUninit::uninit();
             let ret = bpf_map_pop_elem(
-                &mut self.def as *mut _ as *mut _,
+                &self.def as *const _ as *mut _,
                 value.as_mut_ptr() as *mut _,
             );
             (ret == 0).then_some(value.assume_init())

test/integration-ebpf/Cargo.toml

@@ -95,3 +95,7 @@ path = "src/uprobe_cookie.rs"
 [[bin]]
 name = "queue_test"
 path = "src/queue_test.rs"
+
+[[bin]]
+name = "stack_test"
+path = "src/stack_test.rs"

test/integration-ebpf/src/stack_test.rs

@@ -0,0 +1,49 @@
+#![no_std]
+#![no_main]
+
+use aya_ebpf::{
+    cty::c_long,
+    macros::{map, uprobe},
+    maps::{Array, Stack},
+    programs::ProbeContext,
+};
+#[cfg(not(test))]
+extern crate ebpf_panic;
+
+const PEEK_INDEX: u32 = 0;
+const POP_INDEX: u32 = 1;
+
+#[map]
+static RESULT: Array<u64> = Array::<u64>::with_max_entries(2, 0);
+
+#[map]
+static TEST_STACK: Stack<u64> = Stack::with_max_entries(10, 0);
+
+#[uprobe]
+pub fn test_stack_push(ctx: ProbeContext) -> Result<(), c_long> {
+    let value = ctx.arg(0).ok_or(-1)?;
+    TEST_STACK.push(&value, 0)?;
+    Ok(())
+}
+
+#[uprobe]
+pub fn test_stack_pop(_: ProbeContext) -> Result<(), c_long> {
+    let value = TEST_STACK.pop().ok_or(-1)?;
+    result_set(POP_INDEX, value)?;
+    Ok(())
+}
+
+#[uprobe]
+pub fn test_stack_peek(_: ProbeContext) -> Result<(), c_long> {
+    let value = TEST_STACK.peek().ok_or(-1)?;
+    result_set(PEEK_INDEX, value)?;
+    Ok(())
+}
+
+fn result_set(index: u32, value: u64) -> Result<(), c_long> {
+    let ptr = RESULT.get_ptr_mut(index).ok_or(-1)?;
+    let dst = unsafe { ptr.as_mut() };
+    let dst_res = dst.ok_or(-1)?;
+    *dst_res = value;
+    Ok(())
+}

test/integration-test/src/lib.rs

@@ -50,6 +50,7 @@ bpf_file!(
     RING_BUF => "ring_buf",
     SIMPLE_PROG => "simple_prog",
     QUEUE_TEST => "queue_test",
+    STACK_TEST => "stack_test",
     STRNCMP => "strncmp",
     TCX => "tcx",
     TEST => "test",

test/integration-test/src/tests/stack.rs

@@ -0,0 +1,67 @@
+use aya::{EbpfLoader, maps::Array, programs::UProbe};
+
+const PEEK_INDEX: u32 = 0;
+const POP_INDEX: u32 = 1;
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_stack_push(arg: u64) {
+    core::hint::black_box(arg);
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_stack_peek() {
+    core::hint::black_box(trigger_stack_peek);
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_stack_pop() {
+    core::hint::black_box(trigger_stack_pop);
+}
+
+#[test_log::test]
+fn stack_basic() {
+    let mut bpf = EbpfLoader::new().load(crate::stack_TEST).unwrap();
+
+    let prog: &mut UProbe = bpf
+        .program_mut("test_stack_push")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("trigger_stack_push", "/proc/self/exe", None, None)
+        .unwrap();
+
+    let prog: &mut UProbe = bpf
+        .program_mut("test_stack_pop")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("trigger_stack_pop", "/proc/self/exe", None, None)
+        .unwrap();
+
+    let prog: &mut UProbe = bpf
+        .program_mut("test_stack_peek")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("trigger_stack_peek", "/proc/self/exe", None, None)
+        .unwrap();
+
+    let array_map = bpf.map("RESULT").unwrap();
+    let array = Array::<_, u64>::try_from(array_map).unwrap();
+
+    for i in 0..9 {
+        trigger_stack_push(i);
+
+        trigger_stack_peek();
+        assert_eq!(array.get(&PEEK_INDEX, 0).unwrap(), i);
+
+        trigger_stack_pop();
+        assert_eq!(array.get(&POP_INDEX, 0).unwrap(), i);
+    }
+}