From 25c407158967460b23d29167c617d6ae1c4de1bd Mon Sep 17 00:00:00 2001 From: Xiaobo Liu Date: Tue, 12 Aug 2025 20:36:21 +0800 Subject: [PATCH] ebpf: add peak() method to Queue and Stack Add integration tests covering push,pop,peek for both types. Signed-off-by: Xiaobo Liu Signed-off-by: Tamir Duberstein --- ebpf/aya-ebpf/src/maps/queue.rs | 10 +- ebpf/aya-ebpf/src/maps/stack.rs | 21 ++++- test/integration-common/src/lib.rs | 5 + test/integration-ebpf/Cargo.toml | 4 + .../src/linear_data_structures.rs | 71 ++++++++++++++ test/integration-test/src/lib.rs | 1 + test/integration-test/src/tests.rs | 1 + .../src/tests/linear_data_structures.rs | 94 +++++++++++++++++++ xtask/public-api/aya-ebpf.txt | 12 ++- 9 files changed, 209 insertions(+), 10 deletions(-) create mode 100644 test/integration-ebpf/src/linear_data_structures.rs create mode 100644 test/integration-test/src/tests/linear_data_structures.rs diff --git a/ebpf/aya-ebpf/src/maps/queue.rs b/ebpf/aya-ebpf/src/maps/queue.rs index 8c8f0bb1..961ee8a1 100644 --- a/ebpf/aya-ebpf/src/maps/queue.rs +++ b/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 Queue { (ret == 0).then_some(value.assume_init()) } } + + pub fn peek(&self) -> Option { + 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()) + } + } } diff --git a/ebpf/aya-ebpf/src/maps/stack.rs b/ebpf/aya-ebpf/src/maps/stack.rs index 6328693d..05cbc2ee 100644 --- a/ebpf/aya-ebpf/src/maps/stack.rs +++ b/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 Stack { } } - 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 Stack { (ret == 0).then_some(()).ok_or(ret) } - pub fn pop(&mut self) -> Option { + pub fn pop(&self) -> Option { 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 { + 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()) diff --git a/test/integration-common/src/lib.rs b/test/integration-common/src/lib.rs index f75a3939..8efed81d 100644 --- a/test/integration-common/src/lib.rs +++ b/test/integration-common/src/lib.rs @@ -64,3 +64,8 @@ pub mod strncmp { #[cfg(feature = "user")] unsafe impl aya::Pod for TestResult {} } + +pub mod linear_data_structures { + pub const PEEK_INDEX: u32 = 0; + pub const POP_INDEX: u32 = 1; +} diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index b0573882..78c8b89e 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -28,6 +28,10 @@ xtask = { path = "../../xtask" } name = "bpf_probe_read" path = "src/bpf_probe_read.rs" +[[bin]] +name = "linear_data_structures" +path = "src/linear_data_structures.rs" + [[bin]] name = "log" path = "src/log.rs" diff --git a/test/integration-ebpf/src/linear_data_structures.rs b/test/integration-ebpf/src/linear_data_structures.rs new file mode 100644 index 00000000..d8f01709 --- /dev/null +++ b/test/integration-ebpf/src/linear_data_structures.rs @@ -0,0 +1,71 @@ +#![no_std] +#![no_main] + +#[cfg(not(test))] +extern crate ebpf_panic; + +use aya_ebpf::{ + cty::c_long, + macros::{map, uprobe}, + maps::{Array, Queue, Stack}, + programs::ProbeContext, +}; +use integration_common::linear_data_structures::{PEEK_INDEX, POP_INDEX}; + +#[map] +static RESULT: Array = Array::::with_max_entries(2, 0); + +#[inline(always)] +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(()) +} + +macro_rules! define_linear_ds_test { + ($Type:ident, $map_ident:ident, + push_fn: $push_fn:ident, + pop_fn: $pop_fn:ident, + peek_fn: $peek_fn:ident, + ) => { + #[map] + static $map_ident: $Type = $Type::with_max_entries(10, 0); + + #[uprobe] + pub fn $push_fn(ctx: ProbeContext) -> Result<(), c_long> { + let value = ctx.arg(0).ok_or(-1)?; + $map_ident.push(&value, 0)?; + Ok(()) + } + + #[uprobe] + pub fn $pop_fn(_: ProbeContext) -> Result<(), c_long> { + let value = $map_ident.pop().ok_or(-1)?; + result_set(POP_INDEX, value)?; + Ok(()) + } + + #[uprobe] + pub fn $peek_fn(_: ProbeContext) -> Result<(), c_long> { + let value = $map_ident.peek().ok_or(-1)?; + result_set(PEEK_INDEX, value)?; + Ok(()) + } + }; +} + +use define_linear_ds_test; + +define_linear_ds_test!(Stack, TEST_STACK, + push_fn: test_stack_push, + pop_fn: test_stack_pop, + peek_fn: test_stack_peek, +); + +define_linear_ds_test!(Queue, TEST_QUEUE, + push_fn: test_queue_push, + pop_fn: test_queue_pop, + peek_fn: test_queue_peek, +); diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index f7cc30c6..9dbb4d18 100644 --- a/test/integration-test/src/lib.rs +++ b/test/integration-test/src/lib.rs @@ -39,6 +39,7 @@ bpf_file!( VARIABLES_RELOC => "variables_reloc.bpf.o", BPF_PROBE_READ => "bpf_probe_read", + LINEAR_DATA_STRUCTURES => "linear_data_structures", LOG => "log", MAP_TEST => "map_test", MEMMOVE_TEST => "memmove_test", diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index 43894366..deaa63a5 100644 --- a/test/integration-test/src/tests.rs +++ b/test/integration-test/src/tests.rs @@ -4,6 +4,7 @@ mod elf; mod feature_probe; mod info; mod iter; +mod linear_data_structures; mod load; mod log; mod raw_tracepoint; diff --git a/test/integration-test/src/tests/linear_data_structures.rs b/test/integration-test/src/tests/linear_data_structures.rs new file mode 100644 index 00000000..49ddbfdb --- /dev/null +++ b/test/integration-test/src/tests/linear_data_structures.rs @@ -0,0 +1,94 @@ +use aya::{EbpfLoader, maps::Array, programs::UProbe}; +use integration_common::linear_data_structures::{PEEK_INDEX, POP_INDEX}; + +enum Order { + Lifo, + Fifo, +} + +macro_rules! define_linear_ds_host_test { + ( + push_prog: $push_prog:literal, + pop_prog: $pop_prog:literal, + peek_prog: $peek_prog:literal, + push_fn: $push_fn:ident, + pop_fn: $pop_fn:ident, + peek_fn: $peek_fn:ident, + test_fn: $test_fn:ident, + order: $order:expr, + ) => { + #[unsafe(no_mangle)] + #[inline(never)] + pub extern "C" fn $push_fn(arg: u64) { + core::hint::black_box(arg); + } + #[unsafe(no_mangle)] + #[inline(never)] + pub extern "C" fn $peek_fn(marker: u64) -> u64 { + core::hint::black_box($peek_fn); + marker + 1 + } + #[unsafe(no_mangle)] + #[inline(never)] + pub extern "C" fn $pop_fn(marker: u64) -> u64 { + core::hint::black_box($pop_fn); + marker + 2 + } + + #[test_log::test] + fn $test_fn() { + let mut bpf = EbpfLoader::new() + .load(crate::LINEAR_DATA_STRUCTURES) + .unwrap(); + for (prog_name, symbol) in [ + ($push_prog, stringify!($push_fn)), + ($peek_prog, stringify!($peek_fn)), + ($pop_prog, stringify!($pop_fn)), + ] { + let prog: &mut UProbe = bpf.program_mut(prog_name).unwrap().try_into().unwrap(); + prog.load().unwrap(); + prog.attach(symbol, "/proc/self/exe", None, None).unwrap(); + } + let array_map = bpf.map("RESULT").unwrap(); + let array = Array::<_, u64>::try_from(array_map).unwrap(); + let seq = 0..9; + for i in seq.clone() { + $push_fn(i); + } + let mut rev = seq.clone().rev(); + let mut seq = seq; + let iter: &mut dyn Iterator = match $order { + Order::Lifo => &mut rev, + Order::Fifo => &mut seq, + }; + for i in iter { + $peek_fn(i); + assert_eq!(array.get(&PEEK_INDEX, 0).unwrap(), i); + $pop_fn(i); + assert_eq!(array.get(&POP_INDEX, 0).unwrap(), i); + } + } + }; +} + +define_linear_ds_host_test!( + push_prog: "test_stack_push", + pop_prog: "test_stack_pop", + peek_prog: "test_stack_peek", + push_fn: trigger_stack_push, + pop_fn: trigger_stack_pop, + peek_fn: trigger_stack_peek, + test_fn: stack_basic, + order: Order::Lifo, +); + +define_linear_ds_host_test!( + push_prog: "test_queue_push", + pop_prog: "test_queue_pop", + peek_prog: "test_queue_peek", + push_fn: trigger_queue_push, + pop_fn: trigger_queue_pop, + peek_fn: trigger_queue_peek, + test_fn: queue_basic, + order: Order::Fifo, +); diff --git a/xtask/public-api/aya-ebpf.txt b/xtask/public-api/aya-ebpf.txt index 2d70495b..b70493ba 100644 --- a/xtask/public-api/aya-ebpf.txt +++ b/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 impl aya_ebpf::maps::queue::Queue +pub fn aya_ebpf::maps::queue::Queue::peek(&self) -> core::option::Option pub const fn aya_ebpf::maps::queue::Queue::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::queue::Queue pub fn aya_ebpf::maps::queue::Queue::pop(&self) -> core::option::Option pub fn aya_ebpf::maps::queue::Queue::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 impl aya_ebpf::maps::stack::Stack +pub fn aya_ebpf::maps::stack::Stack::peek(&self) -> core::option::Option pub const fn aya_ebpf::maps::stack::Stack::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack -pub fn aya_ebpf::maps::stack::Stack::pop(&mut self) -> core::option::Option -pub fn aya_ebpf::maps::stack::Stack::push(&mut self, value: &T, flags: u64) -> core::result::Result<(), i64> +pub fn aya_ebpf::maps::stack::Stack::pop(&self) -> core::option::Option +pub fn aya_ebpf::maps::stack::Stack::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64> pub const fn aya_ebpf::maps::stack::Stack::with_max_entries(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack impl core::marker::Freeze for aya_ebpf::maps::stack::Stack impl core::marker::Send for aya_ebpf::maps::stack::Stack where T: core::marker::Send @@ -1166,6 +1168,7 @@ impl core::convert::From 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 impl aya_ebpf::maps::queue::Queue +pub fn aya_ebpf::maps::queue::Queue::peek(&self) -> core::option::Option pub const fn aya_ebpf::maps::queue::Queue::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::queue::Queue pub fn aya_ebpf::maps::queue::Queue::pop(&self) -> core::option::Option pub fn aya_ebpf::maps::queue::Queue::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64> @@ -1283,9 +1286,10 @@ impl core::convert::From 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 impl aya_ebpf::maps::stack::Stack +pub fn aya_ebpf::maps::stack::Stack::peek(&self) -> core::option::Option pub const fn aya_ebpf::maps::stack::Stack::pinned(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack -pub fn aya_ebpf::maps::stack::Stack::pop(&mut self) -> core::option::Option -pub fn aya_ebpf::maps::stack::Stack::push(&mut self, value: &T, flags: u64) -> core::result::Result<(), i64> +pub fn aya_ebpf::maps::stack::Stack::pop(&self) -> core::option::Option +pub fn aya_ebpf::maps::stack::Stack::push(&self, value: &T, flags: u64) -> core::result::Result<(), i64> pub const fn aya_ebpf::maps::stack::Stack::with_max_entries(max_entries: u32, flags: u32) -> aya_ebpf::maps::stack::Stack impl core::marker::Freeze for aya_ebpf::maps::stack::Stack impl core::marker::Send for aya_ebpf::maps::stack::Stack where T: core::marker::Send