From b8191195af276fc9c8b481c9f164b60626bdc153 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 Signed-off-by: Xiaobo Liu --- ebpf/aya-ebpf/src/maps/queue.rs | 10 +++- ebpf/aya-ebpf/src/maps/stack.rs | 21 +++++-- test/integration-ebpf/Cargo.toml | 8 +++ test/integration-ebpf/src/queue_test.rs | 49 ++++++++++++++++ test/integration-ebpf/src/stack_test.rs | 49 ++++++++++++++++ test/integration-test/src/lib.rs | 2 + test/integration-test/src/tests.rs | 2 + test/integration-test/src/tests/queue.rs | 74 ++++++++++++++++++++++++ test/integration-test/src/tests/stack.rs | 71 +++++++++++++++++++++++ xtask/public-api/aya-ebpf.txt | 12 ++-- 10 files changed, 288 insertions(+), 10 deletions(-) create mode 100644 test/integration-ebpf/src/queue_test.rs create mode 100644 test/integration-ebpf/src/stack_test.rs create mode 100644 test/integration-test/src/tests/queue.rs create mode 100644 test/integration-test/src/tests/stack.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-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index b0573882..b37742e8 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/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" diff --git a/test/integration-ebpf/src/queue_test.rs b/test/integration-ebpf/src/queue_test.rs new file mode 100644 index 00000000..968894e9 --- /dev/null +++ b/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 = Array::::with_max_entries(2, 0); + +#[map] +static TEST_QUEUE: Queue = 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(()) +} diff --git a/test/integration-ebpf/src/stack_test.rs b/test/integration-ebpf/src/stack_test.rs new file mode 100644 index 00000000..6aca84e0 --- /dev/null +++ b/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 = Array::::with_max_entries(2, 0); + +#[map] +static TEST_STACK: Stack = 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(()) +} diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index f7cc30c6..a3d3875a 100644 --- a/test/integration-test/src/lib.rs +++ b/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", diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index 43894366..67621aa0 100644 --- a/test/integration-test/src/tests.rs +++ b/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; diff --git a/test/integration-test/src/tests/queue.rs b/test/integration-test/src/tests/queue.rs new file mode 100644 index 00000000..e8614ccd --- /dev/null +++ b/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); + } +} diff --git a/test/integration-test/src/tests/stack.rs b/test/integration-test/src/tests/stack.rs new file mode 100644 index 00000000..4547b6f5 --- /dev/null +++ b/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); + } +} 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