macros: add 'map' option to xdp macro

This option allows to place the program in the specific sections to
chain programs with devmaps and cpumaps.

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<XdpMap>,
+}
+
+#[derive(Clone, Copy)]
+pub(crate) enum XdpMap {
+    CpuMap,
+    DevMap,
 }
 
 impl Xdp {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Xdp> {
         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<TokenStream> {
-        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());
+    }
 }

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"

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 {}
+}

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 {}
+}

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;

test/integration-test/src/tests.rs

@@ -6,3 +6,4 @@ mod log;
 mod rbpf;
 mod relocations;
 mod smoke;
+mod xdp;

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::<Vec<_>>();
+        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::<Vec<_>>();
+    assert!(
+        syms.contains(&sym_name),
+        "symbol not found. available symbols in section: {syms:?}"
+    );
+}