diff --git a/aya-ebpf-macros/src/btf_map.rs b/aya-ebpf-macros/src/btf_map.rs new file mode 100644 index 00000000..15bf5e95 --- /dev/null +++ b/aya-ebpf-macros/src/btf_map.rs @@ -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 { + 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()); + } +} diff --git a/aya-ebpf-macros/src/lib.rs b/aya-ebpf-macros/src/lib.rs index d5811d53..91988ea9 100644 --- a/aya-ebpf-macros/src/lib.rs +++ b/aya-ebpf-macros/src/lib.rs @@ -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()) { diff --git a/aya-obj/src/obj.rs b/aya-obj/src/obj.rs index ece738ef..0d6fe314 100644 --- a/aya-obj/src/obj.rs +++ b/aya-obj/src/obj.rs @@ -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 { } } -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" => { diff --git a/ebpf/aya-ebpf/src/btf_maps/array.rs b/ebpf/aya-ebpf/src/btf_maps/array.rs new file mode 100644 index 00000000..8765cd01 --- /dev/null +++ b/ebpf/aya-ebpf/src/btf_maps/array.rs @@ -0,0 +1,58 @@ +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(UnsafeCell>); + +unsafe impl Sync for Array {} + +impl Array { + /// 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 = 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> { + 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) + } +} diff --git a/ebpf/aya-ebpf/src/btf_maps/mod.rs b/ebpf/aya-ebpf/src/btf_maps/mod.rs new file mode 100644 index 00000000..00c383b2 --- /dev/null +++ b/ebpf/aya-ebpf/src/btf_maps/mod.rs @@ -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 { + 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 $name { + pub const fn new() -> $name { + $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(), + } + } + } + }; +} diff --git a/ebpf/aya-ebpf/src/lib.rs b/ebpf/aya-ebpf/src/lib.rs index 55d7517b..0525b810 100644 --- a/ebpf/aya-ebpf/src/lib.rs +++ b/ebpf/aya-ebpf/src/lib.rs @@ -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; diff --git a/test/integration-common/src/lib.rs b/test/integration-common/src/lib.rs index 8efed81d..78cdd927 100644 --- a/test/integration-common/src/lib.rs +++ b/test/integration-common/src/lib.rs @@ -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; diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index 78c8b89e..d4d6da59 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -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" diff --git a/test/integration-ebpf/src/array.rs b/test/integration-ebpf/src/array.rs new file mode 100644 index 00000000..95a82bda --- /dev/null +++ b/test/integration-ebpf/src/array.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 = Array::new(); +#[btf_map] +static ARRAY: Array = Array::new(); + +#[map] +static RESULT_LEGACY: LegacyArray = LegacyArray::with_max_entries(3, 0); +#[map] +static ARRAY_LEGACY: LegacyArray = 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 +); diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index 9dbb4d18..f7da1d9f 100644 --- a/test/integration-test/src/lib.rs +++ b/test/integration-test/src/lib.rs @@ -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", diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index deaa63a5..f9427ab4 100644 --- a/test/integration-test/src/tests.rs +++ b/test/integration-test/src/tests.rs @@ -1,3 +1,4 @@ +mod array; mod bpf_probe_read; mod btf_relocations; mod elf; diff --git a/test/integration-test/src/tests/array.rs b/test/integration-test/src/tests/array.rs new file mode 100644 index 00000000..1b1c6e5c --- /dev/null +++ b/test/integration-test/src/tests/array.rs @@ -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); + } + } +}