From c5328305a1634b451e1f0dfbbb6c4380dd7a84bd Mon Sep 17 00:00:00 2001 From: Tamir Duberstein Date: Thu, 2 Oct 2025 23:31:59 -0400 Subject: [PATCH] maps: use shared helpers --- aya/src/maps/array/array.rs | 21 +++------------- aya/src/maps/array/program_array.rs | 23 +++--------------- aya/src/maps/hash_map/hash_map.rs | 11 ++------- aya/src/maps/hash_map/mod.rs | 11 ++++++++- aya/src/maps/lpm_trie.rs | 29 ++++------------------ aya/src/maps/sock/sock_hash.rs | 10 ++------ aya/src/maps/sock/sock_map.rs | 23 ++++-------------- aya/src/maps/stack_trace.rs | 12 +++------- aya/src/maps/xdp/xsk_map.rs | 37 +++++++---------------------- 9 files changed, 41 insertions(+), 136 deletions(-) diff --git a/aya/src/maps/array/array.rs b/aya/src/maps/array/array.rs index 86eb2f6a..5f62757e 100644 --- a/aya/src/maps/array/array.rs +++ b/aya/src/maps/array/array.rs @@ -1,13 +1,11 @@ use std::{ borrow::{Borrow, BorrowMut}, marker::PhantomData, - os::fd::AsFd as _, }; use crate::{ Pod, - maps::{IterableMap, MapData, MapError, check_bounds, check_kv_size}, - sys::{SyscallError, bpf_map_lookup_elem, bpf_map_update_elem}, + maps::{IterableMap, MapData, MapError, check_bounds, check_kv_size, hash_map}, }; /// A fixed-size array. @@ -63,13 +61,7 @@ impl, V: Pod> Array { pub fn get(&self, index: &u32, flags: u64) -> Result { let data = self.inner.borrow(); check_bounds(data, *index)?; - let fd = data.fd().as_fd(); - - let value = bpf_map_lookup_elem(fd, index, flags).map_err(|io_error| SyscallError { - call: "bpf_map_lookup_elem", - io_error, - })?; - value.ok_or(MapError::KeyNotFound) + hash_map::get(data, index, flags) } /// An iterator over the elements of the array. The iterator item type is `Result, V: Pod> Array { pub fn set(&mut self, index: u32, value: impl Borrow, flags: u64) -> Result<(), MapError> { let data = self.inner.borrow_mut(); check_bounds(data, index)?; - let fd = data.fd().as_fd(); - bpf_map_update_elem(fd, Some(&index), value.borrow(), flags).map_err(|io_error| { - SyscallError { - call: "bpf_map_update_elem", - io_error, - } - })?; - Ok(()) + hash_map::insert(data, &index, value.borrow(), flags) } } diff --git a/aya/src/maps/array/program_array.rs b/aya/src/maps/array/program_array.rs index 42c32c1a..9f8f9340 100644 --- a/aya/src/maps/array/program_array.rs +++ b/aya/src/maps/array/program_array.rs @@ -6,9 +6,8 @@ use std::{ }; use crate::{ - maps::{MapData, MapError, MapKeys, check_bounds, check_kv_size}, + maps::{MapData, MapError, MapKeys, check_bounds, check_kv_size, hash_map}, programs::ProgramFd, - sys::{SyscallError, bpf_map_delete_elem, bpf_map_update_elem}, }; /// An array of eBPF program file descriptors used as a jump table. @@ -74,16 +73,7 @@ impl> ProgramArray { pub fn set(&mut self, index: u32, program: &ProgramFd, flags: u64) -> Result<(), MapError> { let data = self.inner.borrow_mut(); check_bounds(data, index)?; - let fd = data.fd().as_fd(); - let prog_fd = program.as_fd(); - let prog_fd = prog_fd.as_raw_fd(); - - bpf_map_update_elem(fd, Some(&index), &prog_fd, flags) - .map_err(|io_error| SyscallError { - call: "bpf_map_update_elem", - io_error, - }) - .map_err(Into::into) + hash_map::insert(data, &index, &program.as_fd().as_raw_fd(), flags) } /// Clears the value at index in the jump table. @@ -93,13 +83,6 @@ impl> ProgramArray { pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> { let data = self.inner.borrow_mut(); check_bounds(data, *index)?; - let fd = data.fd().as_fd(); - - bpf_map_delete_elem(fd, index) - .map_err(|io_error| SyscallError { - call: "bpf_map_delete_elem", - io_error, - }) - .map_err(Into::into) + hash_map::remove(data, index) } } diff --git a/aya/src/maps/hash_map/hash_map.rs b/aya/src/maps/hash_map/hash_map.rs index 6010e141..cfb8b953 100644 --- a/aya/src/maps/hash_map/hash_map.rs +++ b/aya/src/maps/hash_map/hash_map.rs @@ -1,13 +1,11 @@ use std::{ borrow::{Borrow, BorrowMut}, marker::PhantomData, - os::fd::AsFd as _, }; use crate::{ Pod, maps::{IterableMap, MapData, MapError, MapIter, MapKeys, check_kv_size, hash_map}, - sys::{SyscallError, bpf_map_lookup_elem}, }; /// A hash map that can be shared between eBPF programs and user space. @@ -53,12 +51,7 @@ impl, K: Pod, V: Pod> HashMap { /// Returns a copy of the value associated with the key. pub fn get(&self, key: &K, flags: u64) -> Result { - let fd = self.inner.borrow().fd().as_fd(); - let value = bpf_map_lookup_elem(fd, key, flags).map_err(|io_error| SyscallError { - call: "bpf_map_lookup_elem", - io_error, - })?; - value.ok_or(MapError::KeyNotFound) + hash_map::get(self.inner.borrow(), key, flags) } /// An iterator visiting all key-value pairs in arbitrary order. The @@ -118,7 +111,7 @@ mod tests { Map, test_utils::{self, new_map}, }, - sys::{SysResult, Syscall, override_syscall}, + sys::{SysResult, Syscall, SyscallError, override_syscall}, }; fn new_obj_map() -> aya_obj::Map { diff --git a/aya/src/maps/hash_map/mod.rs b/aya/src/maps/hash_map/mod.rs index 47139ddc..59dc65f5 100644 --- a/aya/src/maps/hash_map/mod.rs +++ b/aya/src/maps/hash_map/mod.rs @@ -4,7 +4,7 @@ use std::os::fd::AsFd as _; use crate::{ Pod, maps::MapError, - sys::{SyscallError, bpf_map_delete_elem, bpf_map_update_elem}, + sys::{SyscallError, bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem}, }; #[expect(clippy::module_inception)] @@ -16,6 +16,15 @@ pub use per_cpu_hash_map::*; use super::MapData; +pub(crate) fn get(map: &MapData, key: &K, flags: u64) -> Result { + let fd = map.fd().as_fd(); + let value = bpf_map_lookup_elem(fd, key, flags).map_err(|io_error| SyscallError { + call: "bpf_map_lookup_elem", + io_error, + })?; + value.ok_or(MapError::KeyNotFound) +} + pub(crate) fn insert( map: &MapData, key: &K, diff --git a/aya/src/maps/lpm_trie.rs b/aya/src/maps/lpm_trie.rs index 7ec495d5..828f51be 100644 --- a/aya/src/maps/lpm_trie.rs +++ b/aya/src/maps/lpm_trie.rs @@ -3,13 +3,11 @@ use std::{ borrow::{Borrow, BorrowMut}, marker::PhantomData, - os::fd::AsFd as _, }; use crate::{ Pod, - maps::{IterableMap, MapData, MapError, MapIter, MapKeys, check_kv_size}, - sys::{SyscallError, bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem}, + maps::{IterableMap, MapData, MapError, MapIter, MapKeys, check_kv_size, hash_map}, }; /// A Longest Prefix Match Trie. @@ -120,12 +118,7 @@ impl, K: Pod, V: Pod> LpmTrie { /// Returns a copy of the value associated with the longest prefix matching key in the LpmTrie. pub fn get(&self, key: &Key, flags: u64) -> Result { - let fd = self.inner.borrow().fd().as_fd(); - let value = bpf_map_lookup_elem(fd, key, flags).map_err(|io_error| SyscallError { - call: "bpf_map_lookup_elem", - io_error, - })?; - value.ok_or(MapError::KeyNotFound) + hash_map::get(self.inner.borrow(), key, flags) } /// An iterator visiting all key-value pairs. The @@ -149,26 +142,14 @@ impl, K: Pod, V: Pod> LpmTrie { value: impl Borrow, flags: u64, ) -> Result<(), MapError> { - let fd = self.inner.borrow().fd().as_fd(); - bpf_map_update_elem(fd, Some(key), value.borrow(), flags) - .map_err(|io_error| SyscallError { - call: "bpf_map_update_elem", - io_error, - }) - .map_err(Into::into) + hash_map::insert(self.inner.borrow_mut(), key, value.borrow(), flags) } /// Removes an element from the map. /// /// Both the prefix and data must match exactly - this method does not do a longest prefix match. pub fn remove(&mut self, key: &Key) -> Result<(), MapError> { - let fd = self.inner.borrow().fd().as_fd(); - bpf_map_delete_elem(fd, key) - .map_err(|io_error| SyscallError { - call: "bpf_map_delete_elem", - io_error, - }) - .map_err(Into::into) + hash_map::remove(self.inner.borrow_mut(), key) } } @@ -199,7 +180,7 @@ mod tests { Map, test_utils::{self, new_map}, }, - sys::{SysResult, Syscall, override_syscall}, + sys::{SysResult, Syscall, SyscallError, override_syscall}, }; fn new_obj_map() -> aya_obj::Map { diff --git a/aya/src/maps/sock/sock_hash.rs b/aya/src/maps/sock/sock_hash.rs index 63efd703..cc5eb03d 100644 --- a/aya/src/maps/sock/sock_hash.rs +++ b/aya/src/maps/sock/sock_hash.rs @@ -1,7 +1,7 @@ use std::{ borrow::{Borrow, BorrowMut}, marker::PhantomData, - os::fd::{AsFd as _, AsRawFd, RawFd}, + os::fd::{AsRawFd, RawFd}, }; use crate::{ @@ -10,7 +10,6 @@ use crate::{ IterableMap, MapData, MapError, MapFd, MapIter, MapKeys, check_kv_size, hash_map, sock::SockMapFd, }, - sys::{SyscallError, bpf_map_lookup_elem}, }; /// A hash map of TCP or UDP sockets. @@ -82,12 +81,7 @@ impl, K: Pod> SockHash { /// Returns the fd of the socket stored at the given key. pub fn get(&self, key: &K, flags: u64) -> Result { - let fd = self.inner.borrow().fd().as_fd(); - let value = bpf_map_lookup_elem(fd, key, flags).map_err(|io_error| SyscallError { - call: "bpf_map_lookup_elem", - io_error, - })?; - value.ok_or(MapError::KeyNotFound) + hash_map::get(self.inner.borrow(), key, flags) } /// An iterator visiting all key-value pairs in arbitrary order. The diff --git a/aya/src/maps/sock/sock_map.rs b/aya/src/maps/sock/sock_map.rs index d139730a..4e55b88a 100644 --- a/aya/src/maps/sock/sock_map.rs +++ b/aya/src/maps/sock/sock_map.rs @@ -2,12 +2,11 @@ use std::{ borrow::{Borrow, BorrowMut}, - os::fd::{AsFd as _, AsRawFd, RawFd}, + os::fd::{AsRawFd, RawFd}, }; -use crate::{ - maps::{MapData, MapError, MapFd, MapKeys, check_bounds, check_kv_size, sock::SockMapFd}, - sys::{SyscallError, bpf_map_delete_elem, bpf_map_update_elem}, +use crate::maps::{ + MapData, MapError, MapFd, MapKeys, check_bounds, check_kv_size, hash_map, sock::SockMapFd, }; /// An array of TCP or UDP sockets. @@ -85,26 +84,14 @@ impl> SockMap { /// Stores a socket into the map. pub fn set(&mut self, index: u32, socket: &I, flags: u64) -> Result<(), MapError> { let data = self.inner.borrow_mut(); - let fd = data.fd().as_fd(); check_bounds(data, index)?; - bpf_map_update_elem(fd, Some(&index), &socket.as_raw_fd(), flags) - .map_err(|io_error| SyscallError { - call: "bpf_map_update_elem", - io_error, - }) - .map_err(Into::into) + hash_map::insert(data, &index, &socket.as_raw_fd(), flags) } /// Removes the socket stored at `index` from the map. pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> { let data = self.inner.borrow_mut(); - let fd = data.fd().as_fd(); check_bounds(data, *index)?; - bpf_map_delete_elem(fd, index) - .map_err(|io_error| SyscallError { - call: "bpf_map_delete_elem", - io_error, - }) - .map_err(Into::into) + hash_map::remove(data, index) } } diff --git a/aya/src/maps/stack_trace.rs b/aya/src/maps/stack_trace.rs index 0c6a49c9..0d324980 100644 --- a/aya/src/maps/stack_trace.rs +++ b/aya/src/maps/stack_trace.rs @@ -11,8 +11,8 @@ use std::{ }; use crate::{ - maps::{IterableMap, MapData, MapError, MapIter, MapKeys}, - sys::{SyscallError, bpf_map_delete_elem, bpf_map_lookup_elem_ptr}, + maps::{IterableMap, MapData, MapError, MapIter, MapKeys, hash_map}, + sys::{SyscallError, bpf_map_lookup_elem_ptr}, }; /// A hash map of kernel or user space stack traces. @@ -167,13 +167,7 @@ impl<'a, T: Borrow> IntoIterator for &'a StackTraceMap { impl> StackTraceMap { /// Removes the stack trace with the given stack_id. pub fn remove(&mut self, stack_id: &u32) -> Result<(), MapError> { - let fd = self.inner.borrow().fd().as_fd(); - bpf_map_delete_elem(fd, stack_id) - .map_err(|io_error| SyscallError { - call: "bpf_map_delete_elem", - io_error, - }) - .map_err(Into::into) + hash_map::remove(self.inner.borrow_mut(), stack_id) } } diff --git a/aya/src/maps/xdp/xsk_map.rs b/aya/src/maps/xdp/xsk_map.rs index fee7be2d..2b3cd57f 100644 --- a/aya/src/maps/xdp/xsk_map.rs +++ b/aya/src/maps/xdp/xsk_map.rs @@ -2,13 +2,10 @@ use std::{ borrow::{Borrow, BorrowMut}, - os::fd::{AsFd as _, AsRawFd, BorrowedFd, RawFd}, + os::fd::{AsRawFd, RawFd}, }; -use crate::{ - maps::{MapData, MapError, check_bounds, check_kv_size}, - sys::{SyscallError, bpf_map_delete_elem, bpf_map_update_elem}, -}; +use crate::maps::{MapData, MapError, check_bounds, check_kv_size, hash_map}; /// An array of AF_XDP sockets. /// @@ -57,16 +54,6 @@ impl> XskMap { } impl> XskMap { - fn with_fd( - &mut self, - index: u32, - f: impl FnOnce(BorrowedFd<'_>) -> Result<(), SyscallError>, - ) -> Result<(), MapError> { - let data = self.inner.borrow_mut(); - check_bounds(data, index)?; - f(data.fd().as_fd()).map_err(Into::into) - } - /// Sets the `AF_XDP` socket at a given index. /// /// When redirecting a packet, the `AF_XDP` socket at `index` will recieve the packet. Note @@ -78,14 +65,9 @@ impl> XskMap { /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`] /// if `bpf_map_update_elem` fails. pub fn set(&mut self, index: u32, socket_fd: impl AsRawFd, flags: u64) -> Result<(), MapError> { - self.with_fd(index, |fd| { - bpf_map_update_elem(fd, Some(&index), &socket_fd.as_raw_fd(), flags).map_err( - |io_error| SyscallError { - call: "bpf_map_update_elem", - io_error, - }, - ) - }) + let data = self.inner.borrow_mut(); + check_bounds(data, index)?; + hash_map::insert(data, &index, &socket_fd.as_raw_fd(), flags) } /// Un-sets the `AF_XDP` socket at a given index. @@ -95,11 +77,8 @@ impl> XskMap { /// Returns [`MapError::OutOfBounds`] if `index` is out of bounds, [`MapError::SyscallError`] /// if `bpf_map_delete_elem` fails. pub fn unset(&mut self, index: u32) -> Result<(), MapError> { - self.with_fd(index, |fd| { - bpf_map_delete_elem(fd, &index).map_err(|io_error| SyscallError { - call: "bpf_map_delete_elem", - io_error, - }) - }) + let data = self.inner.borrow_mut(); + check_bounds(data, index)?; + hash_map::remove(data, &index) } }