mirror of https://github.com/aya-rs/aya
Merge 1a8fdc1c0b into 0237e36dbe
Michal Rostecki
2 days ago
committed by
GitHub
commit
d0fe466787
@ -0,0 +1,150 @@
|
||||
//! Task storage.
|
||||
use std::{
|
||||
borrow::Borrow,
|
||||
marker::PhantomData,
|
||||
os::fd::{AsFd as _, AsRawFd as _},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
Pod,
|
||||
maps::{MapData, MapError, check_kv_size},
|
||||
sys::{PidFd, SyscallError, bpf_map_lookup_elem},
|
||||
};
|
||||
|
||||
/// Task storage is a type of map which uses `task_struct` kernel type as a
|
||||
/// key. When the task (process) stops, the corresponding entry is
|
||||
/// automatically removed.
|
||||
///
|
||||
/// # Minimum kernel version
|
||||
///
|
||||
/// The minimum kernel version required to use this feature is 5.12.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```no_run
|
||||
/// # let mut ebpf = aya::Ebpf::load(&[])?;
|
||||
/// use aya::maps::TaskStorage;
|
||||
///
|
||||
/// let mut task_storage: TaskStorage<_, u32> = TaskStorage::try_from(ebpf.map_mut("TASK_STORAGE").unwrap())?;
|
||||
///
|
||||
/// let pid = 0;
|
||||
/// let value = task_storage.get(&pid, 0)?;
|
||||
/// # Ok::<(), aya::EbpfError>(())
|
||||
/// ```
|
||||
#[doc(alias = "BPF_MAP_TYPE_TASK_STORAGE")]
|
||||
#[derive(Debug)]
|
||||
pub struct TaskStorage<T, V> {
|
||||
pub(crate) inner: T,
|
||||
_v: PhantomData<V>,
|
||||
}
|
||||
|
||||
impl<T: Borrow<MapData>, V: Pod> TaskStorage<T, V> {
|
||||
pub(crate) fn new(map: T) -> Result<Self, MapError> {
|
||||
let data = map.borrow();
|
||||
check_kv_size::<u32, V>(data)?;
|
||||
Ok(Self {
|
||||
inner: map,
|
||||
_v: PhantomData,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the value stored for the given `pid`.
|
||||
pub fn get(&self, pid: &u32, flags: u64) -> Result<V, MapError> {
|
||||
let pidfd = PidFd::open(*pid, 0).map_err(|(_, io_error)| SyscallError {
|
||||
call: "pidfd_open",
|
||||
io_error,
|
||||
})?;
|
||||
let map_fd = self.inner.borrow().fd().as_fd();
|
||||
let value = bpf_map_lookup_elem(map_fd, &pidfd.as_raw_fd(), flags).map_err(|io_error| {
|
||||
SyscallError {
|
||||
call: "bpf_map_lookup_elem",
|
||||
io_error,
|
||||
}
|
||||
})?;
|
||||
value.ok_or(MapError::KeyNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io;
|
||||
|
||||
use assert_matches::assert_matches;
|
||||
use aya_obj::generated::bpf_map_type::BPF_MAP_TYPE_TASK_STORAGE;
|
||||
use libc::EFAULT;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
maps::{
|
||||
Map,
|
||||
test_utils::{self, new_map},
|
||||
},
|
||||
sys::{SysResult, Syscall, override_syscall},
|
||||
};
|
||||
|
||||
fn new_obj_map() -> aya_obj::Map {
|
||||
test_utils::new_obj_map::<u32>(BPF_MAP_TYPE_TASK_STORAGE)
|
||||
}
|
||||
|
||||
fn sys_error(value: i32) -> SysResult {
|
||||
Err((-1, io::Error::from_raw_os_error(value)))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrong_value_size() {
|
||||
let map = new_map(new_obj_map());
|
||||
let map = Map::TaskStorage(map);
|
||||
assert_matches!(
|
||||
TaskStorage::<_, u16>::try_from(&map),
|
||||
Err(MapError::InvalidValueSize {
|
||||
size: 2,
|
||||
expected: 4
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_from_wrong_map() {
|
||||
let map = new_map(new_obj_map());
|
||||
let map = Map::Array(map);
|
||||
assert_matches!(
|
||||
TaskStorage::<_, u32>::try_from(&map),
|
||||
Err(MapError::InvalidMapType { .. })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_new_ok() {
|
||||
let map = new_map(new_obj_map());
|
||||
assert!(TaskStorage::<_, u32>::new(&map).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_try_from_ok() {
|
||||
let map = new_map(new_obj_map());
|
||||
let map = Map::TaskStorage(map);
|
||||
assert!(TaskStorage::<_, u32>::try_from(&map).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_pidfd_syscall_error() {
|
||||
let mut map = new_map(new_obj_map());
|
||||
let map = TaskStorage::<_, u32>::new(&mut map).unwrap();
|
||||
|
||||
override_syscall(|call| match call {
|
||||
Syscall::Ebpf { .. } => Ok(1),
|
||||
Syscall::PidfdOpen { .. } => sys_error(EFAULT),
|
||||
_ => sys_error(EFAULT),
|
||||
});
|
||||
|
||||
assert_matches!(
|
||||
map.get(&1, 0), Err(MapError::SyscallError(
|
||||
SyscallError {
|
||||
call: "pidfd_open",
|
||||
io_error
|
||||
}
|
||||
))
|
||||
if io_error.raw_os_error() == Some(EFAULT)
|
||||
);
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,38 @@
|
||||
// clang-format off
|
||||
#include <vmlinux.h>
|
||||
#include <bpf/bpf_helpers.h>
|
||||
#include <bpf/bpf_core_read.h>
|
||||
#include <bpf/bpf_tracing.h>
|
||||
// clang-format on
|
||||
|
||||
char _license[] SEC("license") = "GPL";
|
||||
|
||||
struct {
|
||||
__uint(type, BPF_MAP_TYPE_TASK_STORAGE);
|
||||
__uint(map_flags, BPF_F_NO_PREALLOC);
|
||||
__type(key, int);
|
||||
__type(value, __u32);
|
||||
} task_storage SEC(".maps");
|
||||
|
||||
void bpf_rcu_read_lock(void) __ksym;
|
||||
void bpf_rcu_read_unlock(void) __ksym;
|
||||
|
||||
SEC("tp_btf/sys_enter")
|
||||
int BPF_PROG(sys_enter, struct pt_regs *regs, long id) {
|
||||
__u32 value = 1;
|
||||
struct task_struct *task = bpf_get_current_task_btf();
|
||||
// This test is triggered by a Rust test, running in a thread. A current task
|
||||
// (the one returned by `bpf_get_current_task()`) represents that thread. If
|
||||
// we create a task storage entry for that task, our user-space test will not
|
||||
// be able to retrieve it as pidfd.
|
||||
// To make retrieval of the map element by pidfd possible, we need to use the
|
||||
// `group_leader` (a `struct task_struct*` instance representing the process)
|
||||
// as the key.
|
||||
bpf_rcu_read_lock();
|
||||
struct task_struct *group_leader = BPF_CORE_READ(task, group_leader);
|
||||
bpf_task_storage_get(&task_storage, group_leader, &value,
|
||||
BPF_LOCAL_STORAGE_GET_F_CREATE);
|
||||
bpf_rcu_read_unlock();
|
||||
|
||||
return 0;
|
||||
}
|
||||
@ -0,0 +1,25 @@
|
||||
use std::{thread::sleep, time::Duration};
|
||||
|
||||
use aya::{Btf, Ebpf, maps::TaskStorage, programs::BtfTracePoint};
|
||||
use test_log::test;
|
||||
|
||||
#[test]
|
||||
fn test_task_storage_get() {
|
||||
let mut ebpf = Ebpf::load(crate::TASK_STORAGE).unwrap();
|
||||
|
||||
let prog: &mut BtfTracePoint = ebpf.program_mut("sys_enter").unwrap().try_into().unwrap();
|
||||
let btf = Btf::from_sys_fs().unwrap();
|
||||
prog.load("sys_enter", &btf).unwrap();
|
||||
prog.attach().unwrap();
|
||||
|
||||
let task_storage: TaskStorage<_, u32> =
|
||||
TaskStorage::try_from(ebpf.map_mut("task_storage").unwrap()).unwrap();
|
||||
|
||||
// Trigger the eBPF program by issuing a syscall (`getpid`).
|
||||
let pid = unsafe { libc::getpid() } as u32;
|
||||
|
||||
sleep(Duration::from_millis(10));
|
||||
|
||||
let value = task_storage.get(&pid, 0).unwrap();
|
||||
assert_eq!(value, 1);
|
||||
}
|
||||
Loading…
Reference in New Issue