From d8f5497884a23bd63f9264dbe3f80081b76f360d Mon Sep 17 00:00:00 2001 From: JPaja <44926986+JPaja@users.noreply.github.com> Date: Mon, 10 Nov 2025 20:12:04 +0100 Subject: [PATCH] feat: add `Ebpf::maps_disjoint_mut` Implemented using `HashMap::get_disjoint_mut` introduced in Rust 1.86.0. --- aya/src/bpf.rs | 29 ++++++++++++ test/integration-ebpf/src/map_test.rs | 21 ++++++++- test/integration-test/src/tests.rs | 1 + .../src/tests/maps_disjoint.rs | 44 +++++++++++++++++++ xtask/public-api/aya.txt | 1 + 5 files changed, 94 insertions(+), 2 deletions(-) create mode 100644 test/integration-test/src/tests/maps_disjoint.rs diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index 609ae553..90ad428a 100644 --- a/aya/src/bpf.rs +++ b/aya/src/bpf.rs @@ -1055,6 +1055,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_disjoint_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_disjoint_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..0135487a 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(array_value) = FOO.get_ptr_mut(0) { + unsafe { + *array_value += 1; + } + } + + if let Some(map_value) = BAZ.get_ptr_mut(0) { + unsafe { + *map_value += 1; + } + } + + 0 +} diff --git a/test/integration-test/src/tests.rs b/test/integration-test/src/tests.rs index b7d4d492..8c605313 100644 --- a/test/integration-test/src/tests.rs +++ b/test/integration-test/src/tests.rs @@ -10,6 +10,7 @@ mod load; mod log; mod lsm; mod map_pin; +mod maps_disjoint; mod raw_tracepoint; mod rbpf; mod relocations; diff --git a/test/integration-test/src/tests/maps_disjoint.rs b/test/integration-test/src/tests/maps_disjoint.rs new file mode 100644 index 00000000..79f6fe23 --- /dev/null +++ b/test/integration-test/src/tests/maps_disjoint.rs @@ -0,0 +1,44 @@ +use aya::{ + Ebpf, + maps::{Array, HashMap}, + programs::UProbe, +}; + +#[unsafe(no_mangle)] +#[inline(never)] +extern "C" fn trigger_ebpf_program_maps_disjoint() { + core::hint::black_box(trigger_ebpf_program_maps_disjoint); +} + +#[test_log::test] +fn test_maps_disjoint() { + 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_disjoint", + "/proc/self/exe", + None, + None, + ) + .unwrap(); + + let [foo, bar, baz] = bpf.maps_disjoint_mut(["FOO", "BAR", "BAZ"]); + + let mut foo: Array<_, u32> = Array::try_from(foo.unwrap()).unwrap(); + let mut bar: HashMap<_, u32, u8> = HashMap::try_from(bar.unwrap()).unwrap(); + assert!(baz.is_none()); + + foo.set(0, 5, 0).unwrap(); + bar.insert(0, 10, 0).unwrap(); + + trigger_ebpf_program_maps_disjoint(); + + assert_eq!(foo.get(&0, 0).unwrap(), 6); + assert_eq!(bar.get(&0, 0).unwrap(), 11); +} diff --git a/xtask/public-api/aya.txt b/xtask/public-api/aya.txt index 0bd2328f..27a51ec5 100644 --- a/xtask/public-api/aya.txt +++ b/xtask/public-api/aya.txt @@ -10629,6 +10629,7 @@ pub fn aya::Ebpf::load_file>(path: P) - pub fn aya::Ebpf::map(&self, name: &str) -> core::option::Option<&aya::maps::Map> pub fn aya::Ebpf::map_mut(&mut self, name: &str) -> core::option::Option<&mut aya::maps::Map> pub fn aya::Ebpf::maps(&self) -> impl core::iter::traits::iterator::Iterator +pub fn aya::Ebpf::maps_disjoint_mut(&mut self, names: [&str; N]) -> [core::option::Option<&mut aya::maps::Map>; N] pub fn aya::Ebpf::maps_mut(&mut self) -> impl core::iter::traits::iterator::Iterator pub fn aya::Ebpf::program(&self, name: &str) -> core::option::Option<&aya::programs::Program> pub fn aya::Ebpf::program_mut(&mut self, name: &str) -> core::option::Option<&mut aya::programs::Program>