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}