Skip to main content

aya_friday/maps/sock/
sock_hash.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    marker::PhantomData,
4    os::fd::{AsRawFd, RawFd},
5};
6
7use crate::{
8    Pod,
9    maps::{
10        IterableMap, MapData, MapError, MapIter, MapKeys, check_kv_size, hash_map, sock::SockMapFd,
11    },
12};
13
14/// A hash map of TCP or UDP sockets.
15///
16/// A `SockHash` is used to store TCP or UDP sockets. eBPF programs can then be
17/// attached to the map to inspect, filter or redirect network buffers on those
18/// sockets.
19///
20/// A `SockHash` can also be used to redirect packets to sockets contained by the
21/// map using `bpf_redirect_map()`, `bpf_sk_redirect_hash()` etc.
22///
23/// # Minimum kernel version
24///
25/// The minimum kernel version required to use this feature is 4.18.
26///
27/// # Examples
28///
29/// ```no_run
30/// # #[derive(Debug, thiserror::Error)]
31/// # enum Error {
32/// #     #[error(transparent)]
33/// #     IO(#[from] std::io::Error),
34/// #     #[error(transparent)]
35/// #     Map(#[from] aya::maps::MapError),
36/// #     #[error(transparent)]
37/// #     Program(#[from] aya::programs::ProgramError),
38/// #     #[error(transparent)]
39/// #     Ebpf(#[from] aya::EbpfError)
40/// # }
41/// # let mut bpf = aya::Ebpf::load(&[])?;
42/// use std::io::Write;
43/// use std::net::TcpStream;
44/// use std::os::fd::AsRawFd;
45/// use aya::maps::SockHash;
46/// use aya::programs::SkMsg;
47///
48/// let mut intercept_egress = SockHash::<_, u32>::try_from(bpf.map("INTERCEPT_EGRESS").unwrap())?;
49/// let map_fd = intercept_egress.fd().try_clone()?;
50///
51/// let prog: &mut SkMsg = bpf.program_mut("intercept_egress_packet").unwrap().try_into()?;
52/// prog.load()?;
53/// prog.attach(&map_fd)?;
54///
55/// let mut client = TcpStream::connect("127.0.0.1:1234")?;
56/// let mut intercept_egress = SockHash::try_from(bpf.map_mut("INTERCEPT_EGRESS").unwrap())?;
57///
58/// intercept_egress.insert(1234, client.as_raw_fd(), 0)?;
59///
60/// // the write will be intercepted
61/// client.write_all(b"foo")?;
62/// # Ok::<(), Error>(())
63/// ```
64#[doc(alias = "BPF_MAP_TYPE_SOCKHASH")]
65pub struct SockHash<T, K> {
66    pub(crate) inner: T,
67    _k: PhantomData<K>,
68}
69
70impl<T: Borrow<MapData>, K: Pod> SockHash<T, K> {
71    pub(crate) fn new(map: T) -> Result<Self, MapError> {
72        let data = map.borrow();
73        check_kv_size::<K, u32>(data)?;
74
75        Ok(Self {
76            inner: map,
77            _k: PhantomData,
78        })
79    }
80
81    /// Returns the fd of the socket stored at the given key.
82    pub fn get(&self, key: &K, flags: u64) -> Result<RawFd, MapError> {
83        hash_map::get(self.inner.borrow(), key, flags)
84    }
85
86    /// An iterator visiting all key-value pairs in arbitrary order. The
87    /// iterator item type is `Result<(K, V), MapError>`.
88    pub fn iter(&self) -> MapIter<'_, K, RawFd, Self> {
89        MapIter::new(self)
90    }
91
92    /// An iterator visiting all keys in arbitrary order. The iterator element
93    /// type is `Result<K, MapError>`.
94    pub fn keys(&self) -> MapKeys<'_, K> {
95        MapKeys::new(self.inner.borrow())
96    }
97
98    /// Returns the map's file descriptor.
99    ///
100    /// The returned file descriptor can be used to attach programs that work with
101    /// socket maps, like [`SkMsg`](crate::programs::SkMsg) and [`SkSkb`](crate::programs::SkSkb).
102    pub fn fd(&self) -> &SockMapFd {
103        let fd = self.inner.borrow().fd();
104        // TODO(https://github.com/rust-lang/rfcs/issues/3066): avoid this unsafe.
105        // SAFETY: `SockMapFd` is #[repr(transparent)] over `MapFd`.
106        unsafe { std::mem::transmute(fd) }
107    }
108}
109
110impl<'a, T: Borrow<MapData>, K: Pod> IntoIterator for &'a SockHash<T, K> {
111    type Item = Result<(K, RawFd), MapError>;
112    type IntoIter = MapIter<'a, K, RawFd, SockHash<T, K>>;
113
114    fn into_iter(self) -> Self::IntoIter {
115        self.iter()
116    }
117}
118
119impl<T: BorrowMut<MapData>, K: Pod> SockHash<T, K> {
120    /// Inserts a socket under the given key.
121    pub fn insert<I: AsRawFd>(
122        &mut self,
123        key: impl Borrow<K>,
124        value: I,
125        flags: u64,
126    ) -> Result<(), MapError> {
127        hash_map::insert(
128            self.inner.borrow_mut(),
129            key.borrow(),
130            &value.as_raw_fd(),
131            flags,
132        )
133    }
134
135    /// Removes a socket from the map.
136    pub fn remove(&mut self, key: &K) -> Result<(), MapError> {
137        hash_map::remove(self.inner.borrow_mut(), key)
138    }
139}
140
141impl<T: Borrow<MapData>, K: Pod> IterableMap<K, RawFd> for SockHash<T, K> {
142    fn map(&self) -> &MapData {
143        self.inner.borrow()
144    }
145
146    fn get(&self, key: &K) -> Result<RawFd, MapError> {
147        Self::get(self, key, 0)
148    }
149}