Use SockMapFd

Create a new type called `SockMapFd` which is solely used when a program needs to attach to a socket map. In the future this same tatic could be used for other use cases so we may make this more generic. Signed-off-by: Andrew Stoycos <astoycos@redhat.com>
aya-rs · Oct 17, 2022 · 898a14d · 898a14d
1 parent 1aefa2e
commit 898a14d
Show file tree

Hide file tree

Showing 6 changed files with 71 additions and 47 deletions.
diff --git a/aya/src/maps/mod.rs b/aya/src/maps/mod.rs
@@ -22,10 +22,12 @@
 //! use aya::maps::SockMap;
 //! use aya::programs::SkMsg;
 //!
-//! let intercept_egress: SockMap<_> = bpf.take_map("INTERCEPT_EGRESS")?.try_into()?;
+//! let mut intercept_egress: SockMap<_> = bpf.map_mut("INTERCEPT_EGRESS")?.try_into()?;
+//! let map_fd = intercept_egress.fd()?;
 //! let prog: &mut SkMsg = bpf.program_mut("intercept_egress_packet").unwrap().try_into()?;
 //! prog.load()?;
-//! prog.attach(&intercept_egress)?;
+//! prog.attach(map_fd)?;
+//!
 //! # Ok::<(), aya::BpfError>(())
 //! ```
 //!

diff --git a/aya/src/maps/sock/mod.rs b/aya/src/maps/sock/mod.rs
@@ -2,15 +2,17 @@
 mod sock_hash;
 mod sock_map;
 
-use std::os::unix::io::RawFd;
-
-use crate::maps::MapError;
-
 pub use sock_hash::SockHash;
 pub use sock_map::SockMap;
 
