diff --git a/Cargo.toml b/Cargo.toml index 73f71433..3d4c27ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,7 +57,7 @@ edition = "2024" homepage = "https://aya-rs.dev" license = "MIT OR Apache-2.0" repository = "https://github.com/aya-rs/aya" -rust-version = "1.85.0" +rust-version = "1.86.0" # NOTE(vadorovsky): Neither cargo-udeps nor cargo-machete are able to detect # unused crates defined in this section. It would be nice to teach either of diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index 609ae553..7da550a1 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_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..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..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); +} diff --git a/xtask/public-api/aya.txt b/xtask/public-api/aya.txt index bc861f4c..dd7ab4f8 100644 --- a/xtask/public-api/aya.txt +++ b/xtask/public-api/aya.txt @@ -10708,6 +10708,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_many_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>