From dad300c88bdfade10f8318fcbbe9a8b5e4de89ce Mon Sep 17 00:00:00 2001 From: Alessandro Decina Date: Mon, 12 Apr 2021 11:43:04 +0000 Subject: [PATCH] aya: maps: add SockHash --- aya/src/maps/mod.rs | 4 +- aya/src/maps/sock/mod.rs | 12 +++ aya/src/maps/sock/sock_hash.rs | 127 ++++++++++++++++++++++++++++ aya/src/maps/{ => sock}/sock_map.rs | 34 ++++++-- aya/src/programs/sk_skb.rs | 10 +-- 5 files changed, 171 insertions(+), 16 deletions(-) create mode 100644 aya/src/maps/sock/mod.rs create mode 100644 aya/src/maps/sock/sock_hash.rs rename aya/src/maps/{ => sock}/sock_map.rs (69%) diff --git a/aya/src/maps/mod.rs b/aya/src/maps/mod.rs index dd0ff2cd..9d85f479 100644 --- a/aya/src/maps/mod.rs +++ b/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)] diff --git a/aya/src/maps/sock/mod.rs b/aya/src/maps/sock/mod.rs new file mode 100644 index 00000000..d8e05ed5 --- /dev/null +++ b/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; +} diff --git a/aya/src/maps/sock/sock_hash.rs b/aya/src/maps/sock/sock_hash.rs new file mode 100644 index 00000000..3301e955 --- /dev/null +++ b/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, K> { + inner: T, + _k: PhantomData, +} + +impl, K: Pod> SockHash { + pub(crate) fn new(map: T) -> Result, 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::(&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 { + 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`. + pub unsafe fn keys(&self) -> MapKeys<'_, K> { + MapKeys::new(&self.inner) + } +} + +impl, K: Pod> SockHash { + /// Inserts a key-value pair into the map. + pub fn insert(&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, K: Pod> IterableMap for SockHash { + fn map(&self) -> &Map { + &self.inner + } + + unsafe fn get(&self, key: &K) -> Result { + SockHash::get(self, key, 0) + } +} + +impl, K: Pod> SocketMap for SockHash { + fn fd_or_err(&self) -> Result { + self.inner.fd_or_err() + } +} + +impl TryFrom for SockHash { + type Error = MapError; + + fn try_from(a: MapRef) -> Result, MapError> { + SockHash::new(a) + } +} + +impl TryFrom for SockHash { + type Error = MapError; + + fn try_from(a: MapRefMut) -> Result, MapError> { + SockHash::new(a) + } +} diff --git a/aya/src/maps/sock_map.rs b/aya/src/maps/sock/sock_map.rs similarity index 69% rename from aya/src/maps/sock_map.rs rename to aya/src/maps/sock/sock_map.rs index 0d9c4db0..f243d49b 100644 --- a/aya/src/maps/sock_map.rs +++ b/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> { pub(crate) inner: T, } @@ -58,20 +66,24 @@ impl> SockMap { } impl + DerefMut> SockMap { - 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(&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 + DerefMut> SockMap { } } +impl + DerefMut> SocketMap for SockMap { + fn fd_or_err(&self) -> Result { + self.inner.fd_or_err() + } +} + impl TryFrom for SockMap { type Error = MapError; diff --git a/aya/src/programs/sk_skb.rs b/aya/src/programs/sk_skb.rs index 1683b579..de6abd3b 100644 --- a/aya/src/programs/sk_skb.rs +++ b/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>( - &mut self, - map: &SockMap, - ) -> Result { + /// Attaches the program to the given sockmap. + pub fn attach(&mut self, map: &dyn SocketMap) -> Result { 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,