mirror of https://github.com/aya-rs/aya
aya::programs::uprobe: add support for cookies
Fixes #1132. Note that this change does not add support in the public API for kprobes or tracepoints, but it's a trivial matter of plumbing. Along the way, the Uprobe::attach API is cleaned up to make the attachment location more coherent. The logic being: if we're going to be breaking the API anyway, may as well clean it up a bit. Furthermore, the aya::sys::bpf_link_attach function is cleaned up by properly modeling the the union in the final field with a rust enum.reviewable/pr1133/r3
Andrew Werner
1 week ago
parent
f34d355d7d
commit
628b7fb022
@ -0,0 +1,26 @@
|
||||
#![no_std]
|
||||
#![no_main]
|
||||
|
||||
use aya_ebpf::{
|
||||
helpers,
|
||||
macros::{map, uprobe},
|
||||
maps::RingBuf,
|
||||
programs::ProbeContext,
|
||||
EbpfContext,
|
||||
};
|
||||
|
||||
#[map]
|
||||
static RING_BUF: RingBuf = RingBuf::with_byte_size(0, 0);
|
||||
|
||||
#[uprobe]
|
||||
pub fn uprobe_cookie(ctx: ProbeContext) {
|
||||
let cookie = unsafe { helpers::bpf_get_attach_cookie(ctx.as_ptr()) };
|
||||
let cookie_bytes = cookie.to_le_bytes();
|
||||
let _res = RING_BUF.output(&cookie_bytes, 0);
|
||||
}
|
||||
|
||||
#[cfg(not(test))]
|
||||
#[panic_handler]
|
||||
fn panic(_info: &core::panic::PanicInfo) -> ! {
|
||||
loop {}
|
||||
}
|
||||
@ -0,0 +1,63 @@
|
||||
use aya::{maps::ring_buf::RingBuf, programs::UProbe, EbpfLoader};
|
||||
use test_log::test;
|
||||
|
||||
#[test]
|
||||
fn test_uprobe_cookie() {
|
||||
const RING_BUF_BYTE_SIZE: u32 = 512; // arbitrary, but big enough
|
||||
|
||||
let mut bpf = EbpfLoader::new()
|
||||
.set_max_entries("RING_BUF", RING_BUF_BYTE_SIZE)
|
||||
.load(crate::UPROBE_COOKIE)
|
||||
.unwrap();
|
||||
let ring_buf = bpf.take_map("RING_BUF").unwrap();
|
||||
let mut ring_buf = RingBuf::try_from(ring_buf).unwrap();
|
||||
let prog: &mut UProbe = bpf
|
||||
.program_mut("uprobe_cookie")
|
||||
.unwrap()
|
||||
.try_into()
|
||||
.unwrap();
|
||||
prog.load().unwrap();
|
||||
const PROG_A: &str = "uprobe_cookie_trigger_ebpf_program_a";
|
||||
const PROG_B: &str = "uprobe_cookie_trigger_ebpf_program_b";
|
||||
let attach = |prog: &mut UProbe, fn_name, cookie| {
|
||||
prog.attach(fn_name, "/proc/self/exe", None, Some(cookie))
|
||||
.unwrap()
|
||||
};
|
||||
|
||||
// Note that the arguments we pass to the functions are meaningless, but we
|
||||
// pass the value we expect to see in the ring buffer from the cookie for
|
||||
// readability.
|
||||
let a = attach(prog, PROG_A, 1);
|
||||
let _b = attach(prog, PROG_B, 2);
|
||||
uprobe_cookie_trigger_ebpf_program_a(1);
|
||||
uprobe_cookie_trigger_ebpf_program_b(2);
|
||||
uprobe_cookie_trigger_ebpf_program_a(1);
|
||||
prog.detach(a).unwrap();
|
||||
let _a = attach(prog, PROG_A, 3);
|
||||
uprobe_cookie_trigger_ebpf_program_a(3);
|
||||
const EXP: &[u64] = &[1, 2, 1, 3];
|
||||
|
||||
let mut seen = Vec::new();
|
||||
while let Some(read) = ring_buf.next() {
|
||||
let read = read.as_ref();
|
||||
match read.try_into() {
|
||||
Ok(read) => seen.push(u64::from_le_bytes(read)),
|
||||
Err(std::array::TryFromSliceError { .. }) => {
|
||||
panic!("invalid ring buffer data: {read:x?}")
|
||||
}
|
||||
}
|
||||
}
|
||||
assert_eq!(seen, EXP);
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
#[inline(never)]
|
||||
pub extern "C" fn uprobe_cookie_trigger_ebpf_program_a(arg: u64) {
|
||||
std::hint::black_box(arg);
|
||||
}
|
||||
|
||||
#[no_mangle]
|
||||
#[inline(never)]
|
||||
pub extern "C" fn uprobe_cookie_trigger_ebpf_program_b(arg: u32) {
|
||||
std::hint::black_box(arg);
|
||||
}
|
||||
Loading…
Reference in New Issue