aya: maps: add SockHash

aya/src/maps/mod.rs

@@ -48,14 +48,14 @@ mod map_lock;
 pub mod array;
 pub mod hash_map;
 pub mod perf;
-pub mod sock_map;
+pub mod sock;
 pub mod stack_trace;
 
 pub use array::{Array, PerCpuArray, ProgramArray};
 pub use hash_map::{HashMap, PerCpuHashMap};
 pub use map_lock::*;
 pub use perf::PerfEventArray;
-pub use sock_map::SockMap;
+pub use sock::SockMap;
 pub use stack_trace::StackTraceMap;
 
 #[derive(Error, Debug)]

aya/src/maps/sock/mod.rs

@@ -0,0 +1,12 @@
+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;
+pub trait SocketMap {
+    fn fd_or_err(&self) -> Result<RawFd, MapError>;
+}

aya/src/maps/sock/sock_hash.rs

@@ -0,0 +1,127 @@
+use std::{
+    convert::TryFrom,
+    marker::PhantomData,
+    ops::{Deref, DerefMut},
+    os::unix::io::{AsRawFd, RawFd},
+};
+
+use crate::{
+    generated::bpf_map_type::BPF_MAP_TYPE_SOCKHASH,
+    maps::{
+        hash_map, sock::SocketMap, IterableMap, Map, MapError, MapIter, MapKeys, MapRef, MapRefMut,
+    },
+    sys::bpf_map_lookup_elem,
+    Pod,
+};
+
+/// A hash map that can be shared between eBPF programs and user space.
+///
+/// It is required that both keys and values implement the [`Pod`] trait.
+///
+/// # Example
+///
+/// ```no_run
+/// # let bpf = aya::Bpf::load(&[], None)?;
+/// use aya::maps::SockHash;
+/// use std::convert::TryFrom;
+///
+/// const CONFIG_KEY_NUM_RETRIES: u8 = 1;
+///
+/// let mut hm = SockHash::try_from(bpf.map_mut("CONFIG")?)?;
+/// hm.insert(CONFIG_KEY_NUM_RETRIES, 3, 0 /* flags */);
+/// # Ok::<(), aya::BpfError>(())
+/// ```
+#[doc(alias = "BPF_MAP_TYPE_SOCKHASH")]
+pub struct SockHash<T: Deref<Target = Map>, K> {
+    inner: T,
+    _k: PhantomData<K>,
+}
+
+impl<T: Deref<Target = Map>, K: Pod> SockHash<T, K> {
+    pub(crate) fn new(map: T) -> Result<SockHash<T, K>, MapError> {
+        let map_type = map.obj.def.map_type;
+
+        // validate the map definition
+        if map_type != BPF_MAP_TYPE_SOCKHASH as u32 {
+            return Err(MapError::InvalidMapType {
+                map_type: map_type as u32,
+            })?;
+        }
+        hash_map::check_kv_size::<K, u32>(&map)?;
+        let _ = map.fd_or_err()?;
+
+        Ok(SockHash {
+            inner: map,
+            _k: PhantomData,
+        })
+    }
+
+    /// Returns a copy of the value associated with the key.
+    pub unsafe fn get(&self, key: &K, flags: u64) -> Result<u32, MapError> {
+        let fd = self.inner.deref().fd_or_err()?;
+        let value = bpf_map_lookup_elem(fd, key, flags).map_err(|(code, io_error)| {
+            MapError::SyscallError {
+                call: "bpf_map_lookup_elem".to_owned(),
+                code,
+                io_error,
+            }
+        })?;
+        value.ok_or(MapError::KeyNotFound)
+    }
+
+    /// An iterator visiting all key-value pairs in arbitrary order. The
+    /// iterator item type is `Result<(K, V), MapError>`.
+    pub unsafe fn iter(&self) -> MapIter<'_, K, u32> {
+        MapIter::new(self)
+    }
+
+    /// An iterator visiting all keys in arbitrary order. The iterator element
+    /// type is `Result<K, MapError>`.
+    pub unsafe fn keys(&self) -> MapKeys<'_, K> {
+        MapKeys::new(&self.inner)
+    }
+}
+
+impl<T: DerefMut<Target = Map>, K: Pod> SockHash<T, K> {
+    /// Inserts a key-value pair into the map.
+    pub fn insert<I: AsRawFd>(&mut self, key: K, value: I, flags: u64) -> Result<(), MapError> {
+        hash_map::insert(&mut self.inner, key, value.as_raw_fd(), flags)
+    }
+
+    /// Removes a key from the map.
+    pub fn remove(&mut self, key: &K) -> Result<(), MapError> {
+        hash_map::remove(&mut self.inner, key)
+    }
+}
+
+impl<T: Deref<Target = Map>, K: Pod> IterableMap<K, u32> for SockHash<T, K> {
+    fn map(&self) -> &Map {
+        &self.inner
+    }
+
+    unsafe fn get(&self, key: &K) -> Result<u32, MapError> {
+        SockHash::get(self, key, 0)
+    }
+}
+
+impl<T: DerefMut<Target = Map>, K: Pod> SocketMap for SockHash<T, K> {
+    fn fd_or_err(&self) -> Result<RawFd, MapError> {
+        self.inner.fd_or_err()
+    }
+}
+
+impl<K: Pod> TryFrom<MapRef> for SockHash<MapRef, K> {
+    type Error = MapError;
+
+    fn try_from(a: MapRef) -> Result<SockHash<MapRef, K>, MapError> {
+        SockHash::new(a)
+    }
+}
+
+impl<K: Pod> TryFrom<MapRefMut> for SockHash<MapRefMut, K> {
+    type Error = MapError;
+
+    fn try_from(a: MapRefMut) -> Result<SockHash<MapRefMut, K>, MapError> {
+        SockHash::new(a)
+    }
+}

