Add: maps_many_mut

aya/src/bpf.rs

@@ -990,6 +990,35 @@ impl Ebpf {
         self.maps.iter_mut().map(|(name, map)| (name.as_str(), map))
     }
 
+    /// Attempts to get mutable references to `N` maps at once.
+    ///
+    /// Returns an array of length `N` with the results of each query, in the same order
+    /// as the requested map names. For soundness, at most one mutable reference will be
+    /// returned to any map. `None` will be used if a map with the given name is missing.
+    ///
+    /// This method performs a check to ensure that there are no duplicate map names,
+    /// which currently has a time-complexity of *O(n²)*. Be careful when passing a large
+    /// number of names.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any names are duplicated.
+    ///
+    /// # Examples
+    /// ```no_run
+    /// # let mut bpf = aya::Ebpf::load(&[])?;
+    /// match bpf.maps_many_mut(["MAP1", "MAP2"]) {
+    ///     [Some(m1), Some(m2)] => println!("Got MAP1 and MAP2"),
+    ///     [Some(m1), None] => println!("Got only MAP1"),
+    ///     [None, Some(m2)] => println!("Got only MAP2"),
+    ///     [None, None] => println!("No maps"),
+    /// }
+    /// # Ok::<(), aya::EbpfError>(())
+    /// ```
+    pub fn maps_many_mut<const N: usize>(&mut self, names: [&str; N]) -> [Option<&mut Map>; N] {
+        self.maps.get_disjoint_mut(names)
+    }
+
     /// Returns a reference to the program with the given name.
     ///
     /// You can use this to inspect a program and its properties. To load and attach a program, use

test/integration-ebpf/src/map_test.rs

@@ -3,9 +3,9 @@
 #![expect(unused_crate_dependencies, reason = "used in other bins")]
 
 use aya_ebpf::{
-    macros::{map, socket_filter},
+    macros::{map, socket_filter, uprobe},
     maps::{Array, HashMap},
-    programs::SkBuffContext,
+    programs::{ProbeContext, SkBuffContext},
 };
 #[cfg(not(test))]
 extern crate ebpf_panic;
@@ -35,3 +35,20 @@ fn simple_prog(_ctx: SkBuffContext) -> i64 {
 
     0
 }
+
+#[uprobe]
+fn simple_prog_mut(_ctx: ProbeContext) -> i64 {
+    if let Some(foo) = FOO.get_ptr_mut(0) {
+        unsafe {
+            *foo += 1;
+        }
+    }
+
+    if let Some(baz) = BAZ.get_ptr_mut(0) {
+        unsafe {
+            *baz += 1;
+        }
+    }
+
+    0
+}

test/integration-test/src/tests.rs

@@ -9,6 +9,7 @@ mod linear_data_structures;
 mod load;
 mod log;
 mod lsm;
+mod map_many;
 mod map_pin;
 mod raw_tracepoint;
 mod rbpf;

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

@@ -0,0 +1,51 @@
+use aya::{
+    Ebpf,
+    maps::{Array, HashMap},
+    programs::UProbe,
+};
+
+#[unsafe(no_mangle)]
+#[inline(never)]
+extern "C" fn trigger_ebpf_program_maps_many() {
+    core::hint::black_box(trigger_ebpf_program_maps_many);
+}
+
+#[test_log::test]
+fn test_maps_many() {
+    let mut bpf: Ebpf = Ebpf::load(crate::MAP_TEST).unwrap();
+    let prog: &mut UProbe = bpf
+        .program_mut("simple_prog_mut")
+        .unwrap()
+        .try_into()
+        .unwrap();
+
+    prog.load().unwrap();
+    prog.attach(
+        "trigger_ebpf_program_maps_many",
+        "/proc/self/exe",
+        None,
+        None,
+    )
+    .unwrap();
+
+    let [foo, bar, baz] = bpf.maps_many_mut(["FOO", "BAR", "BAZ"]);
+    assert!(foo.is_some());
+    assert!(bar.is_some());
+    assert!(baz.is_none());
+
+    let mut foo: Array<_, u32> = Array::try_from(foo.unwrap()).unwrap();
+    let mut bar: HashMap<_, u32, u8> = HashMap::try_from(bar.unwrap()).unwrap();
+    foo.set(0, 5, 0).unwrap();
+    bar.insert(0, 10, 0).unwrap();
+
+    trigger_ebpf_program_maps_many();
+
+    assert_eq!(foo.get(&0, 0).unwrap(), 6);
+    assert_eq!(bar.get(&0, 0).unwrap(), 11);
+
+    trigger_ebpf_program_maps_many();
+    trigger_ebpf_program_maps_many();
+
+    assert_eq!(foo.get(&0, 0).unwrap(), 8);
+    assert_eq!(bar.get(&0, 0).unwrap(), 13);
+}