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;

aya-obj/src/obj.rs

@@ -184,6 +184,8 @@ pub enum ProgramSection {
     Xdp {
         name: String,
         frags_supported: bool,
+        cpumap: bool,
+        devmap: bool,
     },
     SkMsg {
         name: String,
@@ -316,10 +318,38 @@ impl FromStr for ProgramSection {
             "xdp" => Xdp {
                 name,
                 frags_supported: false,
+                cpumap: false,
+                devmap: false,
             },
             "xdp.frags" => Xdp {
                 name,
                 frags_supported: true,
+                cpumap: false,
+                devmap: false,
+            },
+            "xdp_cpumap" => Xdp {
+                name,
+                frags_supported: false,
+                cpumap: true,
+                devmap: false,
+            },
+            "xdp_devmap" => Xdp {
+                name,
+                frags_supported: false,
+                cpumap: false,
+                devmap: true,
+            },
+            "xdp.frags/cpumap" => Xdp {
+                name,
+                frags_supported: true,
+                cpumap: true,
+                devmap: false,
+            },
+            "xdp.frags/devmap" => Xdp {
+                name,
+                frags_supported: true,
+                cpumap: false,
+                devmap: true,
             },
             "tp_btf" => BtfTracePoint { name },
             _ if kind.starts_with("tracepoint") || kind.starts_with("tp") => {

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,35 @@
+#![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"));
+probe!(
+    frags_cm_named,
+    (name = "baz", frags = "true", map = "cpumap")
+);
+probe!(
+    frags_dm_named,
+    (name = "bat", 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,53 @@ 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"));
+    assert!(has_symbol(
+        &obj_file,
+        "xdp.frags/cpumap/baz",
+        "frags_cm_named"
+    ));
+    assert!(has_symbol(
+        &obj_file,
+        "xdp.frags/devmap/bat",
+        "frags_dm_named"
+    ));
+}
+
+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))
+}
+
+#[integration_test]
+fn map_load() {
+    let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/xdp_sec");
+    let bpf = Bpf::load(bytes).unwrap();
+
+    bpf.program("xdp").unwrap(); // plain
+    bpf.program("foo").unwrap(); // named
+    bpf.program("xdp.frags").unwrap(); // frags
+    bpf.program("bar").unwrap(); // named_frags
+    bpf.program("cpumap").unwrap();
+    bpf.program("devmap").unwrap();
+    bpf.program("xdp_cpumap").unwrap(); // frags_cm
+    bpf.program("xdp_devmap").unwrap(); // frags_dm
+    bpf.program("baz").unwrap(); // frags_cm_named
+    bpf.program("bat").unwrap(); // frags_dm_named
+}