mirror of https://github.com/aya-rs/aya
aya-ebpf: Add BTF array definition
Before this change, Aya supported only legacy BPF map definitions, which are instances of the `bpf_map_def` struct and end up in the `maps` ELF section. This change introduces a BTF map definition for arrays, with custom structs indicating the metadata of the map, which end up in the `.maps` section. Co-authored-by: Tamir Duberstein <tamird@gmail.com>main
Michal R
2 weeks ago
parent
e0ceb6214b
commit
0b2a544ddd
@ -0,0 +1,75 @@
|
||||
use std::borrow::Cow;
|
||||
|
||||
use proc_macro2::TokenStream;
|
||||
use quote::quote;
|
||||
use syn::{ItemStatic, Result};
|
||||
|
||||
use crate::args::name_arg;
|
||||
|
||||
pub(crate) struct BtfMap {
|
||||
item: ItemStatic,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl BtfMap {
|
||||
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<BtfMap> {
|
||||
let item: ItemStatic = syn::parse2(item)?;
|
||||
let mut args = syn::parse2(attrs)?;
|
||||
let name = name_arg(&mut args).unwrap_or_else(|| item.ident.to_string());
|
||||
Ok(BtfMap { item, name })
|
||||
}
|
||||
|
||||
pub(crate) fn expand(&self) -> TokenStream {
|
||||
let section_name: Cow<'_, _> = ".maps".into();
|
||||
let name = &self.name;
|
||||
let item = &self.item;
|
||||
quote! {
|
||||
#[unsafe(link_section = #section_name)]
|
||||
#[unsafe(export_name = #name)]
|
||||
#item
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use syn::parse_quote;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_map_with_name() {
|
||||
let map = BtfMap::parse(
|
||||
parse_quote!(name = "foo"),
|
||||
parse_quote!(
|
||||
static BAR: HashMap<&'static str, u32> = HashMap::new();
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
let expanded = map.expand();
|
||||
let expected = quote!(
|
||||
#[unsafe(link_section = ".maps")]
|
||||
#[unsafe(export_name = "foo")]
|
||||
static BAR: HashMap<&'static str, u32> = HashMap::new();
|
||||
);
|
||||
assert_eq!(expected.to_string(), expanded.to_string());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_map_no_name() {
|
||||
let map = BtfMap::parse(
|
||||
parse_quote!(),
|
||||
parse_quote!(
|
||||
static BAR: HashMap<&'static str, u32> = HashMap::new();
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
let expanded = map.expand();
|
||||
let expected = quote!(
|
||||
#[unsafe(link_section = ".maps")]
|
||||
#[unsafe(export_name = "BAR")]
|
||||
static BAR: HashMap<&'static str, u32> = HashMap::new();
|
||||
);
|
||||
assert_eq!(expected.to_string(), expanded.to_string());
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,57 @@
|
||||
use core::{cell::UnsafeCell, ptr::NonNull};
|
||||
|
||||
use crate::{bindings::bpf_map_type::BPF_MAP_TYPE_ARRAY, btf_map_def, cty::c_long, insert, lookup};
|
||||
|
||||
btf_map_def!(ArrayDef, BPF_MAP_TYPE_ARRAY);
|
||||
|
||||
#[repr(transparent)]
|
||||
pub struct Array<T, const M: usize, const F: usize = 0>(UnsafeCell<ArrayDef<u32, T, M, F>>);
|
||||
|
||||
unsafe impl<T: Sync, const M: usize, const F: usize> Sync for Array<T, M, F> {}
|
||||
|
||||
impl<T, const M: usize, const F: usize> Array<T, M, F> {
|
||||
/// Creates a new [`Array`] instance with elements of type `T`, maximum
|
||||
/// capacity of `M` and additional flags `F`.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```rust
|
||||
/// use aya_ebpf::{btf_maps::Array, macros::btf_map};
|
||||
///
|
||||
/// #[btf_map]
|
||||
/// static ARRAY: Array<u32, 10 /* max_elements */, 0> = Array::new();
|
||||
/// ```
|
||||
#[expect(
|
||||
clippy::new_without_default,
|
||||
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||
)]
|
||||
pub const fn new() -> Self {
|
||||
Array(UnsafeCell::new(ArrayDef::new()))
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn get(&self, index: u32) -> Option<&T> {
|
||||
unsafe { self.lookup(index).map(|p| p.as_ref()) }
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn get_ptr(&self, index: u32) -> Option<*const T> {
|
||||
unsafe { self.lookup(index).map(|p| p.as_ptr() as *const T) }
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
pub fn get_ptr_mut(&self, index: u32) -> Option<*mut T> {
|
||||
unsafe { self.lookup(index).map(|p| p.as_ptr()) }
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
unsafe fn lookup(&self, index: u32) -> Option<NonNull<T>> {
|
||||
lookup(self.0.get().cast(), &index)
|
||||
}
|
||||
|
||||
/// Sets the value of the element at the given index.
|
||||
#[inline(always)]
|
||||
pub fn set(&self, index: u32, value: &T, flags: u64) -> Result<(), c_long> {
|
||||
insert(self.0.get().cast(), &index, value, flags)
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,53 @@
|
||||
use core::marker::PhantomData;
|
||||
|
||||
pub mod array;
|
||||
|
||||
pub use array::Array;
|
||||
|
||||
/// A marker used to remove names of annotated types in LLVM debug info and
|
||||
/// therefore also in BTF.
|
||||
#[repr(transparent)]
|
||||
pub(crate) struct AyaBtfMapMarker(PhantomData<()>);
|
||||
|
||||
impl AyaBtfMapMarker {
|
||||
pub(crate) const fn new() -> Self {
|
||||
Self(PhantomData)
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_export]
|
||||
macro_rules! btf_map_def {
|
||||
($name:ident, $t:ident) => {
|
||||
#[expect(
|
||||
dead_code,
|
||||
reason = "These fields exist only for BTF metadata exposure. None of them are actually used."
|
||||
)]
|
||||
pub struct $name<K, V, const M: usize, const F: usize = 0> {
|
||||
r#type: *const [i32; $t as usize],
|
||||
key: *const K,
|
||||
value: *const V,
|
||||
max_entries: *const [i32; M],
|
||||
map_flags: *const [i32; F],
|
||||
|
||||
// Anonymize the struct.
|
||||
_anon: $crate::btf_maps::AyaBtfMapMarker,
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::new_without_default,
|
||||
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||
)]
|
||||
impl<K, V, const M: usize, const F: usize> $name<K, V, M, F> {
|
||||
pub const fn new() -> $name<K, V, M, F> {
|
||||
$name {
|
||||
r#type: &[0i32; $t as usize],
|
||||
key: ::core::ptr::null(),
|
||||
value: ::core::ptr::null(),
|
||||
max_entries: &[0i32; M],
|
||||
map_flags: &[0i32; F],
|
||||
_anon: $crate::btf_maps::AyaBtfMapMarker::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
@ -0,0 +1,90 @@
|
||||
#![no_std]
|
||||
#![no_main]
|
||||
|
||||
#[cfg(not(test))]
|
||||
extern crate ebpf_panic;
|
||||
|
||||
use aya_ebpf::{
|
||||
btf_maps::Array,
|
||||
cty::c_long,
|
||||
macros::{btf_map, map, uprobe},
|
||||
maps::Array as LegacyArray,
|
||||
programs::ProbeContext,
|
||||
};
|
||||
use integration_common::array::{GET_INDEX, GET_PTR_INDEX, GET_PTR_MUT_INDEX};
|
||||
|
||||
#[btf_map]
|
||||
static RESULT: Array<u32, 3 /* max_elements */, 0> = Array::new();
|
||||
#[btf_map]
|
||||
static ARRAY: Array<u32, 10 /* max_elements */, 0> = Array::new();
|
||||
|
||||
#[map]
|
||||
static RESULT_LEGACY: LegacyArray<u32> = LegacyArray::with_max_entries(3, 0);
|
||||
#[map]
|
||||
static ARRAY_LEGACY: LegacyArray<u32> = LegacyArray::with_max_entries(10, 0);
|
||||
|
||||
macro_rules! define_array_test {
|
||||
(
|
||||
$result_map:ident,
|
||||
$array_map:ident,
|
||||
$result_set_fn:ident,
|
||||
$set_prog:ident,
|
||||
$get_prog:ident,
|
||||
$get_ptr_prog:ident,
|
||||
$get_ptr_mut_prog:ident
|
||||
$(,)?
|
||||
) => {
|
||||
#[inline(always)]
|
||||
fn $result_set_fn(index: u32, value: u32) -> Result<(), c_long> {
|
||||
let ptr = $result_map.get_ptr_mut(index).ok_or(-1)?;
|
||||
let dst = unsafe { ptr.as_mut() };
|
||||
let dst_res = dst.ok_or(-1)?;
|
||||
*dst_res = value;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[uprobe]
|
||||
pub fn $set_prog(ctx: ProbeContext) -> Result<(), c_long> {
|
||||
let index = ctx.arg(0).ok_or(-1)?;
|
||||
let value = ctx.arg(1).ok_or(-1)?;
|
||||
$array_map.set(index, &value, 0)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[uprobe]
|
||||
pub fn $get_prog(ctx: ProbeContext) -> Result<(), c_long> {
|
||||
let index = ctx.arg(0).ok_or(-1)?;
|
||||
let value = $array_map.get(index).ok_or(-1)?;
|
||||
$result_set_fn(GET_INDEX, *value)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[uprobe]
|
||||
pub fn $get_ptr_prog(ctx: ProbeContext) -> Result<(), c_long> {
|
||||
let index = ctx.arg(0).ok_or(-1)?;
|
||||
let value = $array_map.get_ptr(index).ok_or(-1)?;
|
||||
$result_set_fn(GET_PTR_INDEX, unsafe { *value })?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[uprobe]
|
||||
pub fn $get_ptr_mut_prog(ctx: ProbeContext) -> Result<(), c_long> {
|
||||
let index = ctx.arg(0).ok_or(-1)?;
|
||||
let ptr = $array_map.get_ptr_mut(index).ok_or(-1)?;
|
||||
let value = unsafe { *ptr };
|
||||
$result_set_fn(GET_PTR_MUT_INDEX, value)?;
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
define_array_test!(RESULT, ARRAY, result_set, set, get, get_ptr, get_ptr_mut);
|
||||
define_array_test!(
|
||||
RESULT_LEGACY,
|
||||
ARRAY_LEGACY,
|
||||
result_set_legacy,
|
||||
set_legacy,
|
||||
get_legacy,
|
||||
get_ptr_legacy,
|
||||
get_ptr_mut_legacy,
|
||||
);
|
||||
@ -0,0 +1,88 @@
|
||||
use aya::{EbpfLoader, maps::Array, programs::UProbe};
|
||||
use integration_common::array::{GET_INDEX, GET_PTR_INDEX, GET_PTR_MUT_INDEX};
|
||||
|
||||
#[unsafe(no_mangle)]
|
||||
#[inline(never)]
|
||||
pub extern "C" fn set(index: u32, value: u32) {
|
||||
std::hint::black_box((index, value));
|
||||
}
|
||||
|
||||
#[unsafe(no_mangle)]
|
||||
#[inline(never)]
|
||||
pub extern "C" fn get(index: u32) {
|
||||
std::hint::black_box(index);
|
||||
}
|
||||
|
||||
#[unsafe(no_mangle)]
|
||||
#[inline(never)]
|
||||
pub extern "C" fn get_ptr(index: u32) {
|
||||
std::hint::black_box(index);
|
||||
}
|
||||
|
||||
#[unsafe(no_mangle)]
|
||||
#[inline(never)]
|
||||
pub extern "C" fn get_ptr_mut(index: u32) {
|
||||
std::hint::black_box(index);
|
||||
}
|
||||
|
||||
#[test_log::test]
|
||||
fn test_array() {
|
||||
let mut ebpf = EbpfLoader::new().load(crate::ARRAY).unwrap();
|
||||
for (result_map, array_map, progs_and_symbols) in [
|
||||
// BTF map definitions.
|
||||
(
|
||||
"RESULT",
|
||||
"ARRAY",
|
||||
[
|
||||
("set", "set"),
|
||||
("get", "get"),
|
||||
("get_ptr", "get_ptr"),
|
||||
("get_ptr_mut", "get_ptr_mut"),
|
||||
],
|
||||
),
|
||||
// Legacy map definitions.
|
||||
(
|
||||
"RESULT_LEGACY",
|
||||
"ARRAY_LEGACY",
|
||||
[
|
||||
("set_legacy", "set"),
|
||||
("get_legacy", "get"),
|
||||
("get_ptr_legacy", "get_ptr"),
|
||||
("get_ptr_mut_legacy", "get_ptr_mut"),
|
||||
],
|
||||
),
|
||||
] {
|
||||
for (prog_name, symbol) in progs_and_symbols {
|
||||
let prog: &mut UProbe = ebpf.program_mut(prog_name).unwrap().try_into().unwrap();
|
||||
prog.load().unwrap();
|
||||
prog.attach(symbol, "/proc/self/exe", None, None).unwrap();
|
||||
}
|
||||
let result_array = ebpf.map(result_map).unwrap();
|
||||
let result_array = Array::<_, u32>::try_from(result_array).unwrap();
|
||||
let array = ebpf.map(array_map).unwrap();
|
||||
let array = Array::<_, u32>::try_from(array).unwrap();
|
||||
let seq = 0..9;
|
||||
for i in seq.clone() {
|
||||
set(i, i.pow(2));
|
||||
}
|
||||
for i in seq.clone() {
|
||||
// Assert the value returned by user-space API.
|
||||
let expected_value = i.pow(2);
|
||||
let value = array.get(&i, 0).unwrap();
|
||||
assert_eq!(value, expected_value);
|
||||
// Assert the value returned by eBPF in-kernel API.
|
||||
get(i);
|
||||
let result = result_array.get(&GET_INDEX, 0).unwrap();
|
||||
assert_eq!(result, expected_value);
|
||||
get_ptr(i);
|
||||
let result = result_array.get(&GET_PTR_INDEX, 0).unwrap();
|
||||
assert_eq!(result, expected_value);
|
||||
}
|
||||
for i in seq.clone() {
|
||||
let value = i.pow(2);
|
||||
get_ptr_mut(i);
|
||||
let result = result_array.get(&GET_PTR_MUT_INDEX, 0).unwrap();
|
||||
assert_eq!(result, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue