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/expand.rs

@@ -201,6 +201,13 @@ pub struct Xdp {
     item: ItemFn,
     name: Option<String>,
     frags: bool,
+    map: Option<XdpMap>,
+}
+
+#[derive(Clone, Copy)]
+pub enum XdpMap {
+    CpuMap,
+    DevMap,
 }
 
 impl Xdp {
@@ -217,16 +224,42 @@ impl Xdp {
                 ));
             }
         }
+        let map = match pop_arg(&mut args, "map").as_deref() {
+            Some("cpumap") => Some(XdpMap::CpuMap),
+            Some("devmap") => Some(XdpMap::DevMap),
+            Some(_) => {
+                return Err(Error::new_spanned(
+                    "map",
+                    "invalid value. should be 'cpumap' or 'devmap'",
+                ))
+            }
+            None => None,
+        };
+
         err_on_unknown_args(&args)?;
-        Ok(Xdp { item, name, frags })
+        Ok(Xdp {
+            item,
+            name,
+            frags,
+            map,
+        })
     }
 
     pub fn expand(&self) -> Result<TokenStream> {
-        let section_prefix = if self.frags { "xdp.frags" } else { "xdp" };
+        let (prefix, delimiter) = if self.frags {
+            ("xdp.frags", "/")
+        } else {
+            ("xdp", "_")
+        };
+        let section_prefix = match self.map {
+            None => prefix.to_owned(),
+            Some(XdpMap::CpuMap) => format!("{prefix}{delimiter}cpumap"),
+            Some(XdpMap::DevMap) => format!("{prefix}{delimiter}devmap"),
+        };
         let section_name = if let Some(name) = &self.name {
             format!("{section_prefix}/{name}")
         } else {
-            section_prefix.to_string()
+            section_prefix
         };
         let fn_vis = &self.item.vis;
         let fn_name = &self.item.sig.ident;

test/integration-ebpf/Cargo.toml

@@ -26,3 +26,7 @@ path = "src/test.rs"
 [[bin]]
 name = "redirect"
 path = "src/redirect.rs"
+
+[[bin]]
+name = "xdp_sec"
+path = "src/xdp_sec.rs"

test/integration-ebpf/src/xdp_sec.rs

@@ -0,0 +1,27 @@
+#![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!(plain, ());
+probe!(named, (name = "foo"));
+probe!(frags, (frags = "true"));
+probe!(named_frags, (name = "bar", frags = "true"));
+probe!(cpumap, (map = "cpumap"));
+probe!(devmap, (map = "devmap"));
+probe!(frags_cm, (frags = "true", map = "cpumap"));
+probe!(frags_dm, (frags = "true", map = "devmap"));
+
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    unsafe { core::hint::unreachable_unchecked() }
+}

test/integration-test/src/tests/xdp.rs

@@ -6,6 +6,7 @@ use aya::{
     programs::{Xdp, XdpFlags},
     Bpf,
 };
+use object::{Object, ObjectSection, ObjectSymbol, SymbolSection};
 use xsk_rs::{
     config::{LibbpfFlags, SocketConfigBuilder},
     Socket, Umem,
@@ -58,3 +59,26 @@ fn af_xdp() {
     assert_eq!(&udp[2..4], 1777u16.to_be_bytes().as_slice()); // Dest
     assert_eq!(payload, b"hello AF_XDP");
 }
+
+#[integration_test]
+fn prog_sections() {
+    let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/xdp_sec");
+    let obj_file = object::File::parse(bytes).unwrap();
+
+    assert!(has_symbol(&obj_file, "xdp", "plain"));
+    assert!(has_symbol(&obj_file, "xdp/foo", "named"));
+    assert!(has_symbol(&obj_file, "xdp.frags", "frags"));
+    assert!(has_symbol(&obj_file, "xdp.frags/bar", "named_frags"));
+    assert!(has_symbol(&obj_file, "xdp_cpumap", "cpumap"));
+    assert!(has_symbol(&obj_file, "xdp_devmap", "devmap"));
+    assert!(has_symbol(&obj_file, "xdp.frags/cpumap", "frags_cm"));
+    assert!(has_symbol(&obj_file, "xdp.frags/devmap", "frags_dm"));
+}
+
+fn has_symbol(obj_file: &object::File, sec_name: &str, sym_name: &str) -> bool {
+    let sec = obj_file.section_by_name(sec_name).unwrap();
+    let sec = SymbolSection::Section(sec.index());
+    obj_file
+        .symbols()
+        .any(|sym| sym.section() == sec && sym.name() == Ok(sym_name))
+}