xdp: add support for chained xdp programs in {cpu,dev}map

set/insert functions can now take an optional bpf program fd to run once
the packet has been redirected from the main probe

aya/src/maps/xdp/cpu_map.rs

@@ -1,10 +1,17 @@
 //! An array of available CPUs.
 
-use std::borrow::{Borrow, BorrowMut};
+use std::{
+    borrow::{Borrow, BorrowMut},
+    num::NonZeroU32,
+    os::fd::AsRawFd,
+};
+
+use aya_obj::generated::{bpf_cpumap_val, bpf_cpumap_val__bindgen_ty_1};
 
 use crate::{
     maps::{check_bounds, check_kv_size, IterableMap, MapData, MapError},
     sys::{bpf_map_lookup_elem, bpf_map_update_elem, SyscallError},
+    Pod,
 };
 
 /// An array of available CPUs.
@@ -29,7 +36,7 @@ use crate::{
 /// let flags = 0;
 /// let queue_size = 2048;
 /// for i in 0u32..8u32 {
-///     cpumap.set(i, queue_size, flags);
+///     cpumap.set(i, queue_size, None::<i32>, flags);
 /// }
 ///
 /// # Ok::<(), aya::BpfError>(())
@@ -42,7 +49,7 @@ pub struct CpuMap<T> {
 impl<T: Borrow<MapData>> CpuMap<T> {
     pub(crate) fn new(map: T) -> Result<Self, MapError> {
         let data = map.borrow();
-        check_kv_size::<u32, u32>(data)?;
+        check_kv_size::<u32, bpf_cpumap_val>(data)?;
 
         Ok(Self { inner: map })
     }
@@ -60,7 +67,7 @@ impl<T: Borrow<MapData>> CpuMap<T> {
     ///
     /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`]
     /// if `bpf_map_lookup_elem` fails.
-    pub fn get(&self, index: u32, flags: u64) -> Result<u32, MapError> {
+    pub fn get(&self, index: u32, flags: u64) -> Result<CpuMapValue, MapError> {
         let data = self.inner.borrow();
         check_bounds(data, index)?;
         let fd = data.fd().as_fd();
@@ -70,11 +77,18 @@ impl<T: Borrow<MapData>> CpuMap<T> {
                 call: "bpf_map_lookup_elem",
                 io_error,
             })?;
-        value.ok_or(MapError::KeyNotFound)
+        let value: bpf_cpumap_val = value.ok_or(MapError::KeyNotFound)?;
+
+        // SAFETY: map writes use fd, map reads use id.
+        // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/include/uapi/linux/bpf.h#L6241
+        Ok(CpuMapValue {
+            qsize: value.qsize,
+            prog_id: NonZeroU32::new(unsafe { value.bpf_prog.id }),
+        })
     }
 
     /// An iterator over the elements of the map.
-    pub fn iter(&self) -> impl Iterator<Item = Result<u32, MapError>> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = Result<CpuMapValue, MapError>> + '_ {
         (0..self.len()).map(move |i| self.get(i, 0))
     }
 }
@@ -86,10 +100,25 @@ impl<T: BorrowMut<MapData>> CpuMap<T> {
     ///
     /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`]
     /// if `bpf_map_update_elem` fails.
-    pub fn set(&mut self, index: u32, value: u32, flags: u64) -> Result<(), MapError> {
+    pub fn set(
+        &mut self,
+        index: u32,
+        value: u32,
+        program: Option<impl AsRawFd>,
+        flags: u64,
+    ) -> Result<(), MapError> {
         let data = self.inner.borrow_mut();
         check_bounds(data, index)?;
         let fd = data.fd().as_fd();
+
+        let value = bpf_cpumap_val {
+            qsize: value,
+            bpf_prog: bpf_cpumap_val__bindgen_ty_1 {
+                // Default is valid as the kernel will only consider fd > 0:
+                // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/kernel/bpf/cpumap.c#L466
+                fd: program.map(|prog| prog.as_raw_fd()).unwrap_or_default(),
+            },
+        };
         bpf_map_update_elem(fd, Some(&index), &value, flags).map_err(|(_, io_error)| {
             SyscallError {
                 call: "bpf_map_update_elem",
@@ -100,12 +129,20 @@ impl<T: BorrowMut<MapData>> CpuMap<T> {
     }
 }
 
-impl<T: Borrow<MapData>> IterableMap<u32, u32> for CpuMap<T> {
+impl<T: Borrow<MapData>> IterableMap<u32, CpuMapValue> for CpuMap<T> {
     fn map(&self) -> &MapData {
         self.inner.borrow()
     }
 
-    fn get(&self, key: &u32) -> Result<u32, MapError> {
+    fn get(&self, key: &u32) -> Result<CpuMapValue, MapError> {
         self.get(*key, 0)
     }
 }
+
+unsafe impl Pod for bpf_cpumap_val {}
+
+#[derive(Clone, Copy, Debug)]
+pub struct CpuMapValue {
+    pub qsize: u32,
+    pub prog_id: Option<NonZeroU32>,
+}

aya/src/maps/xdp/dev_map.rs

@@ -1,10 +1,17 @@
 //! An array of network devices.
 
-use std::borrow::{Borrow, BorrowMut};
+use std::{
+    borrow::{Borrow, BorrowMut},
+    num::NonZeroU32,
+    os::fd::AsRawFd,
+};
+
+use aya_obj::generated::{bpf_devmap_val, bpf_devmap_val__bindgen_ty_1};
 
 use crate::{
     maps::{check_bounds, check_kv_size, IterableMap, MapData, MapError},
     sys::{bpf_map_lookup_elem, bpf_map_update_elem, SyscallError},
+    Pod,
 };
 
 /// An array of network devices.
@@ -24,7 +31,7 @@ use crate::{
 /// let mut devmap = DevMap::try_from(bpf.map_mut("IFACES").unwrap())?;
 /// let source = 32u32;
 /// let dest = 42u32;
-/// devmap.set(source, dest, 0);
+/// devmap.set(source, dest, None::<i32>, 0);
 ///
 /// # Ok::<(), aya::BpfError>(())
 /// ```
@@ -36,7 +43,7 @@ pub struct DevMap<T> {
 impl<T: Borrow<MapData>> DevMap<T> {
     pub(crate) fn new(map: T) -> Result<Self, MapError> {
         let data = map.borrow();
-        check_kv_size::<u32, u32>(data)?;
+        check_kv_size::<u32, bpf_devmap_val>(data)?;
 
         Ok(Self { inner: map })
     }
@@ -54,7 +61,7 @@ impl<T: Borrow<MapData>> DevMap<T> {
     ///
     /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`]
     /// if `bpf_map_lookup_elem` fails.
-    pub fn get(&self, index: u32, flags: u64) -> Result<u32, MapError> {
+    pub fn get(&self, index: u32, flags: u64) -> Result<DevMapValue, MapError> {
         let data = self.inner.borrow();
         check_bounds(data, index)?;
         let fd = data.fd().as_fd();
@@ -64,11 +71,18 @@ impl<T: Borrow<MapData>> DevMap<T> {
                 call: "bpf_map_lookup_elem",
                 io_error,
             })?;
-        value.ok_or(MapError::KeyNotFound)
+        let value: bpf_devmap_val = value.ok_or(MapError::KeyNotFound)?;
+
+        // SAFETY: map writes use fd, map reads use id.
+        // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/include/uapi/linux/bpf.h#L6228
+        Ok(DevMapValue {
+            ifindex: value.ifindex,
+            prog_id: NonZeroU32::new(unsafe { value.bpf_prog.id }),
+        })
     }
 
     /// An iterator over the elements of the array.
-    pub fn iter(&self) -> impl Iterator<Item = Result<u32, MapError>> + '_ {
+    pub fn iter(&self) -> impl Iterator<Item = Result<DevMapValue, MapError>> + '_ {
         (0..self.len()).map(move |i| self.get(i, 0))
     }
 }
@@ -80,10 +94,26 @@ impl<T: BorrowMut<MapData>> DevMap<T> {
     ///
     /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`]
     /// if `bpf_map_update_elem` fails.
-    pub fn set(&mut self, index: u32, value: u32, flags: u64) -> Result<(), MapError> {
+    pub fn set(
+        &mut self,
+        index: u32,
+        value: u32,
+        program: Option<impl AsRawFd>,
+        flags: u64,
+    ) -> Result<(), MapError> {
         let data = self.inner.borrow_mut();
         check_bounds(data, index)?;
         let fd = data.fd().as_fd();
+
+        let value = bpf_devmap_val {
+            ifindex: value,
+            bpf_prog: bpf_devmap_val__bindgen_ty_1 {
+                // Default is valid as the kernel will only consider fd > 0:
+                // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/kernel/bpf/devmap.c#L866
+                // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/kernel/bpf/devmap.c#L918
+                fd: program.map(|prog| prog.as_raw_fd()).unwrap_or_default(),
+            },
+        };
         bpf_map_update_elem(fd, Some(&index), &value, flags).map_err(|(_, io_error)| {
             SyscallError {
                 call: "bpf_map_update_elem",
@@ -94,12 +124,20 @@ impl<T: BorrowMut<MapData>> DevMap<T> {
     }
 }
 
-impl<T: Borrow<MapData>> IterableMap<u32, u32> for DevMap<T> {
+impl<T: Borrow<MapData>> IterableMap<u32, DevMapValue> for DevMap<T> {
     fn map(&self) -> &MapData {
         self.inner.borrow()
     }
 
-    fn get(&self, key: &u32) -> Result<u32, MapError> {
+    fn get(&self, key: &u32) -> Result<DevMapValue, MapError> {
         self.get(*key, 0)
     }
 }
+
+unsafe impl Pod for bpf_devmap_val {}
+
+#[derive(Clone, Copy, Debug)]
+pub struct DevMapValue {
+    pub ifindex: u32,
+    pub prog_id: Option<NonZeroU32>,
+}

aya/src/maps/xdp/dev_map_hash.rs

@@ -1,12 +1,20 @@
 //! An hashmap of network devices.
 
-use std::borrow::{Borrow, BorrowMut};
+use std::{
+    borrow::{Borrow, BorrowMut},
+    num::NonZeroU32,
+    os::fd::AsRawFd,
+};
+
+use aya_obj::generated::{bpf_devmap_val, bpf_devmap_val__bindgen_ty_1};
 
 use crate::{
     maps::{check_kv_size, hash_map, IterableMap, MapData, MapError, MapIter, MapKeys},
     sys::{bpf_map_lookup_elem, SyscallError},
 };
 
+use super::dev_map::DevMapValue;
+
 /// An hashmap of network devices.
 ///
 /// XDP programs can use this map to redirect to other network
@@ -24,7 +32,7 @@ use crate::{
 /// let mut devmap = DevMapHash::try_from(bpf.map_mut("IFACES").unwrap())?;
 /// let flags = 0;
 /// let ifindex = 32u32;
-/// devmap.insert(ifindex, ifindex, flags);
+/// devmap.insert(ifindex, ifindex, None::<i32>, flags);
 ///
 /// # Ok::<(), aya::BpfError>(())
 /// ```
@@ -36,7 +44,7 @@ pub struct DevMapHash<T> {
 impl<T: Borrow<MapData>> DevMapHash<T> {
     pub(crate) fn new(map: T) -> Result<Self, MapError> {
         let data = map.borrow();
-        check_kv_size::<u32, u32>(data)?;
+        check_kv_size::<u32, bpf_devmap_val>(data)?;
 
         Ok(Self { inner: map })
     }
@@ -47,18 +55,25 @@ impl<T: Borrow<MapData>> DevMapHash<T> {
     ///
     /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`]
     /// if `bpf_map_lookup_elem` fails.
-    pub fn get(&self, index: u32, flags: u64) -> Result<u32, MapError> {
+    pub fn get(&self, index: u32, flags: u64) -> Result<DevMapValue, MapError> {
         let fd = self.inner.borrow().fd().as_fd();
         let value =
             bpf_map_lookup_elem(fd, &index, flags).map_err(|(_, io_error)| SyscallError {
                 call: "bpf_map_lookup_elem",
                 io_error,
             })?;
-        value.ok_or(MapError::KeyNotFound)
+        let value: bpf_devmap_val = value.ok_or(MapError::KeyNotFound)?;
+
+        // SAFETY: map writes use fd, map reads use id.
+        // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/include/uapi/linux/bpf.h#L6228
+        Ok(DevMapValue {
+            ifindex: value.ifindex,
+            prog_id: NonZeroU32::new(unsafe { value.bpf_prog.id }),
+        })
     }
 
     /// An iterator over the elements of the devmap in arbitrary order.
-    pub fn iter(&self) -> MapIter<'_, u32, u32, Self> {
+    pub fn iter(&self) -> MapIter<'_, u32, DevMapValue, Self> {
         MapIter::new(self)
     }
 
@@ -74,7 +89,22 @@ impl<T: BorrowMut<MapData>> DevMapHash<T> {
     /// # Errors
     ///
     /// Returns [`MapError::SyscallError`] if `bpf_map_update_elem` fails.
-    pub fn insert(&mut self, index: u32, value: u32, flags: u64) -> Result<(), MapError> {
+    pub fn insert(
+        &mut self,
+        index: u32,
+        value: u32,
+        program: Option<impl AsRawFd>,
+        flags: u64,
+    ) -> Result<(), MapError> {
+        let value = bpf_devmap_val {
+            ifindex: value,
+            bpf_prog: bpf_devmap_val__bindgen_ty_1 {
+                // Default is valid as the kernel will only consider fd > 0:
+                // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/kernel/bpf/devmap.c#L866
+                // https://github.com/torvalds/linux/blob/2dde18cd1d8fac735875f2e4987f11817cc0bc2c/kernel/bpf/devmap.c#L918
+                fd: program.map(|prog| prog.as_raw_fd()).unwrap_or_default(),
+            },
+        };
         hash_map::insert(self.inner.borrow_mut(), &index, &value, flags)
     }
 
@@ -88,12 +118,12 @@ impl<T: BorrowMut<MapData>> DevMapHash<T> {
     }
 }
 
-impl<T: Borrow<MapData>> IterableMap<u32, u32> for DevMapHash<T> {
+impl<T: Borrow<MapData>> IterableMap<u32, DevMapValue> for DevMapHash<T> {
     fn map(&self) -> &MapData {
         self.inner.borrow()
     }
 
-    fn get(&self, key: &u32) -> Result<u32, MapError> {
+    fn get(&self, key: &u32) -> Result<DevMapValue, MapError> {
         self.get(*key, 0)
     }
 }

bpf/aya-bpf/src/maps/xdp/cpu_map.rs

@@ -1,5 +1,7 @@
 use core::{cell::UnsafeCell, mem};
 
+use aya_bpf_bindings::bindings::bpf_cpumap_val;
+
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_CPUMAP},
     helpers::bpf_redirect_map,
@@ -19,7 +21,7 @@ impl CpuMap {
             def: UnsafeCell::new(bpf_map_def {
                 type_: BPF_MAP_TYPE_CPUMAP,
                 key_size: mem::size_of::<u32>() as u32,
-                value_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<bpf_cpumap_val>() as u32,
                 max_entries,
                 map_flags: flags,
                 id: 0,
@@ -33,7 +35,7 @@ impl CpuMap {
             def: UnsafeCell::new(bpf_map_def {
                 type_: BPF_MAP_TYPE_CPUMAP,
                 key_size: mem::size_of::<u32>() as u32,
-                value_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<bpf_cpumap_val>() as u32,
                 max_entries,
                 map_flags: flags,
                 id: 0,

bpf/aya-bpf/src/maps/xdp/dev_map.rs

@@ -1,5 +1,6 @@
 use core::{cell::UnsafeCell, mem, ptr::NonNull};
 
+use aya_bpf_bindings::bindings::bpf_devmap_val;
 use aya_bpf_cty::c_void;
 
 use crate::{
@@ -21,7 +22,7 @@ impl DevMap {
             def: UnsafeCell::new(bpf_map_def {
                 type_: BPF_MAP_TYPE_DEVMAP,
                 key_size: mem::size_of::<u32>() as u32,
-                value_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<bpf_devmap_val>() as u32,
                 max_entries,
                 map_flags: flags,
                 id: 0,
@@ -35,7 +36,7 @@ impl DevMap {
             def: UnsafeCell::new(bpf_map_def {
                 type_: BPF_MAP_TYPE_DEVMAP,
                 key_size: mem::size_of::<u32>() as u32,
-                value_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<bpf_devmap_val>() as u32,
                 max_entries,
                 map_flags: flags,
                 id: 0,

bpf/aya-bpf/src/maps/xdp/dev_map_hash.rs

@@ -22,7 +22,7 @@ impl DevMapHash {
             def: UnsafeCell::new(bpf_map_def {
                 type_: BPF_MAP_TYPE_DEVMAP_HASH,
                 key_size: mem::size_of::<u32>() as u32,
-                value_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<bpf_devmap_val>() as u32,
                 max_entries,
                 map_flags: flags,
                 id: 0,
@@ -36,7 +36,7 @@ impl DevMapHash {
             def: UnsafeCell::new(bpf_map_def {
                 type_: BPF_MAP_TYPE_DEVMAP_HASH,
                 key_size: mem::size_of::<u32>() as u32,
-                value_size: mem::size_of::<u32>() as u32,
+                value_size: mem::size_of::<bpf_devmap_val>() as u32,
                 max_entries,
                 map_flags: flags,
                 id: 0,

test/integration-ebpf/src/redirect.rs

@@ -4,7 +4,7 @@
 use aya_bpf::{
     bindings::xdp_action,
     macros::{map, xdp},
-    maps::{CpuMap, DevMap, DevMapHash, XskMap},
+    maps::{Array, CpuMap, DevMap, DevMapHash, XskMap},
     programs::XdpContext,
 };
 
@@ -17,6 +17,13 @@ static DEVS_HASH: DevMapHash = DevMapHash::with_max_entries(1, 0);
 #[map]
 static CPUS: CpuMap = CpuMap::with_max_entries(1, 0);
 
+/// Hits of a probe, used to test program chaining through CpuMap/DevMap.
+/// The first slot counts how many times the "raw" xdp program got executed, while the second slot
+/// counts how many times the map programs got executed.
+/// This allows the test harness to assert that a specific step got executed.
+#[map]
+static mut HITS: Array<u32> = Array::with_max_entries(2, 0);
+
 #[xdp]
 pub fn redirect_sock(_ctx: XdpContext) -> u32 {
     SOCKS.redirect(0, xdp_action::XDP_ABORTED as u64)
@@ -24,19 +31,41 @@ pub fn redirect_sock(_ctx: XdpContext) -> u32 {
 
 #[xdp]
 pub fn redirect_dev(_ctx: XdpContext) -> u32 {
+    inc_hit(0);
     DEVS.redirect(0, xdp_action::XDP_ABORTED as u64)
 }
 
 #[xdp]
 pub fn redirect_dev_hash(_ctx: XdpContext) -> u32 {
+    inc_hit(0);
     DEVS_HASH.redirect(10, xdp_action::XDP_ABORTED as u64)
 }
 
 #[xdp]
 pub fn redirect_cpu(_ctx: XdpContext) -> u32 {
+    inc_hit(0);
     CPUS.redirect(0, xdp_action::XDP_ABORTED as u64)
 }
 
+#[xdp(map = "cpumap")]
+pub fn redirect_cpu_chain(_ctx: XdpContext) -> u32 {
+    inc_hit(1);
+    xdp_action::XDP_PASS
+}
+
+#[xdp(map = "devmap")]
+pub fn redirect_dev_chain(_ctx: XdpContext) -> u32 {
+    inc_hit(1);
+    xdp_action::XDP_PASS
+}
+
+#[inline(always)]
+fn inc_hit(index: u32) {
+    if let Some(hit) = unsafe { HITS.get_ptr_mut(index) } {
+        unsafe { *hit += 1 };
+    }
+}
+
 #[cfg(not(test))]
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {

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

@@ -1,6 +1,14 @@
-use aya::Bpf;
+use std::{net::UdpSocket, os::fd::AsFd, time::Duration};
+
+use aya::{
+    maps::{Array, CpuMap},
+    programs::{Xdp, XdpFlags},
+    Bpf,
+};
 use object::{Object, ObjectSection, ObjectSymbol, SymbolSection};
 
+use crate::utils::NetNsGuard;
+
 #[test]
 fn prog_sections() {
     let obj_file = object::File::parse(crate::XDP_SEC).unwrap();
@@ -46,3 +54,45 @@ fn map_load() {
     bpf.program("xdp_frags_cpumap").unwrap();
     bpf.program("xdp_frags_devmap").unwrap();
 }
+
+#[test]
+fn cpumap_chain() {
+    let _netns = NetNsGuard::new();
+
+    let mut bpf = Bpf::load(crate::REDIRECT).unwrap();
+
+    // Load our cpumap and our canary map
+    let mut cpus: CpuMap<_> = bpf.take_map("CPUS").unwrap().try_into().unwrap();
+    let hits: Array<_, u32> = bpf.take_map("HITS").unwrap().try_into().unwrap();
+
+    let xdp_chain_fd = {
+        // Load the chained program to run on the target CPU
+        let xdp: &mut Xdp = bpf
+            .program_mut("redirect_cpu_chain")
+            .unwrap()
+            .try_into()
+            .unwrap();
+        xdp.load().unwrap();
+        xdp.fd().unwrap()
+    };
+    cpus.set(0, 2048, Some(xdp_chain_fd.as_fd()), 0).unwrap();
+
+    // Load the main program
+    let xdp: &mut Xdp = bpf.program_mut("redirect_cpu").unwrap().try_into().unwrap();
+    xdp.load().unwrap();
+    xdp.attach("lo", XdpFlags::default()).unwrap();
+
+    let sock = UdpSocket::bind("127.0.0.1:1777").unwrap();
+    sock.set_read_timeout(Some(Duration::from_millis(1)))
+        .unwrap();
+    sock.send_to(b"hello cpumap", "127.0.0.1:1777").unwrap();
+
+    // Read back the packet to ensure it wenth through the entire network stack, including our two
+    // probes.
+    let mut buf = vec![0u8; 1000];
+    let n = sock.recv(&mut buf).unwrap();
+
+    assert_eq!(&buf[..n], b"hello cpumap");
+    assert_eq!(hits.get(&0, 0).unwrap(), 1);
+    assert_eq!(hits.get(&1, 0).unwrap(), 1);
+}