Skip to content

Commit

Permalink
Use SockMapFd
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
astoycos committed Oct 17, 2022
1 parent 1aefa2e commit 898a14d
Show file tree
Hide file tree
Showing 6 changed files with 71 additions and 47 deletions.
6 changes: 4 additions & 2 deletions aya/src/maps/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>(())
//! ```
//!
Expand Down
18 changes: 10 additions & 8 deletions aya/src/maps/sock/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
22 changes: 13 additions & 9 deletions aya/src/maps/sock/sock_hash.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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> {
Expand All @@ -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()
}
}
39 changes: 21 additions & 18 deletions aya/src/maps/sock/sock_map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
};

Expand All @@ -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,
}

Expand Down Expand Up @@ -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 {
Expand All @@ -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(|_| ())
Expand All @@ -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(())
}
}
}
17 changes: 12 additions & 5 deletions aya/src/programs/sk_msg.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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 {
Expand Down
16 changes: 11 additions & 5 deletions aya/src/programs/sk_skb.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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>(())
/// ```
///
Expand All @@ -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,
Expand Down

0 comments on commit 898a14d

Please sign in to comment.