mirror of https://github.com/aya-rs/aya
aya,ebpf: add BPF_MAP_TYPE_SK_STORAGE
This map type requires BTF, and we can finally do it!reviewable/pr1357/r18
Tamir Duberstein
2 weeks ago
parent
6babf17969
commit
de42b80c74
@ -0,0 +1,65 @@
|
||||
//! A socket local storage map backed by `BPF_MAP_TYPE_SK_STORAGE`.
|
||||
|
||||
use std::{
|
||||
borrow::{Borrow, BorrowMut},
|
||||
marker::PhantomData,
|
||||
os::fd::AsRawFd,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
Pod,
|
||||
maps::{MapData, MapError, check_kv_size, hash_map},
|
||||
};
|
||||
|
||||
/// A socket local storage map backed by `BPF_MAP_TYPE_SK_STORAGE`.
|
||||
///
|
||||
/// This map type stores values that are owned by individual sockets. The map keys are socket file
|
||||
/// descriptors and the values can be accessed both from eBPF using [`bpf_sk_storage_get`] and from
|
||||
/// user space through the methods on this type.
|
||||
///
|
||||
/// [`bpf_sk_storage_get`]: https://elixir.bootlin.com/linux/v6.12/source/include/uapi/linux/bpf.h#L4064-L4093
|
||||
#[doc(alias = "BPF_MAP_TYPE_SK_STORAGE")]
|
||||
#[derive(Debug)]
|
||||
pub struct SkStorage<T, V: Pod> {
|
||||
pub(crate) inner: T,
|
||||
_v: PhantomData<V>,
|
||||
}
|
||||
|
||||
impl<T: Borrow<MapData>, V: Pod> SkStorage<T, V> {
|
||||
pub(crate) fn new(map: T) -> Result<Self, MapError> {
|
||||
let data = map.borrow();
|
||||
check_kv_size::<i32, V>(data)?;
|
||||
|
||||
Ok(Self {
|
||||
inner: map,
|
||||
_v: PhantomData,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the value associated with `socket`.
|
||||
pub fn get(&self, socket: &impl AsRawFd, flags: u64) -> Result<V, MapError> {
|
||||
hash_map::get(self.inner.borrow(), &socket.as_raw_fd(), flags)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: BorrowMut<MapData>, V: Pod> SkStorage<T, V> {
|
||||
/// Creates or updates the value associated with `socket`.
|
||||
pub fn insert(
|
||||
&mut self,
|
||||
socket: &impl AsRawFd,
|
||||
value: impl Borrow<V>,
|
||||
flags: u64,
|
||||
) -> Result<(), MapError> {
|
||||
hash_map::insert(
|
||||
self.inner.borrow_mut(),
|
||||
&socket.as_raw_fd(),
|
||||
value.borrow(),
|
||||
flags,
|
||||
)
|
||||
}
|
||||
|
||||
/// Removes the storage associated with `socket`.
|
||||
pub fn remove(&mut self, socket: &impl AsRawFd) -> Result<(), MapError> {
|
||||
hash_map::remove(self.inner.borrow_mut(), &socket.as_raw_fd())
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,87 @@
|
||||
use core::{cell::UnsafeCell, ptr};
|
||||
|
||||
use aya_ebpf_bindings::bindings::{
|
||||
BPF_F_NO_PREALLOC, BPF_SK_STORAGE_GET_F_CREATE, bpf_map_type::BPF_MAP_TYPE_SK_STORAGE, bpf_sock,
|
||||
};
|
||||
use aya_ebpf_cty::{c_long, c_void};
|
||||
|
||||
use crate::{
|
||||
btf_map_def,
|
||||
helpers::generated::{bpf_sk_storage_delete, bpf_sk_storage_get},
|
||||
programs::sock_addr::SockAddrContext,
|
||||
};
|
||||
|
||||
btf_map_def!(SkStorageDef, BPF_MAP_TYPE_SK_STORAGE);
|
||||
|
||||
// TODO(https://github.com/rust-lang/rust/issues/76560): this should be:
|
||||
//
|
||||
// { F | BPF_F_NO_PREALLOC as usize }.
|
||||
#[repr(transparent)]
|
||||
pub struct SkStorage<T>(UnsafeCell<SkStorageDef<i32, T, 0, { BPF_F_NO_PREALLOC as usize }>>);
|
||||
|
||||
unsafe impl<T: Sync> Sync for SkStorage<T> {}
|
||||
|
||||
impl<T> SkStorage<T> {
|
||||
#[expect(
|
||||
clippy::new_without_default,
|
||||
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||
)]
|
||||
pub const fn new() -> Self {
|
||||
Self(UnsafeCell::new(SkStorageDef::new()))
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn as_ptr(&self) -> *mut c_void {
|
||||
let Self(inner) = self;
|
||||
|
||||
inner.get().cast()
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
fn get_ptr(&self, ctx: &SockAddrContext, value: *mut T, flags: u64) -> *mut T {
|
||||
let sock_addr = unsafe { &*ctx.sock_addr };
|
||||
let sk = unsafe { sock_addr.__bindgen_anon_1.sk };
|
||||
unsafe { bpf_sk_storage_get(self.as_ptr(), sk.cast(), value.cast(), flags) }.cast::<T>()
|
||||
}
|
||||
|
||||
/// Gets a mutable reference to the value associated with `sk`.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// This function may dereference the pointer `sk`.
|
||||
#[inline(always)]
|
||||
pub unsafe fn get_ptr_mut(&self, ctx: &SockAddrContext) -> *mut T {
|
||||
self.get_ptr(ctx, ptr::null_mut(), 0)
|
||||
}
|
||||
|
||||
/// Gets a mutable reference to the value associated with `sk`.
|
||||
///
|
||||
/// If no value is associated with `sk`, `value` will be inserted.`
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// This function may dereference the pointer `sk`.
|
||||
#[inline(always)]
|
||||
pub unsafe fn get_or_insert_ptr_mut(
|
||||
&self,
|
||||
ctx: &SockAddrContext,
|
||||
value: Option<&mut T>,
|
||||
) -> *mut T {
|
||||
self.get_ptr(
|
||||
ctx,
|
||||
value.map_or(ptr::null_mut(), ptr::from_mut),
|
||||
BPF_SK_STORAGE_GET_F_CREATE.into(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Deletes the value associated with `sk`.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// This function may dereference the pointer `sk`.
|
||||
#[inline(always)]
|
||||
pub unsafe fn delete(&self, sk: *mut bpf_sock) -> Result<(), c_long> {
|
||||
let ret = unsafe { bpf_sk_storage_delete(self.as_ptr(), sk.cast()) };
|
||||
if ret == 0 { Ok(()) } else { Err(ret) }
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,78 @@
|
||||
#![no_std]
|
||||
#![no_main]
|
||||
#![expect(unused_crate_dependencies, reason = "used in other bins")]
|
||||
|
||||
use aya_ebpf::{
|
||||
bindings::sk_action,
|
||||
btf_maps::SkStorage,
|
||||
macros::{btf_map, cgroup_sock_addr},
|
||||
programs::SockAddrContext,
|
||||
};
|
||||
use integration_common::sk_storage::{Ip, Value};
|
||||
#[cfg(not(test))]
|
||||
extern crate ebpf_panic;
|
||||
|
||||
#[btf_map]
|
||||
static SOCKET_STORAGE: SkStorage<Value> = SkStorage::new();
|
||||
|
||||
#[cgroup_sock_addr(connect4)]
|
||||
pub(crate) fn sk_storage_connect4(ctx: SockAddrContext) -> i32 {
|
||||
let sock_addr = unsafe { &*ctx.sock_addr };
|
||||
|
||||
let storage = unsafe { SOCKET_STORAGE.get_or_insert_ptr_mut(&ctx, None) };
|
||||
if !storage.is_null() {
|
||||
unsafe {
|
||||
*storage = Value {
|
||||
user_family: sock_addr.user_family,
|
||||
user_ip: Ip::V4(sock_addr.user_ip4),
|
||||
user_port: sock_addr.user_port,
|
||||
family: sock_addr.family,
|
||||
type_: sock_addr.type_,
|
||||
protocol: sock_addr.protocol,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sk_action::SK_PASS as _
|
||||
}
|
||||
|
||||
#[cgroup_sock_addr(connect6)]
|
||||
pub(crate) fn sk_storage_connect6(ctx: SockAddrContext) -> i32 {
|
||||
let sock_addr = unsafe { &*ctx.sock_addr };
|
||||
|
||||
let storage = unsafe { SOCKET_STORAGE.get_or_insert_ptr_mut(&ctx, None) };
|
||||
if !storage.is_null() {
|
||||
let mut user_ip6 = [0u32; 4];
|
||||
unsafe {
|
||||
// Verifier dislikes the naive thing:
|
||||
//
|
||||
// ; let sk = unsafe { sock_addr.__bindgen_anon_1.sk };
|
||||
// 0: (79) r2 = *(u64 *)(r1 +64) ; R1=ctx(off=0,imm=0) R2_w=sock(off=0,imm=0)
|
||||
// ; user_family: sock_addr.user_family,
|
||||
// 1: (61) r3 = *(u32 *)(r1 +0) ; R1=ctx(off=0,imm=0) R3_w=scalar(umax=4294967295,var_off=(0x0; 0xffffffff))
|
||||
// ; user_ip: Ip::V6(sock_addr.user_ip6),
|
||||
// 2: (bf) r4 = r1 ; R1=ctx(off=0,imm=0) R4_w=ctx(off=0,imm=0)
|
||||
// 3: (07) r4 += 8 ; R4_w=ctx(off=8,imm=0)
|
||||
// ; let mut value = Value {
|
||||
// 4: (bf) r5 = r10 ; R5_w=fp0 R10=fp0
|
||||
// 5: (07) r5 += -32 ; R5_w=fp-32
|
||||
// ; user_ip: Ip::V6(sock_addr.user_ip6),
|
||||
// 6: (61) r0 = *(u32 *)(r4 +0)
|
||||
// dereference of modified ctx ptr R4 off=8 disallowed
|
||||
user_ip6[0] = sock_addr.user_ip6[0];
|
||||
user_ip6[1] = sock_addr.user_ip6[1];
|
||||
user_ip6[2] = sock_addr.user_ip6[2];
|
||||
user_ip6[3] = sock_addr.user_ip6[3];
|
||||
*storage = Value {
|
||||
user_family: sock_addr.user_family,
|
||||
user_ip: Ip::V6(user_ip6),
|
||||
user_port: sock_addr.user_port,
|
||||
family: sock_addr.family,
|
||||
type_: sock_addr.type_,
|
||||
protocol: sock_addr.protocol,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
sk_action::SK_PASS as _
|
||||
}
|
||||
@ -0,0 +1,98 @@
|
||||
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener, TcpStream};
|
||||
|
||||
use assert_matches::assert_matches;
|
||||
use aya::{
|
||||
EbpfLoader,
|
||||
maps::{MapError, SkStorage},
|
||||
programs::{CgroupAttachMode, CgroupSockAddr},
|
||||
};
|
||||
use integration_common::sk_storage::{Ip, Value};
|
||||
use libc::{self};
|
||||
use test_log::test;
|
||||
|
||||
use crate::utils::{Cgroup, NetNsGuard};
|
||||
|
||||
#[test]
|
||||
fn sk_storage_connect() {
|
||||
let mut ebpf = EbpfLoader::new().load(crate::SK_STORAGE).unwrap();
|
||||
|
||||
let storage = ebpf.take_map("SOCKET_STORAGE").unwrap();
|
||||
let mut storage = SkStorage::<_, Value>::try_from(storage).unwrap();
|
||||
|
||||
let _netns = NetNsGuard::new();
|
||||
let root_cgroup = Cgroup::root();
|
||||
let cgroup = root_cgroup.create_child("aya-test-sk-storage");
|
||||
let cgroup_fd = cgroup.fd();
|
||||
|
||||
let guards = ebpf
|
||||
.programs_mut()
|
||||
.map(|(name, prog)| {
|
||||
let prog: &mut CgroupSockAddr = prog.try_into().expect(name);
|
||||
prog.load().expect(name);
|
||||
let link_id = prog
|
||||
.attach(cgroup_fd, CgroupAttachMode::Single)
|
||||
.expect(name);
|
||||
scopeguard::guard((), |()| {
|
||||
prog.detach(link_id).expect(name);
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let cgroup = cgroup.into_cgroup();
|
||||
cgroup.write_pid(std::process::id());
|
||||
|
||||
let listener4 = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
|
||||
let addr4 = listener4.local_addr().unwrap();
|
||||
let listener6 = TcpListener::bind((Ipv6Addr::LOCALHOST, 0)).unwrap();
|
||||
let addr6 = listener6.local_addr().unwrap();
|
||||
|
||||
{
|
||||
let client4 = TcpStream::connect(addr4).unwrap();
|
||||
assert_matches!(storage.get(&client4, 0), Ok(value4) => {
|
||||
assert_eq!(value4, expected_value(&addr4));
|
||||
});
|
||||
storage.remove(&client4).unwrap();
|
||||
assert_matches!(storage.get(&client4, 0), Err(MapError::KeyNotFound));
|
||||
|
||||
let client6 = TcpStream::connect(addr6).unwrap();
|
||||
assert_matches!(storage.get(&client6, 0), Ok(value6) => {
|
||||
assert_eq!(value6, expected_value(&addr6));
|
||||
});
|
||||
storage.remove(&client6).unwrap();
|
||||
assert_matches!(storage.get(&client6, 0), Err(MapError::KeyNotFound));
|
||||
}
|
||||
|
||||
// Detach.
|
||||
drop(guards);
|
||||
|
||||
{
|
||||
let client4 = TcpStream::connect(addr4).unwrap();
|
||||
assert_matches!(storage.get(&client4, 0), Err(MapError::KeyNotFound));
|
||||
|
||||
let client6 = TcpStream::connect(addr6).unwrap();
|
||||
assert_matches!(storage.get(&client6, 0), Err(MapError::KeyNotFound));
|
||||
}
|
||||
}
|
||||
|
||||
fn expected_value(addr: &SocketAddr) -> Value {
|
||||
match addr {
|
||||
SocketAddr::V4(addr) => Value {
|
||||
user_family: libc::AF_INET as u32,
|
||||
user_ip: Ip::V4(u32::from_ne_bytes(addr.ip().octets())),
|
||||
user_port: u32::from(addr.port().to_be()),
|
||||
family: libc::AF_INET as u32,
|
||||
type_: libc::SOCK_STREAM as u32,
|
||||
protocol: libc::IPPROTO_TCP as u32,
|
||||
},
|
||||
SocketAddr::V6(addr) => Value {
|
||||
user_family: libc::AF_INET6 as u32,
|
||||
user_ip: Ip::V6(unsafe {
|
||||
core::mem::transmute::<[u8; 16], [u32; 4]>(addr.ip().octets())
|
||||
}),
|
||||
user_port: u32::from(addr.port().to_be()),
|
||||
family: libc::AF_INET6 as u32,
|
||||
type_: libc::SOCK_STREAM as u32,
|
||||
protocol: libc::IPPROTO_TCP as u32,
|
||||
},
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue