ebpf: add peak() method to Queue and Stack

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

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

@@ -2,7 +2,7 @@ use core::{cell::UnsafeCell, marker::PhantomData, mem};
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_QUEUE},
-    helpers::{bpf_map_pop_elem, bpf_map_push_elem},
+    helpers::{bpf_map_peek_elem, bpf_map_pop_elem, bpf_map_push_elem},
     maps::PinningType,
 };
 
@@ -63,4 +63,12 @@ impl<T> Queue<T> {
             (ret == 0).then_some(value.assume_init())
         }
     }
+
+    pub fn peek(&self) -> Option<T> {
+        unsafe {
+            let mut value = mem::MaybeUninit::uninit();
+            let ret = bpf_map_peek_elem(self.def.get() as *mut _, value.as_mut_ptr() as *mut _);
+            (ret == 0).then_some(value.assume_init())
+        }
+    }
 }

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

@@ -2,7 +2,7 @@ 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},
+    helpers::{bpf_map_peek_elem, bpf_map_pop_elem, bpf_map_push_elem},
     maps::PinningType,
 };
 
@@ -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,22 @@ 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())
+        }
+    }
+
+    pub fn peek(&self) -> Option<T> {
+        unsafe {
+            let mut value = mem::MaybeUninit::uninit();
+            let ret = bpf_map_peek_elem(
+                &self.def as *const _ as *mut _,
                 value.as_mut_ptr() as *mut _,
             );
             (ret == 0).then_some(value.assume_init())

test/integration-ebpf/Cargo.toml

@@ -91,3 +91,11 @@ path = "src/xdp_sec.rs"
 [[bin]]
 name = "uprobe_cookie"
 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/queue_test.rs

@@ -0,0 +1,49 @@
+#![no_std]
+#![no_main]
+
+use aya_ebpf::{
+    cty::c_long,
+    macros::{map, uprobe},
+    maps::{Array, Queue},
+    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_QUEUE: Queue<u64> = Queue::with_max_entries(10, 0);
+
+#[uprobe]
+pub fn test_queue_push(ctx: ProbeContext) -> Result<(), c_long> {
+    let value = ctx.arg(0).ok_or(-1)?;
+    TEST_QUEUE.push(&value, 0)?;
+    Ok(())
+}
+
+#[uprobe]
+pub fn test_queue_pop(_: ProbeContext) -> Result<(), c_long> {
+    let value = TEST_QUEUE.pop().ok_or(-1)?;
+    result_set(POP_INDEX, value)?;
+    Ok(())
+}
+
+#[uprobe]
+pub fn test_queue_peek(_: ProbeContext) -> Result<(), c_long> {
+    let value = TEST_QUEUE.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-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

@@ -49,6 +49,8 @@ bpf_file!(
     RELOCATIONS => "relocations",
     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.rs

@@ -6,11 +6,13 @@ mod info;
 mod iter;
 mod load;
 mod log;
+mod queue;
 mod raw_tracepoint;
 mod rbpf;
 mod relocations;
 mod ring_buf;
 mod smoke;
+mod stack;
 mod strncmp;
 mod tcx;
 mod uprobe_cookie;

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

@@ -0,0 +1,74 @@
+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_queue_push(arg: u64) {
+    core::hint::black_box(arg);
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_queue_peek(marker: u64) -> u64 {
+    core::hint::black_box(trigger_queue_peek);
+    marker + 1
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_queue_pop(marker: u64) -> u64 {
+    core::hint::black_box(trigger_queue_pop);
+    marker + 2
+}
+
+#[test_log::test]
+fn queue_basic() {
+    let mut bpf = EbpfLoader::new().load(crate::QUEUE_TEST).unwrap();
+
+    let prog: &mut UProbe = bpf
+        .program_mut("test_queue_push")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("trigger_queue_push", "/proc/self/exe", None, None)
+        .unwrap();
+
+    let prog: &mut UProbe = bpf
+        .program_mut("test_queue_pop")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("trigger_queue_pop", "/proc/self/exe", None, None)
+        .unwrap();
+
+    let prog: &mut UProbe = bpf
+        .program_mut("test_queue_peek")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    prog.attach("trigger_queue_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_queue_push(i);
+    }
+
+    for i in 0..9 {
+        trigger_queue_peek(i);
+        let peek_value = array.get(&PEEK_INDEX, 0).unwrap();
+
+        trigger_queue_pop(i);
+        let pop_value = array.get(&POP_INDEX, 0).unwrap();
+
+        assert_eq!(peek_value, pop_value);
+        assert_eq!(pop_value, i);
+    }
+}

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

@@ -0,0 +1,71 @@
+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(marker: u64) -> u64 {
+    core::hint::black_box(trigger_stack_peek);
+    marker + 1
+}
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+pub extern "C" fn trigger_stack_pop(marker: u64) -> u64 {
+    core::hint::black_box(trigger_stack_pop);
+    marker + 2
+}
+
+#[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);
+    }
+
+    for i in (0..9).rev() {
+        trigger_stack_peek(i);
+        assert_eq!(array.get(&PEEK_INDEX, 0).unwrap(), i);
+
+        trigger_stack_pop(i);
+        assert_eq!(array.get(&POP_INDEX, 0).unwrap(), i);
+    }
+}

xtask/public-api/aya-ebpf.txt

@@ -438,6 +438,7 @@ pub fn aya_ebpf::maps::program_array::ProgramArray::from(t: T) -> T
 pub mod aya_ebpf::maps::queue
 #[repr(transparent)] pub struct aya_ebpf::maps::queue::Queue<T>
 impl<T> aya_ebpf::maps::queue::Queue<T>
+pub fn aya_ebpf::maps::queue::Queue<T>::peek(&self) -> core::option::Option<T>
 pub const fn aya_ebpf::maps::queue::Queue<T>::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::queue::Queue<T>
 pub fn aya_ebpf::maps::queue::Queue<T>::pop(&self) -> core::option::Option<T>
 pub fn aya_ebpf::maps::queue::Queue<T>::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64>
@@ -592,9 +593,10 @@ pub fn aya_ebpf::maps::sock_map::SockMap::from(t: T) -> T
 pub mod aya_ebpf::maps::stack
 #[repr(transparent)] pub struct aya_ebpf::maps::stack::Stack<T>
 impl<T> aya_ebpf::maps::stack::Stack<T>
+pub fn aya_ebpf::maps::stack::Stack<T>::peek(&self) -> core::option::Option<T>
 pub const fn aya_ebpf::maps::stack::Stack<T>::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack<T>
-pub fn aya_ebpf::maps::stack::Stack<T>::pop(&mut self) -> core::option::Option<T>
-pub fn aya_ebpf::maps::stack::Stack<T>::push(&mut self, value: &T, flags: u64) -> core::result::Result<(), i64>
+pub fn aya_ebpf::maps::stack::Stack<T>::pop(&self) -> core::option::Option<T>
+pub fn aya_ebpf::maps::stack::Stack<T>::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64>
 pub const fn aya_ebpf::maps::stack::Stack<T>::with_max_entries(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack<T>
 impl<T> core::marker::Freeze for aya_ebpf::maps::stack::Stack<T>
 impl<T> core::marker::Send for aya_ebpf::maps::stack::Stack<T> where T: core::marker::Send
@@ -1166,6 +1168,7 @@ impl<T> core::convert::From<T> for aya_ebpf::maps::program_array::ProgramArray
 pub fn aya_ebpf::maps::program_array::ProgramArray::from(t: T) -> T
 #[repr(transparent)] pub struct aya_ebpf::maps::Queue<T>
 impl<T> aya_ebpf::maps::queue::Queue<T>
+pub fn aya_ebpf::maps::queue::Queue<T>::peek(&self) -> core::option::Option<T>
 pub const fn aya_ebpf::maps::queue::Queue<T>::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::queue::Queue<T>
 pub fn aya_ebpf::maps::queue::Queue<T>::pop(&self) -> core::option::Option<T>
 pub fn aya_ebpf::maps::queue::Queue<T>::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64>
@@ -1283,9 +1286,10 @@ impl<T> core::convert::From<T> for aya_ebpf::maps::sock_map::SockMap
 pub fn aya_ebpf::maps::sock_map::SockMap::from(t: T) -> T
 #[repr(transparent)] pub struct aya_ebpf::maps::Stack<T>
 impl<T> aya_ebpf::maps::stack::Stack<T>
+pub fn aya_ebpf::maps::stack::Stack<T>::peek(&self) -> core::option::Option<T>
 pub const fn aya_ebpf::maps::stack::Stack<T>::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack<T>
-pub fn aya_ebpf::maps::stack::Stack<T>::pop(&mut self) -> core::option::Option<T>
-pub fn aya_ebpf::maps::stack::Stack<T>::push(&mut self, value: &T, flags: u64) -> core::result::Result<(), i64>
+pub fn aya_ebpf::maps::stack::Stack<T>::pop(&self) -> core::option::Option<T>
+pub fn aya_ebpf::maps::stack::Stack<T>::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64>
 pub const fn aya_ebpf::maps::stack::Stack<T>::with_max_entries(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack<T>
 impl<T> core::marker::Freeze for aya_ebpf::maps::stack::Stack<T>
 impl<T> core::marker::Send for aya_ebpf::maps::stack::Stack<T> where T: core::marker::Send