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.
reviewable/pr1340/r3
Michal R 4 days ago
parent efb08a8850
commit d530cc8547

@ -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());
}
}

@ -1,4 +1,5 @@
pub(crate) mod args;
mod btf_map;
mod btf_tracepoint;
mod cgroup_device;
mod cgroup_skb;
@ -24,6 +25,7 @@ mod tracepoint;
mod uprobe;
mod xdp;
use btf_map::BtfMap;
use btf_tracepoint::BtfTracePoint;
use cgroup_device::CgroupDevice;
use cgroup_skb::CgroupSkb;
@ -50,6 +52,15 @@ use tracepoint::TracePoint;
use uprobe::{UProbe, UProbeKind};
use xdp::Xdp;
#[proc_macro_attribute]
pub fn btf_map(attrs: TokenStream, item: TokenStream) -> TokenStream {
match BtfMap::parse(attrs.into(), item.into()) {
Ok(prog) => prog.expand(),
Err(err) => err.into_compile_error(),
}
.into()
}
#[proc_macro_attribute]
pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream {
match Map::parse(attrs.into(), item.into()) {

@ -758,7 +758,7 @@ impl Object {
if type_name == section.name {
// each btf_var_secinfo contains a map
for info in &datasec.entries {
let (map_name, def) = parse_btf_map_def(btf, info)?;
let (map_name, def) = parse_btf_map(btf, info)?;
let symbol_index =
maps.get(&map_name)
.ok_or_else(|| ParseError::SymbolNotFound {
@ -1240,7 +1240,7 @@ fn parse_map_def(name: &str, data: &[u8]) -> Result<bpf_map_def, ParseError> {
}
}
fn parse_btf_map_def(btf: &Btf, info: &DataSecEntry) -> Result<(String, BtfMapDef), BtfError> {
fn parse_btf_map(btf: &Btf, info: &DataSecEntry) -> Result<(String, BtfMapDef), BtfError> {
let ty = match btf.type_by_id(info.btf_type)? {
BtfType::Var(var) => var,
other => {
@ -1250,11 +1250,18 @@ fn parse_btf_map_def(btf: &Btf, info: &DataSecEntry) -> Result<(String, BtfMapDe
}
};
let map_name = btf.string_at(ty.name_offset)?;
let mut map_def = BtfMapDef::default();
// Safety: union
let root_type = btf.resolve_type(ty.btf_type)?;
let s = match btf.type_by_id(root_type)? {
parse_btf_map_def(btf, &map_name, root_type)
}
fn parse_btf_map_def(
btf: &Btf,
map_name: &str,
btf_type_id: u32,
) -> Result<(String, BtfMapDef), BtfError> {
// Safety: union
let s = match btf.type_by_id(btf_type_id)? {
BtfType::Struct(s) => s,
other => {
return Err(BtfError::UnexpectedBtfType {
@ -1263,8 +1270,26 @@ fn parse_btf_map_def(btf: &Btf, info: &DataSecEntry) -> Result<(String, BtfMapDe
}
};
let mut map_def = BtfMapDef::default();
for m in &s.members {
// In aya-ebpf, the BTF map definition types are not used directly.
// Instead, they are wrapped in structs provided by aya-ebpf (e.g.
// `HashMap`, `Array`), which then wrap the definition type in
// `UnsafeCell`.
// To retrieve the definition type, we need to walk through all the
// wrapper types:
//
// - aya-ebpf wrappers like `HashMap`, `Array` etc. They wrap an
// `UnsafeCell` in a tuple struct with one member, hence the field
// name is `__0`.
// - `UnsafeCell`, which wraps the BTF map definition inside a `value`
// field.
match btf.string_at(m.name_offset)?.as_ref() {
"__0" => {
let unsafe_cell_id = btf.resolve_type(m.btf_type)?;
return parse_btf_map_def(btf, map_name, unsafe_cell_id);
}
"type" => {
map_def.map_type = get_map_field(btf, m.btf_type)?;
}
@ -1284,14 +1309,31 @@ fn parse_btf_map_def(btf: &Btf, info: &DataSecEntry) -> Result<(String, BtfMapDe
map_def.key_size = get_map_field(btf, m.btf_type)?;
}
"value" => {
if let BtfType::Ptr(pty) = btf.type_by_id(m.btf_type)? {
let t = pty.btf_type;
map_def.value_size = btf.type_size(t)? as u32;
map_def.btf_value_type_id = t;
} else {
return Err(BtfError::UnexpectedBtfType {
type_id: m.btf_type,
});
// There are two cases to handle:
//
// 1. We are parsing an actual BTF map type with fields like
// `type`, `value`, `key`. In this case, `value` is a
// pointer.
// 2. We are parsing an `UnsafeCell`, which wraps an actual BTF
// map type, and the `value` field of `UnsafeCell` is a
// struct which we want to parse.
match btf.type_by_id(m.btf_type)? {
// BTF map with `value` as a pointer field.
BtfType::Ptr(pty) => {
let t = pty.btf_type;
map_def.value_size = btf.type_size(t)? as u32;
map_def.btf_value_type_id = t;
}
// `UnsafeCell` wrapping a BTF map in a `value field`.
BtfType::Struct(_) => {
let map_type_id = btf.resolve_type(m.btf_type)?;
return parse_btf_map_def(btf, map_name, map_type_id);
}
_ => {
return Err(BtfError::UnexpectedBtfType {
type_id: m.btf_type,
});
}
}
}
"value_size" => {

@ -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();
/// ```
// BPF maps are always used as static variables, therefore this method has
// to be `const`. `Default::default` is not const.
#[allow(clippy::new_without_default)]
pub const fn new() -> Self {
Array(UnsafeCell::new(ArrayDef::new()))
}
#[inline(always)]
pub fn get(&self, index: u32) -> Option<&T> {
// FIXME: alignment
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,50 @@
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) => {
#[allow(dead_code)]
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,
}
// Implementing `Default` makes no sense in this case. Maps are always
// global variables, so they need to be instantiated with a `const`
// method. `Default::default` method is not `const`.
#[allow(clippy::new_without_default)]
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(),
}
}
}
};
}

@ -23,6 +23,7 @@ pub use aya_ebpf_bindings::bindings;
mod args;
pub use args::{PtRegs, RawTracepointArgs};
pub mod btf_maps;
#[expect(clippy::missing_safety_doc, unsafe_op_in_unsafe_fn)]
pub mod helpers;
pub mod maps;

@ -1,5 +1,11 @@
#![no_std]
pub mod array {
pub const GET_INDEX: u32 = 0;
pub const GET_PTR_INDEX: u32 = 1;
pub const GET_PTR_MUT_INDEX: u32 = 2;
}
pub mod bpf_probe_read {
pub const RESULT_BUF_LEN: usize = 1024;

@ -24,6 +24,10 @@ network-types = { workspace = true }
which = { workspace = true, features = ["real-sys"] }
xtask = { path = "../../xtask" }
[[bin]]
name = "array"
path = "src/array.rs"
[[bin]]
name = "bpf_probe_read"
path = "src/bpf_probe_read.rs"

@ -0,0 +1,89 @@
#![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
);

@ -38,6 +38,7 @@ bpf_file!(
TEXT_64_64_RELOC => "text_64_64_reloc.o",
VARIABLES_RELOC => "variables_reloc.bpf.o",
ARRAY => "array",
BPF_PROBE_READ => "bpf_probe_read",
LINEAR_DATA_STRUCTURES => "linear_data_structures",
LOG => "log",

@ -1,3 +1,4 @@
mod array;
mod bpf_probe_read;
mod btf_relocations;
mod elf;

@ -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…
Cancel
Save