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-rs · Feb 23, 2023 · f039aa8 · f039aa8
1 parent 8745f04
commit f039aa8
Show file tree

Hide file tree

Showing 8 changed files with 196 additions and 30 deletions.
diff --git a/aya/src/maps/xdp/cpu_map.rs b/aya/src/maps/xdp/cpu_map.rs
@@ -1,8 +1,13 @@
 //! An array of available CPUs.
 
+use std::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},
+    Pod,
 };
 
 /// An array of available CPUs.
@@ -36,7 +41,7 @@ pub struct CpuMap<T> {
 impl<T: AsRef<MapData>> CpuMap<T> {
     pub(crate) fn new(map: T) -> Result<CpuMap<T>, MapError> {
         let data = map.as_ref();
-        check_kv_size::<u32, u32>(data)?;
+        check_kv_size::<u32, bpf_cpumap_val>(data)?;
 
         let _fd = data.fd_or_err()?;
 
@@ -56,7 +61,7 @@ impl<T: AsRef<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.as_ref();
         check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
@@ -67,12 +72,13 @@ impl<T: AsRef<MapData>> CpuMap<T> {
                 io_error,
             }
         })?;
-        value.ok_or(MapError::KeyNotFound)
+        let value: bpf_cpumap_val = value.ok_or(MapError::KeyNotFound)?;
+        Ok(value.into())
     }
 
     /// An iterator over the elements of the map. The iterator item type is `Result<u32,
     /// MapError>`.
-    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))
     }
 }
@@ -84,10 +90,23 @@ impl<T: AsMut<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.as_mut();
         check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
+
+        let value = bpf_cpumap_val {
+            qsize: value,
+            bpf_prog: bpf_cpumap_val__bindgen_ty_1 {
+                fd: program.map(|prog| prog.as_raw_fd()).unwrap_or_default(),
+            },
+        };
         bpf_map_update_elem(fd, Some(&index), &value, flags).map_err(|(_, io_error)| {
             MapError::SyscallError {
                 call: "bpf_map_update_elem".to_owned(),
@@ -98,12 +117,31 @@ impl<T: AsMut<MapData>> CpuMap<T> {
     }
 }
 
-impl<T: AsRef<MapData>> IterableMap<u32, u32> for CpuMap<T> {
+impl<T: AsRef<MapData>> IterableMap<u32, CpuMapValue> for CpuMap<T> {
     fn map(&self) -> &MapData {
         self.inner.as_ref()
     }
 
-    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: u32,
+}
+
+impl From<bpf_cpumap_val> for CpuMapValue {
+    fn from(value: bpf_cpumap_val) -> Self {
+        // SAFETY: map writes use fd, map reads use id.
+        // https://elixir.bootlin.com/linux/v6.2/source/include/uapi/linux/bpf.h#L6149
+        CpuMapValue {
+            qsize: value.qsize,
+            prog_id: unsafe { value.bpf_prog.id },
+        }
+    }
+}
diff --git a/aya/src/maps/xdp/dev_map.rs b/aya/src/maps/xdp/dev_map.rs
@@ -1,8 +1,13 @@
 //! An array of network devices.
 
+use std::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},
+    Pod,
 };
 
 /// An array of network devices.
@@ -34,7 +39,7 @@ pub struct DevMap<T> {
 impl<T: AsRef<MapData>> DevMap<T> {
     pub(crate) fn new(map: T) -> Result<DevMap<T>, MapError> {
         let data = map.as_ref();
-        check_kv_size::<u32, u32>(data)?;
+        check_kv_size::<u32, bpf_devmap_val>(data)?;
 
         let _fd = data.fd_or_err()?;
 
@@ -54,7 +59,7 @@ impl<T: AsRef<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.as_ref();
         check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
@@ -65,12 +70,13 @@ impl<T: AsRef<MapData>> DevMap<T> {
                 io_error,
             }
         })?;
-        value.ok_or(MapError::KeyNotFound)
+        let value: bpf_devmap_val = value.ok_or(MapError::KeyNotFound)?;
+        Ok(value.into())
     }
 
     /// An iterator over the elements of the array. The iterator item type is `Result<u32,
     /// MapError>`.
-    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))
     }
 }
@@ -82,10 +88,23 @@ impl<T: AsMut<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.as_mut();
         check_bounds(data, index)?;
         let fd = data.fd_or_err()?;
+
+        let value = bpf_devmap_val {
+            ifindex: value,
+            bpf_prog: bpf_devmap_val__bindgen_ty_1 {
+                fd: program.map(|prog| prog.as_raw_fd()).unwrap_or_default(),
+            },
+        };
         bpf_map_update_elem(fd, Some(&index), &value, flags).map_err(|(_, io_error)| {
             MapError::SyscallError {
                 call: "bpf_map_update_elem".to_owned(),
@@ -96,12 +115,31 @@ impl<T: AsMut<MapData>> DevMap<T> {
     }
 }
 
-impl<T: AsRef<MapData>> IterableMap<u32, u32> for DevMap<T> {
+impl<T: AsRef<MapData>> IterableMap<u32, DevMapValue> for DevMap<T> {
     fn map(&self) -> &MapData {
         self.inner.as_ref()
     }
 
-    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: u32,
+}
+
+impl From<bpf_devmap_val> for DevMapValue {
+    fn from(value: bpf_devmap_val) -> Self {
+        // SAFETY: map writes use fd, map reads use id.
+        // https://elixir.bootlin.com/linux/v6.2/source/include/uapi/linux/bpf.h#L6136
+        DevMapValue {
+            ifindex: value.ifindex,
+            prog_id: unsafe { value.bpf_prog.id },
+        }
+    }
+}
diff --git a/aya/src/maps/xdp/dev_map_hash.rs b/aya/src/maps/xdp/dev_map_hash.rs
@@ -1,10 +1,16 @@
 //! An hashmap of network devices.
 
+use std::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,
 };
 
+use super::dev_map::DevMapValue;
+
 /// An hashmap of network devices.
 ///
 /// XDP programs can use this map to redirect to other network
@@ -34,7 +40,7 @@ pub struct DevMapHash<T> {
 impl<T: AsRef<MapData>> DevMapHash<T> {
     pub(crate) fn new(map: T) -> Result<DevMapHash<T>, MapError> {
         let data = map.as_ref();
-        check_kv_size::<u32, u32>(data)?;
+        check_kv_size::<u32, bpf_devmap_val>(data)?;
 
         let _fd = data.fd_or_err()?;
 
@@ -47,20 +53,21 @@ impl<T: AsRef<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.as_ref().fd_or_err()?;
         let value = bpf_map_lookup_elem(fd, &index, flags).map_err(|(_, io_error)| {
             MapError::SyscallError {
                 call: "bpf_map_lookup_elem".to_owned(),
                 io_error,
             }
         })?;
-        value.ok_or(MapError::KeyNotFound)
+        let value: bpf_devmap_val = value.ok_or(MapError::KeyNotFound)?;
+        Ok(value.into())
     }
 
     /// An iterator over the elements of the devmap in arbitrary order. The iterator item type is
     /// `Result<(u32, u32), MapError>`.
-    pub fn iter(&self) -> MapIter<'_, u32, u32, Self> {
+    pub fn iter(&self) -> MapIter<'_, u32, DevMapValue, Self> {
         MapIter::new(self)
     }
 
@@ -77,7 +84,19 @@ impl<T: AsMut<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 {
+                fd: program.map(|prog| prog.as_raw_fd()).unwrap_or_default(),
+            },
+        };
         hash_map::insert(self.inner.as_mut(), &index, &value, flags)
     }
 
@@ -91,12 +110,12 @@ impl<T: AsMut<MapData>> DevMapHash<T> {
     }
 }
 
-impl<T: AsRef<MapData>> IterableMap<u32, u32> for DevMapHash<T> {
+impl<T: AsRef<MapData>> IterableMap<u32, DevMapValue> for DevMapHash<T> {
     fn map(&self) -> &MapData {
         self.inner.as_ref()
     }
 
-    fn get(&self, key: &u32) -> Result<u32, MapError> {
+    fn get(&self, key: &u32) -> Result<DevMapValue, MapError> {
         self.get(*key, 0)
     }
 }
diff --git a/bpf/aya-bpf/src/maps/xdp/cpu_map.rs b/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,

diff --git a/bpf/aya-bpf/src/maps/xdp/dev_map.rs b/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,

diff --git a/bpf/aya-bpf/src/maps/xdp/dev_map_hash.rs b/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,