diff --git a/aya-bpf-macros/src/xdp.rs b/aya-bpf-macros/src/xdp.rs index 9a5ff36a..6a8f61e7 100644 --- a/aya-bpf-macros/src/xdp.rs +++ b/aya-bpf-macros/src/xdp.rs @@ -1,31 +1,52 @@ -use std::borrow::Cow; - use proc_macro2::TokenStream; use quote::quote; -use syn::{ItemFn, Result}; +use syn::{Error, ItemFn, Result}; -use crate::args::{err_on_unknown_args, pop_bool_arg, Args}; +use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg, Args}; pub(crate) struct Xdp { item: ItemFn, frags: bool, + map: Option, +} + +#[derive(Clone, Copy)] +pub(crate) enum XdpMap { + CpuMap, + DevMap, } impl Xdp { pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { let item = syn::parse2(item)?; let mut args: Args = syn::parse2(attrs)?; + let frags = pop_bool_arg(&mut args, "frags"); + let map = match pop_string_arg(&mut args, "map").as_deref() { + Some("cpumap") => Some(XdpMap::CpuMap), + Some("devmap") => Some(XdpMap::DevMap), + Some(name) => { + return Err(Error::new_spanned( + "map", + format!("invalid value. expected 'cpumap' or 'devmap', found '{name}'"), + )) + } + None => None, + }; + err_on_unknown_args(&args)?; - Ok(Xdp { item, frags }) + Ok(Xdp { item, frags, map }) } pub(crate) fn expand(&self) -> Result { - let section_name: Cow<'_, _> = if self.frags { - "xdp.frags".into() - } else { - "xdp".into() + let mut section_name = vec![if self.frags { "xdp.frags" } else { "xdp" }]; + match self.map { + Some(XdpMap::CpuMap) => section_name.push("cpumap"), + Some(XdpMap::DevMap) => section_name.push("devmap"), + None => (), }; + let section_name = section_name.join("/"); + let fn_vis = &self.item.vis; let fn_name = self.item.sig.ident.clone(); let item = &self.item; @@ -97,4 +118,122 @@ mod tests { }; assert_eq!(expected.to_string(), expanded.to_string()); } + + #[test] + fn test_xdp_cpumap() { + let prog = Xdp::parse( + parse_quote! { map = "cpumap" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "xdp/cpumap"] + fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return prog(::aya_bpf::programs::XdpContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn test_xdp_devmap() { + let prog = Xdp::parse( + parse_quote! { map = "devmap" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "xdp/devmap"] + fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return prog(::aya_bpf::programs::XdpContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + #[should_panic(expected = "Invalid value. Expected 'cpumap' or 'devmap', found 'badmap'")] + fn test_xdp_bad_map() { + Xdp::parse( + parse_quote! { map = "badmap" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + } + + #[test] + fn test_xdp_frags_cpumap() { + let prog = Xdp::parse( + parse_quote! { frags, map = "cpumap" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "xdp.frags/cpumap"] + fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return prog(::aya_bpf::programs::XdpContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } + + #[test] + fn test_xdp_frags_devmap() { + let prog = Xdp::parse( + parse_quote! { frags, map = "devmap" }, + parse_quote! { + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + }, + ) + .unwrap(); + let expanded = prog.expand().unwrap(); + let expected = quote! { + #[no_mangle] + #[link_section = "xdp.frags/devmap"] + fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 { + return prog(::aya_bpf::programs::XdpContext::new(ctx)); + + fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 { + 0 + } + } + }; + assert_eq!(expected.to_string(), expanded.to_string()); + } } diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index 12b04546..e72e0a87 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -43,3 +43,11 @@ path = "src/bpf_probe_read.rs" [[bin]] name = "two_progs" path = "src/two_progs.rs" + +[[bin]] +name = "redirect" +path = "src/redirect.rs" + +[[bin]] +name = "xdp_sec" +path = "src/xdp_sec.rs" diff --git a/test/integration-ebpf/src/redirect.rs b/test/integration-ebpf/src/redirect.rs new file mode 100644 index 00000000..827ea451 --- /dev/null +++ b/test/integration-ebpf/src/redirect.rs @@ -0,0 +1,44 @@ +#![no_std] +#![no_main] + +use aya_bpf::{ + bindings::xdp_action, + macros::{map, xdp}, + maps::{CpuMap, DevMap, DevMapHash, XskMap}, + programs::XdpContext, +}; + +#[map] +static SOCKS: XskMap = XskMap::with_max_entries(1, 0); +#[map] +static DEVS: DevMap = DevMap::with_max_entries(1, 0); +#[map] +static DEVS_HASH: DevMapHash = DevMapHash::with_max_entries(1, 0); +#[map] +static CPUS: CpuMap = CpuMap::with_max_entries(1, 0); + +#[xdp] +pub fn redirect_sock(_ctx: XdpContext) -> u32 { + SOCKS.redirect(0, xdp_action::XDP_ABORTED as u64) +} + +#[xdp] +pub fn redirect_dev(_ctx: XdpContext) -> u32 { + DEVS.redirect(0, xdp_action::XDP_ABORTED as u64) +} + +#[xdp] +pub fn redirect_dev_hash(_ctx: XdpContext) -> u32 { + DEVS_HASH.redirect(10, xdp_action::XDP_ABORTED as u64) +} + +#[xdp] +pub fn redirect_cpu(_ctx: XdpContext) -> u32 { + CPUS.redirect(0, xdp_action::XDP_ABORTED as u64) +} + +#[cfg(not(test))] +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + loop {} +} diff --git a/test/integration-ebpf/src/xdp_sec.rs b/test/integration-ebpf/src/xdp_sec.rs new file mode 100644 index 00000000..5e64f56a --- /dev/null +++ b/test/integration-ebpf/src/xdp_sec.rs @@ -0,0 +1,26 @@ +#![no_std] +#![no_main] + +use aya_bpf::{bindings::xdp_action::XDP_PASS, macros::xdp, programs::XdpContext}; + +macro_rules! probe { + ($name:ident, ($($arg:ident $(= $value:literal)?),*) ) => { + #[xdp($($arg $(= $value)?),*)] + pub fn $name(_ctx: XdpContext) -> u32 { + XDP_PASS + } + }; +} + +probe!(xdp_plain, ()); +probe!(xdp_frags, (frags)); +probe!(xdp_cpumap, (map = "cpumap")); +probe!(xdp_devmap, (map = "devmap")); +probe!(xdp_frags_cpumap, (frags, map = "cpumap")); +probe!(xdp_frags_devmap, (frags, map = "devmap")); + +#[cfg(not(test))] +#[panic_handler] +fn panic(_info: &core::panic::PanicInfo) -> ! { + loop {} +} diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index 08911a4e..79be5fd3 100644 --- a/test/integration-test/src/lib.rs +++ b/test/integration-test/src/lib.rs @@ -19,6 +19,8 @@ pub const RELOCATIONS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), " pub const TWO_PROGS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/two_progs")); pub const BPF_PROBE_READ: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/bpf_probe_read")); +pub const REDIRECT: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/redirect")); +pub const XDP_SEC: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/xdp_sec")); #[cfg(test)] mod tests; diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index dd8565b0..3a995180 100644 --- a/test/integration-test/src/tests.rs +++ b/test/integration-test/src/tests.rs @@ -6,3 +6,4 @@ mod log; mod rbpf; mod relocations; mod smoke; +mod xdp; diff --git a/test/integration-test/src/tests/xdp.rs b/test/integration-test/src/tests/xdp.rs new file mode 100644 index 00000000..0c2f56da --- /dev/null +++ b/test/integration-test/src/tests/xdp.rs @@ -0,0 +1,35 @@ +use object::{Object, ObjectSection, ObjectSymbol, SymbolSection}; + +#[test] +fn prog_sections() { + let obj_file = object::File::parse(crate::XDP_SEC).unwrap(); + + ensure_symbol(&obj_file, "xdp", "xdp_plain"); + ensure_symbol(&obj_file, "xdp.frags", "xdp_frags"); + ensure_symbol(&obj_file, "xdp/cpumap", "xdp_cpumap"); + ensure_symbol(&obj_file, "xdp/devmap", "xdp_devmap"); + ensure_symbol(&obj_file, "xdp.frags/cpumap", "xdp_frags_cpumap"); + ensure_symbol(&obj_file, "xdp.frags/devmap", "xdp_frags_devmap"); +} + +#[track_caller] +fn ensure_symbol(obj_file: &object::File, sec_name: &str, sym_name: &str) { + let sec = obj_file.section_by_name(sec_name).unwrap_or_else(|| { + let secs = obj_file + .sections() + .flat_map(|sec| sec.name().ok().map(|name| name.to_owned())) + .collect::>(); + panic!("section {sec_name} not found. available sections: {secs:?}"); + }); + let sec = SymbolSection::Section(sec.index()); + + let syms = obj_file + .symbols() + .filter(|sym| sym.section() == sec) + .filter_map(|sym| sym.name().ok()) + .collect::>(); + assert!( + syms.contains(&sym_name), + "symbol not found. available symbols in section: {syms:?}" + ); +}