diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index 5055e923..9a448240 100644 --- a/aya/src/bpf.rs +++ b/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(&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 diff --git a/test/integration-ebpf/src/map_test.rs b/test/integration-ebpf/src/map_test.rs index eac1823d..4728ca05 100644 --- a/test/integration-ebpf/src/map_test.rs +++ b/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 +} diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index b7d4d492..51e7a2bb 100644 --- a/test/integration-test/src/tests.rs +++ b/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; diff --git a/test/integration-test/src/tests/map_many.rs b/test/integration-test/src/tests/map_many.rs new file mode 100644 index 00000000..b04bdc79 --- /dev/null +++ b/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); +}