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
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