aya-ebpf: extract insert,remove,lookup

These functions (and more) are duplicated all over the place.

ebpf/aya-ebpf/src/lib.rs

@@ -24,12 +24,15 @@ pub mod helpers;
 pub mod maps;
 pub mod programs;
 
-use core::ffi::c_void;
+use core::{ffi::c_void, ptr::NonNull};
 
 pub use aya_ebpf_cty as cty;
 pub use aya_ebpf_macros as macros;
 use cty::{c_int, c_long};
-use helpers::{bpf_get_current_comm, bpf_get_current_pid_tgid, bpf_get_current_uid_gid};
+use helpers::{
+    bpf_get_current_comm, bpf_get_current_pid_tgid, bpf_get_current_uid_gid, bpf_map_delete_elem,
+    bpf_map_lookup_elem, bpf_map_update_elem,
+};
 
 pub const TASK_COMM_LEN: usize = 16;
 
@@ -125,3 +128,33 @@ pub fn check_bounds_signed(value: i64, lower: i64, upper: i64) -> bool {
         unimplemented!()
     }
 }
+
+#[inline]
+fn insert<K, V>(
+    def: *mut bindings::bpf_map_def,
+    key: &K,
+    value: &V,
+    flags: u64,
+) -> Result<(), c_long> {
+    let key: *const _ = key;
+    let value: *const _ = value;
+    match unsafe { bpf_map_update_elem(def.cast(), key.cast(), value.cast(), flags) } {
+        0 => Ok(()),
+        ret => Err(ret),
+    }
+}
+
+#[inline]
+fn remove<K>(def: *mut bindings::bpf_map_def, key: &K) -> Result<(), c_long> {
+    let key: *const _ = key;
+    match unsafe { bpf_map_delete_elem(def.cast(), key.cast()) } {
+        0 => Ok(()),
+        ret => Err(ret),
+    }
+}
+
+#[inline]
+fn lookup<K, V>(def: *mut bindings::bpf_map_def, key: &K) -> Option<NonNull<V>> {
+    let key: *const _ = key;
+    NonNull::new(unsafe { bpf_map_lookup_elem(def.cast(), key.cast()) }.cast())
+}

ebpf/aya-ebpf/src/maps/array.rs

@@ -1,10 +1,8 @@
 use core::{cell::UnsafeCell, marker::PhantomData, mem, ptr::NonNull};
 
-use aya_ebpf_cty::c_void;
-
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_ARRAY},
-    helpers::bpf_map_lookup_elem,
+    lookup,
     maps::PinningType,
 };
 
@@ -65,10 +63,6 @@ impl<T> Array<T> {
 
     #[inline(always)]
     unsafe fn lookup(&self, index: u32) -> Option<NonNull<T>> {
-        let ptr = bpf_map_lookup_elem(
+        lookup(self.def.get(), &index)
-            self.def.get() as *mut _,
-            &index as *const _ as *const c_void,
-        );
-        NonNull::new(ptr as *mut T)
     }
 }

ebpf/aya-ebpf/src/maps/hash_map.rs

@@ -1,14 +1,15 @@
-use core::{cell::UnsafeCell, marker::PhantomData, mem, ptr::NonNull};
+use core::{cell::UnsafeCell, marker::PhantomData, mem};
 
 use aya_ebpf_bindings::bindings::bpf_map_type::{
     BPF_MAP_TYPE_LRU_HASH, BPF_MAP_TYPE_LRU_PERCPU_HASH, BPF_MAP_TYPE_PERCPU_HASH,
 };
-use aya_ebpf_cty::{c_long, c_void};
+use aya_ebpf_cty::c_long;
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_HASH},
-    helpers::{bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem},
+    insert, lookup,
     maps::PinningType,
+    remove,
 };
 
 #[repr(transparent)]
@@ -321,11 +322,7 @@ const fn build_def<K, V>(ty: u32, max_entries: u32, flags: u32, pin: PinningType
 
 #[inline]
 fn get_ptr_mut<K, V>(def: *mut bpf_map_def, key: &K) -> Option<*mut V> {
-    unsafe {
+    lookup(def, key).map(|p| p.as_ptr())
-        let value = bpf_map_lookup_elem(def as *mut _, key as *const _ as *const c_void);
-        // FIXME: alignment
-        NonNull::new(value as *mut V).map(|p| p.as_ptr())
-    }
 }
 
 #[inline]
@@ -337,22 +334,3 @@ fn get_ptr<K, V>(def: *mut bpf_map_def, key: &K) -> Option<*const V> {
 unsafe fn get<'a, K, V>(def: *mut bpf_map_def, key: &K) -> Option<&'a V> {
     get_ptr(def, key).map(|p| &*p)
 }
-
-#[inline]
-fn insert<K, V>(def: *mut bpf_map_def, key: &K, value: &V, flags: u64) -> Result<(), c_long> {
-    let ret = unsafe {
-        bpf_map_update_elem(
-            def as *mut _,
-            key as *const _ as *const _,
-            value as *const _ as *const _,
-            flags,
-        )
-    };
-    (ret == 0).then_some(()).ok_or(ret)
-}
-
-#[inline]
-fn remove<K>(def: *mut bpf_map_def, key: &K) -> Result<(), c_long> {
-    let ret = unsafe { bpf_map_delete_elem(def as *mut _, key as *const _ as *const c_void) };
-    (ret == 0).then_some(()).ok_or(ret)
-}

ebpf/aya-ebpf/src/maps/lpm_trie.rs

@@ -1,12 +1,13 @@
-use core::{cell::UnsafeCell, marker::PhantomData, mem, ptr::NonNull};
+use core::{cell::UnsafeCell, marker::PhantomData, mem};
 
 use aya_ebpf_bindings::bindings::BPF_F_NO_PREALLOC;
-use aya_ebpf_cty::{c_long, c_void};
+use aya_ebpf_cty::c_long;
 
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_LPM_TRIE},
-    helpers::{bpf_map_delete_elem, bpf_map_lookup_elem, bpf_map_update_elem},
+    insert, lookup,
     maps::PinningType,
+    remove,
 };
 
 #[repr(transparent)]
@@ -63,33 +64,17 @@ impl<K, V> LpmTrie<K, V> {
 
     #[inline]
     pub fn get(&self, key: &Key<K>) -> Option<&V> {
-        unsafe {
+        lookup(self.def.get(), key).map(|p| unsafe { p.as_ref() })
-            let value =
-                bpf_map_lookup_elem(self.def.get() as *mut _, key as *const _ as *const c_void);
-            // FIXME: alignment
-            NonNull::new(value as *mut V).map(|p| p.as_ref())
-        }
     }
 
     #[inline]
     pub fn insert(&self, key: &Key<K>, value: &V, flags: u64) -> Result<(), c_long> {
-        let ret = unsafe {
+        insert(self.def.get(), key, value, flags)
-            bpf_map_update_elem(
-                self.def.get() as *mut _,
-                key as *const _ as *const _,
-                value as *const _ as *const _,
-                flags,
-            )
-        };
-        (ret == 0).then_some(()).ok_or(ret)
     }
 
     #[inline]
     pub fn remove(&self, key: &Key<K>) -> Result<(), c_long> {
-        let ret = unsafe {
+        remove(self.def.get(), key)
-            bpf_map_delete_elem(self.def.get() as *mut _, key as *const _ as *const c_void)
-        };
-        (ret == 0).then_some(()).ok_or(ret)
     }
 }
 

ebpf/aya-ebpf/src/maps/per_cpu_array.rs

@@ -1,10 +1,8 @@
 use core::{cell::UnsafeCell, marker::PhantomData, mem, ptr::NonNull};
 
-use aya_ebpf_cty::c_void;
-
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_PERCPU_ARRAY},
-    helpers::bpf_map_lookup_elem,
+    lookup,
     maps::PinningType,
 };
 
@@ -67,10 +65,6 @@ impl<T> PerCpuArray<T> {
 
     #[inline(always)]
     unsafe fn lookup(&self, index: u32) -> Option<NonNull<T>> {
-        let ptr = bpf_map_lookup_elem(
+        lookup(self.def.get(), &index)
-            self.def.get() as *mut _,
-            &index as *const _ as *const c_void,
-        );
-        NonNull::new(ptr as *mut T)
     }
 }