aya/src/maps/sock_map.rs → aya/src/maps/sock/sock_map.rs

@@ -4,15 +4,23 @@ use std::{
     convert::TryFrom,
     mem,
     ops::{Deref, DerefMut},
-    os::unix::prelude::RawFd,
+    os::unix::{io::AsRawFd, prelude::RawFd},
 };
 
 use crate::{
     generated::bpf_map_type::BPF_MAP_TYPE_SOCKMAP,
-    maps::{Map, MapError, MapKeys, MapRef, MapRefMut},
+    maps::{sock::SocketMap, Map, MapError, MapKeys, MapRef, MapRefMut},
     sys::{bpf_map_delete_elem, bpf_map_update_elem},
 };
 
+/// An array of TCP or UDP sock objects. Primarly used for doing socket redirect with eBPF helpers.
+///
+/// A sock map can have two eBPF programs attached: one to parse packets and one to provide a
+/// redirect decision on packets. Whenever a sock object is added to the map, the map's programs
+/// are automatically attached to the socket.
+///
+/// # Example
+///
 pub struct SockMap<T: Deref<Target = Map>> {
     pub(crate) inner: T,
 }
@@ -58,20 +66,24 @@ impl<T: Deref<Target = Map>> SockMap<T> {
 }
 
 impl<T: Deref<Target = Map> + DerefMut<Target = Map>> SockMap<T> {
-    pub fn set(&mut self, index: u32, tcp_fd: RawFd, flags: u64) -> Result<(), MapError> {
+    /// Stores a TCP socket into the map.
+    ///
+    /// eBPF programs can then pass `index` to the `bpf_sk_redirect_map()` helper to redirect
+    /// packets to the corresponding socket.
+    pub fn set<I: AsRawFd>(&mut self, index: u32, socket: &I, flags: u64) -> Result<(), MapError> {
         let fd = self.inner.fd_or_err()?;
         self.check_bounds(index)?;
-        bpf_map_update_elem(fd, &index, &tcp_fd, flags).map_err(|(code, io_error)| {
-            MapError::SyscallError {
+        bpf_map_update_elem(fd, &index, &socket.as_raw_fd(), flags).map_err(
+            |(code, io_error)| MapError::SyscallError {
                 call: "bpf_map_update_elem".to_owned(),
                 code,
                 io_error,
-            }
-        })?;
+            },
+        )?;
         Ok(())
     }
 
-    /// Clears the value at index in the jump table.
+    /// Removes the TCP socket stored at `index` from the map.
     pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> {
         let fd = self.inner.fd_or_err()?;
         self.check_bounds(*index)?;
@@ -85,6 +97,12 @@ impl<T: Deref<Target = Map> + DerefMut<Target = Map>> SockMap<T> {
     }
 }
 
+impl<T: Deref<Target = Map> + DerefMut<Target = Map>> SocketMap for SockMap<T> {
+    fn fd_or_err(&self) -> Result<RawFd, MapError> {
+        self.inner.fd_or_err()
+    }
+}
+
 impl TryFrom<MapRef> for SockMap<MapRef> {
     type Error = MapError;
 

aya/src/programs/sk_skb.rs

@@ -5,7 +5,7 @@ use crate::{
         bpf_attach_type::{BPF_SK_SKB_STREAM_PARSER, BPF_SK_SKB_STREAM_VERDICT},
         bpf_prog_type::BPF_PROG_TYPE_SK_SKB,
     },
-    maps::{Map, SockMap},
+    maps::{sock::SocketMap, Map, SockMap},
     programs::{load_program, LinkRef, ProgAttachLink, ProgramData, ProgramError},
     sys::bpf_prog_attach,
 };
@@ -35,12 +35,10 @@ impl SkSkb {
         self.data.name.to_string()
     }
 
-    pub fn attach<T: Deref<Target = Map>>(
-        &mut self,
-        map: &SockMap<T>,
-    ) -> Result<LinkRef, ProgramError> {
+    /// Attaches the program to the given sockmap.
+    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<LinkRef, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
-        let map_fd = map.inner.fd_or_err()?;
+        let map_fd = map.fd_or_err()?;
 
         let attach_type = match self.kind {
             SkSkbKind::StreamParser => BPF_SK_SKB_STREAM_PARSER,