-/// Shared behaviour between [`SockHash`] and [`SockMap`]
-pub trait SocketMap {
-    /// Returns a [`Result`] containg the map fd or an error if there is none
-    fn fd_or_err(&self) -> Result<RawFd, MapError>;
+use std::os::unix::io::{AsRawFd, RawFd};
+
+/// A [`SocketMap`] file descriptor.
+#[derive(Copy, Clone)]
+pub struct SockMapFd(RawFd);
+
+impl AsRawFd for SockMapFd {
+    fn as_raw_fd(&self) -> RawFd {
+        self.0
+    }
 }
diff --git a/aya/src/maps/sock/sock_hash.rs b/aya/src/maps/sock/sock_hash.rs
@@ -5,7 +5,7 @@ use std::{
 };
 
 use crate::{
-    maps::{hash_map, sock::SocketMap, IterableMap, MapData, MapError, MapIter, MapKeys},
+    maps::{hash_map, sock::SockMapFd, IterableMap, MapData, MapError, MapIter, MapKeys},
     sys::bpf_map_lookup_elem,
     Pod,
 };
@@ -44,12 +44,16 @@ use crate::{
 /// use aya::maps::SockHash;
 /// use aya::programs::SkMsg;
 ///
-/// let mut intercept_egress: SockHash<_, u32> = bpf.take_map("INTERCEPT_EGRESS")?.try_into()?;
+/// let intercept_egress: SockHash<_, u32> = bpf.map("INTERCEPT_EGRESS")?.try_into()?;
+/// let map_fd = intercept_egress.fd()?;
+///
 /// let prog: &mut SkMsg = bpf.program_mut("intercept_egress_packet").unwrap().try_into()?;
 /// prog.load()?;
-/// prog.attach(&intercept_egress)?;
+/// prog.attach(map_fd)?;
 ///
 /// let mut client = TcpStream::connect("127.0.0.1:1234")?;
+/// let mut intercept_egress: SockHash<_, u32> = bpf.map_mut("INTERCEPT_EGRESS")?.try_into()?;
+///
 /// intercept_egress.insert(1234, client.as_raw_fd(), 0)?;
 ///
 /// // the write will be intercepted
@@ -97,6 +101,12 @@ impl<T: AsRef<MapData>, K: Pod> SockHash<T, K> {
     pub fn keys(&self) -> MapKeys<'_, K> {
         MapKeys::new(self.inner.as_ref())
     }
+
+    /// Returns the map's file descriptor, used for instances where programs
+    /// are attached to maps.
+    pub fn fd(&self) -> Result<SockMapFd, MapError> {
+        Ok(SockMapFd(self.inner.as_ref().fd_or_err()?))
+    }
 }
 
 impl<T: AsMut<MapData>, K: Pod> SockHash<T, K> {
@@ -120,9 +130,3 @@ impl<T: AsRef<MapData>, K: Pod> IterableMap<K, RawFd> for SockHash<T, K> {
         SockHash::get(self, key, 0)
     }
 }
-
-impl<T: AsRef<MapData>, K: Pod> SocketMap for SockHash<T, K> {
-    fn fd_or_err(&self) -> Result<RawFd, MapError> {
-        self.inner.as_ref().fd_or_err()
-    }
-}
diff --git a/aya/src/maps/sock/sock_map.rs b/aya/src/maps/sock/sock_map.rs
@@ -7,7 +7,7 @@ use std::{
 };
 
 use crate::{
-    maps::{sock::SocketMap, MapData, MapError, MapKeys},
+    maps::{sock::SockMapFd, MapData, MapError, MapKeys},
     sys::{bpf_map_delete_elem, bpf_map_update_elem},
 };
 
@@ -31,14 +31,17 @@ use crate::{
 /// use aya::maps::SockMap;
 /// use aya::programs::SkSkb;
 ///
-/// let intercept_ingress: SockMap<_> = bpf.take_map("INTERCEPT_INGRESS")?.try_into()?;
+/// let intercept_ingress: SockMap<_> = bpf.map("INTERCEPT_INGRESS")?.try_into()?;
+/// let map_fd = intercept_ingress.fd()?;
+///
 /// let prog: &mut SkSkb = bpf.program_mut("intercept_ingress_packet").unwrap().try_into()?;
 /// prog.load()?;
-/// prog.attach(&intercept_ingress)?;
+/// prog.attach(map_fd)?;
+///
 /// # Ok::<(), aya::BpfError>(())
 /// ```
 #[doc(alias = "BPF_MAP_TYPE_SOCKMAP")]
-pub struct SockMap<T: AsRef<MapData>> {
+pub struct SockMap<T> {
     pub(crate) inner: T,
 }
 
@@ -67,20 +70,17 @@ impl<T: AsRef<MapData>> SockMap<T> {
         MapKeys::new(self.inner.as_ref())
     }
 
-    fn check_bounds(&self, index: u32) -> Result<(), MapError> {
-        let max_entries = self.inner.as_ref().obj.max_entries();
-        if index >= self.inner.as_ref().obj.max_entries() {
-            Err(MapError::OutOfBounds { index, max_entries })
-        } else {
-            Ok(())
-        }
+    /// Returns the map's file descriptor, used for instances where programs
+    /// are attached to maps.
+    pub fn fd(&self) -> Result<SockMapFd, MapError> {
+        Ok(SockMapFd(self.inner.as_ref().fd_or_err()?))
     }
 }
 
-impl<T: AsRef<MapData> + AsMut<MapData>> SockMap<T> {
+impl<T: AsMut<MapData>> SockMap<T> {
     /// Stores a socket into the map.
     pub fn set<I: AsRawFd>(&mut self, index: u32, socket: &I, flags: u64) -> Result<(), MapError> {
-        let fd = self.inner.as_ref().fd_or_err()?;
+        let fd = self.inner.as_mut().fd_or_err()?;
         self.check_bounds(index)?;
         bpf_map_update_elem(fd, Some(&index), &socket.as_raw_fd(), flags).map_err(
             |(_, io_error)| MapError::SyscallError {
@@ -93,7 +93,7 @@ impl<T: AsRef<MapData> + AsMut<MapData>> SockMap<T> {
 
     /// Removes the socket stored at `index` from the map.
     pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> {
-        let fd = self.inner.as_ref().fd_or_err()?;
+        let fd = self.inner.as_mut().fd_or_err()?;
         self.check_bounds(*index)?;
         bpf_map_delete_elem(fd, index)
             .map(|_| ())
@@ -102,10 +102,13 @@ impl<T: AsRef<MapData> + AsMut<MapData>> SockMap<T> {
                 io_error,
             })
     }
-}
 
-impl<T: AsRef<MapData> + AsMut<MapData>> SocketMap for SockMap<T> {
-    fn fd_or_err(&self) -> Result<RawFd, MapError> {
-        self.inner.as_ref().fd_or_err()
+    fn check_bounds(&mut self, index: u32) -> Result<(), MapError> {
+        let max_entries = self.inner.as_mut().obj.max_entries();
+        if index >= self.inner.as_mut().obj.max_entries() {
+            Err(MapError::OutOfBounds { index, max_entries })
+        } else {
+            Ok(())
+        }
     }
 }
diff --git a/aya/src/programs/sk_msg.rs b/aya/src/programs/sk_msg.rs
@@ -1,7 +1,10 @@
 //! Skmsg programs.
+
+use std::os::unix::io::AsRawFd;
+
 use crate::{
     generated::{bpf_attach_type::BPF_SK_MSG_VERDICT, bpf_prog_type::BPF_PROG_TYPE_SK_MSG},
-    maps::sock::SocketMap,
+    maps::sock::SockMapFd,
     programs::{
         define_link_wrapper, load_program, ProgAttachLink, ProgAttachLinkId, ProgramData,
         ProgramError,
@@ -40,12 +43,16 @@ use crate::{
 /// use aya::maps::SockHash;
 /// use aya::programs::SkMsg;
 ///
-/// let mut intercept_egress: SockHash<_, u32> = bpf.take_map("INTERCEPT_EGRESS")?.try_into()?;
+/// let intercept_egress: SockHash<_, u32> = bpf.map("INTERCEPT_EGRESS")?.try_into()?;
+/// let map_fd = intercept_egress.fd()?;
+///
 /// let prog: &mut SkMsg = bpf.program_mut("intercept_egress_packet").unwrap().try_into()?;
 /// prog.load()?;
-/// prog.attach(&intercept_egress)?;
+/// prog.attach(map_fd)?;
 ///
 /// let mut client = TcpStream::connect("127.0.0.1:1234")?;
+/// let mut intercept_egress: SockHash<_, u32> = bpf.map_mut("INTERCEPT_EGRESS")?.try_into()?;
+///
 /// intercept_egress.insert(1234, client.as_raw_fd(), 0)?;
 ///
 /// // the write will be intercepted
@@ -71,9 +78,9 @@ impl SkMsg {
     /// Attaches the program to the given sockmap.
     ///
     /// The returned value can be used to detach, see [SkMsg::detach].
-    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<SkMsgLinkId, ProgramError> {
+    pub fn attach(&mut self, map: SockMapFd) -> Result<SkMsgLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
-        let map_fd = map.fd_or_err()?;
+        let map_fd = map.as_raw_fd();
 
         bpf_prog_attach(prog_fd, map_fd, BPF_SK_MSG_VERDICT).map_err(|(_, io_error)| {
             ProgramError::SyscallError {

diff --git a/aya/src/programs/sk_skb.rs b/aya/src/programs/sk_skb.rs
@@ -1,10 +1,13 @@
 //! Skskb programs.
+
+use std::os::unix::io::AsRawFd;
+
 use crate::{
     generated::{
         bpf_attach_type::{BPF_SK_SKB_STREAM_PARSER, BPF_SK_SKB_STREAM_VERDICT},
         bpf_prog_type::BPF_PROG_TYPE_SK_SKB,
     },
-    maps::sock::SocketMap,
+    maps::sock::SockMapFd,
     programs::{
         define_link_wrapper, load_program, ProgAttachLink, ProgAttachLinkId, ProgramData,
         ProgramError,
@@ -38,10 +41,13 @@ pub enum SkSkbKind {
 /// use aya::maps::SockMap;
 /// use aya::programs::SkSkb;
 ///
-/// let intercept_ingress: SockMap<_> = bpf.take_map("INTERCEPT_INGRESS")?.try_into()?;
+/// let intercept_ingress: SockMap<_> = bpf.map("INTERCEPT_INGRESS")?.try_into()?;
+/// let map_fd = intercept_ingress.fd()?;
+///
 /// let prog: &mut SkSkb = bpf.program_mut("intercept_ingress_packet").unwrap().try_into()?;
 /// prog.load()?;
-/// prog.attach(&intercept_ingress)?;
+/// prog.attach(map_fd)?;
+///
 /// # Ok::<(), aya::BpfError>(())
 /// ```
 ///
@@ -64,9 +70,9 @@ impl SkSkb {
     /// Attaches the program to the given socket map.
     ///
     /// The returned value can be used to detach, see [SkSkb::detach].
-    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<SkSkbLinkId, ProgramError> {
+    pub fn attach(&mut self, map: SockMapFd) -> Result<SkSkbLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
-        let map_fd = map.fd_or_err()?;
+        let map_fd = map.as_raw_fd();
 
         let attach_type = match self.kind {
             SkSkbKind::StreamParser => BPF_SK_SKB_STREAM_PARSER,