ebpf/aya-ebpf/src/maps/sock_hash.rs

@@ -5,9 +5,10 @@ use aya_ebpf_cty::c_void;
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKHASH, bpf_sock_ops},
     helpers::{
-        bpf_map_lookup_elem, bpf_msg_redirect_hash, bpf_sk_assign, bpf_sk_redirect_hash,
+        bpf_msg_redirect_hash, bpf_sk_assign, bpf_sk_redirect_hash, bpf_sk_release,
-        bpf_sk_release, bpf_sock_hash_update,
+        bpf_sock_hash_update,
     },
+    lookup,
     maps::PinningType,
     programs::{SkBuffContext, SkLookupContext, SkMsgContext},
     EbpfContext,
@@ -92,17 +93,12 @@ impl<K> SockHash<K> {
         key: impl Borrow<K>,
         flags: u64,
     ) -> Result<(), u32> {
-        unsafe {
+        let sk = lookup(self.def.get(), key.borrow()).ok_or(1u32)?;
-            let sk = bpf_map_lookup_elem(
+        let ret = unsafe { bpf_sk_assign(ctx.as_ptr().cast(), sk.as_ptr(), flags) };
-                &mut self.def as *mut _ as *mut _,
+        unsafe { bpf_sk_release(sk.as_ptr()) };
-                &key as *const _ as *const c_void,
+        match ret {
-            );
+            0 => Ok(()),
-            if sk.is_null() {
+            _ret => Err(1),
-                return Err(1);
-            }
-            let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags);
-            bpf_sk_release(sk);
-            (ret == 0).then_some(()).ok_or(1)
         }
     }
 }

ebpf/aya-ebpf/src/maps/sock_map.rs

@@ -5,9 +5,10 @@ use aya_ebpf_cty::c_void;
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKMAP, bpf_sock_ops},
     helpers::{
-        bpf_map_lookup_elem, bpf_msg_redirect_map, bpf_sk_assign, bpf_sk_redirect_map,
+        bpf_msg_redirect_map, bpf_sk_assign, bpf_sk_redirect_map, bpf_sk_release,
-        bpf_sk_release, bpf_sock_map_update,
+        bpf_sock_map_update,
     },
+    lookup,
     maps::PinningType,
     programs::{SkBuffContext, SkLookupContext, SkMsgContext},
     EbpfContext,
@@ -92,17 +93,12 @@ impl SockMap {
         index: u32,
         flags: u64,
     ) -> Result<(), u32> {
-        unsafe {
+        let sk = lookup(self.def.get(), &index).ok_or(1u32)?;
-            let sk = bpf_map_lookup_elem(
+        let ret = unsafe { bpf_sk_assign(ctx.as_ptr().cast(), sk.as_ptr(), flags) };
-                &mut self.def as *mut _ as *mut _,
+        unsafe { bpf_sk_release(sk.as_ptr()) };
-                &index as *const _ as *const c_void,
+        match ret {
-            );
+            0 => Ok(()),
-            if sk.is_null() {
+            _ret => Err(1),
-                return Err(1);
-            }
-            let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags);
-            bpf_sk_release(sk);
-            (ret == 0).then_some(()).ok_or(1)
         }
     }
 }

ebpf/aya-ebpf/src/maps/xdp/dev_map.rs

@@ -1,12 +1,11 @@
-use core::{cell::UnsafeCell, mem, num::NonZeroU32, ptr::NonNull};
+use core::{cell::UnsafeCell, mem, num::NonZeroU32};
 
 use aya_ebpf_bindings::bindings::bpf_devmap_val;
-use aya_ebpf_cty::c_void;
 
 use super::try_redirect_map;
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_DEVMAP},
-    helpers::bpf_map_lookup_elem,
+    lookup,
     maps::PinningType,
 };
 
@@ -106,18 +105,14 @@ impl DevMap {
     /// ```
     #[inline(always)]
     pub fn get(&self, index: u32) -> Option<DevMapValue> {
-        unsafe {
+        let value = lookup(self.def.get(), &index)?;
-            let value = bpf_map_lookup_elem(
+        let value: &bpf_devmap_val = unsafe { value.as_ref() };
-                self.def.get() as *mut _,
+        Some(DevMapValue {
-                &index as *const _ as *const c_void,
+            if_index: value.ifindex,
-            );
+            // SAFETY: map writes use fd, map reads use id.
-            NonNull::new(value as *mut bpf_devmap_val).map(|p| DevMapValue {
+            // https://elixir.bootlin.com/linux/v6.2/source/include/uapi/linux/bpf.h#L6136
-                if_index: p.as_ref().ifindex,
+            prog_id: NonZeroU32::new(unsafe { value.bpf_prog.id }),
-                // SAFETY: map writes use fd, map reads use id.
+        })
-                // https://elixir.bootlin.com/linux/v6.2/source/include/uapi/linux/bpf.h#L6136
-                prog_id: NonZeroU32::new(p.as_ref().bpf_prog.id),
-            })
-        }
     }
 
     /// Redirects the current packet on the interface at `index`.

ebpf/aya-ebpf/src/maps/xdp/dev_map_hash.rs

@@ -1,12 +1,11 @@
-use core::{cell::UnsafeCell, mem, num::NonZeroU32, ptr::NonNull};
+use core::{cell::UnsafeCell, mem, num::NonZeroU32};
 
 use aya_ebpf_bindings::bindings::bpf_devmap_val;
-use aya_ebpf_cty::c_void;
 
 use super::{dev_map::DevMapValue, try_redirect_map};
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_DEVMAP_HASH},
-    helpers::bpf_map_lookup_elem,
+    lookup,
     maps::PinningType,
 };
 
@@ -108,16 +107,14 @@ impl DevMapHash {
     /// ```
     #[inline(always)]
     pub fn get(&self, key: u32) -> Option<DevMapValue> {
-        unsafe {
+        let value = lookup(self.def.get(), &key)?;
-            let value =
+        let value: &bpf_devmap_val = unsafe { value.as_ref() };
-                bpf_map_lookup_elem(self.def.get() as *mut _, &key as *const _ as *const c_void);
+        Some(DevMapValue {
-            NonNull::new(value as *mut bpf_devmap_val).map(|p| DevMapValue {
+            if_index: value.ifindex,
-                if_index: p.as_ref().ifindex,
+            // SAFETY: map writes use fd, map reads use id.
-                // SAFETY: map writes use fd, map reads use id.
+            // https://elixir.bootlin.com/linux/v6.2/source/include/uapi/linux/bpf.h#L6136
-                // https://elixir.bootlin.com/linux/v6.2/source/include/uapi/linux/bpf.h#L6136
+            prog_id: NonZeroU32::new(unsafe { value.bpf_prog.id }),
-                prog_id: NonZeroU32::new(p.as_ref().bpf_prog.id),
+        })
-            })
-        }
     }
 
     /// Redirects the current packet on the interface at `key`.

ebpf/aya-ebpf/src/maps/xdp/xsk_map.rs

@@ -1,12 +1,11 @@
-use core::{cell::UnsafeCell, mem, ptr::NonNull};
+use core::{cell::UnsafeCell, mem};
 
 use aya_ebpf_bindings::bindings::bpf_xdp_sock;
-use aya_ebpf_cty::c_void;
 
 use super::try_redirect_map;
 use crate::{
     bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_XSKMAP},
-    helpers::bpf_map_lookup_elem,
+    lookup,
     maps::PinningType,
 };
 
@@ -125,13 +124,9 @@ impl XskMap {
     /// ```
     #[inline(always)]
     pub fn get(&self, index: u32) -> Option<u32> {
-        unsafe {
+        let value = lookup(self.def.get(), &index)?;
-            let value = bpf_map_lookup_elem(
+        let value: &bpf_xdp_sock = unsafe { value.as_ref() };
-                self.def.get() as *mut _,
+        Some(value.queue_id)
-                &index as *const _ as *const c_void,
-            );
-            NonNull::new(value as *mut bpf_xdp_sock).map(|p| p.as_ref().queue_id)
-        }
     }
 
     /// Redirects the current packet to the AF_XDP socket at `index`.