diff --git a/ebpf/aya-ebpf/src/btf_maps/mod.rs b/ebpf/aya-ebpf/src/btf_maps/mod.rs index b84e4a7f1..51813dd62 100644 --- a/ebpf/aya-ebpf/src/btf_maps/mod.rs +++ b/ebpf/aya-ebpf/src/btf_maps/mod.rs @@ -1,10 +1,12 @@ pub mod array; pub mod ring_buf; pub mod sk_storage; +pub mod task_storage; pub use array::Array; pub use ring_buf::RingBuf; pub use sk_storage::SkStorage; +pub use task_storage::TaskStorage; /// Defines a BTF-compatible map struct with flat `#[repr(C)]` layout. /// diff --git a/ebpf/aya-ebpf/src/btf_maps/task_storage.rs b/ebpf/aya-ebpf/src/btf_maps/task_storage.rs new file mode 100644 index 000000000..034548723 --- /dev/null +++ b/ebpf/aya-ebpf/src/btf_maps/task_storage.rs @@ -0,0 +1,74 @@ +use core::ptr; + +use aya_ebpf_bindings::bindings::{BPF_F_NO_PREALLOC, BPF_LOCAL_STORAGE_GET_F_CREATE, task_struct}; + +use crate::{ + btf_maps::btf_map_def, + helpers::{bpf_get_current_task_btf, bpf_task_storage_delete, bpf_task_storage_get}, +}; + +btf_map_def!( + /// A BTF-compatible BPF task storage map. + /// + /// Task storage maps require `BPF_F_NO_PREALLOC` flag and `max_entries: 0`. + pub struct TaskStorage, + map_type: BPF_MAP_TYPE_TASK_STORAGE, + max_entries: 0, + map_flags: BPF_F_NO_PREALLOC as usize, + key_type: i32, + value_type: T, +); + +impl TaskStorage { + #[inline(always)] + fn get_ptr(&self, task: Option<*mut task_struct>, value: *mut T, flags: u64) -> *mut T { + let task = task.unwrap_or_else(|| unsafe { bpf_get_current_task_btf() }); + unsafe { bpf_task_storage_get(self.as_ptr(), task.cast(), value.cast(), flags) }.cast() + } + + /// Gets a mutable reference to the value associated with `task`. + /// + /// If `task` is `None`, the current task is used. + /// + /// # Safety + /// + /// This function may dereference the pointer `task`. + #[inline(always)] + pub unsafe fn get_ptr_mut(&self, task: Option<*mut task_struct>) -> *mut T { + self.get_ptr(task, ptr::null_mut(), 0) + } + + /// Gets a mutable reference to the value associated with `task`. + /// + /// If no value is associated with `task`, `value` will be inserted. If + /// `task` is `None`, the current task is used. + /// + /// # Safety + /// + /// This function may dereference the pointer `task`. + #[inline(always)] + pub unsafe fn get_or_insert_ptr_mut( + &self, + task: Option<*mut task_struct>, + value: Option<&mut T>, + ) -> *mut T { + self.get_ptr( + task, + value.map_or(ptr::null_mut(), ptr::from_mut), + BPF_LOCAL_STORAGE_GET_F_CREATE.into(), + ) + } + + /// Deletes the value associated with `task`. If `task` is `None`, the + /// current task is used. + /// + /// # Safety + /// + /// This function may dereference the pointer `task`. + #[inline(always)] + pub unsafe fn delete(&self, task: Option<*mut task_struct>) -> Result<(), i32> { + let task = task.unwrap_or_else(|| unsafe { bpf_get_current_task_btf() }); + let ret = unsafe { bpf_task_storage_delete(self.as_ptr(), task.cast()) }; + if ret == 0 { Ok(()) } else { Err(ret as i32) } + } +} diff --git a/test/integration-ebpf/Cargo.toml b/test/integration-ebpf/Cargo.toml index 185c37c39..c220070af 100644 --- a/test/integration-ebpf/Cargo.toml +++ b/test/integration-ebpf/Cargo.toml @@ -92,6 +92,10 @@ path = "src/simple_prog.rs" name = "strncmp" path = "src/strncmp.rs" +[[bin]] +name = "task_storage" +path = "src/task_storage.rs" + [[bin]] name = "tcx" path = "src/tcx.rs" diff --git a/test/integration-ebpf/src/task_storage.rs b/test/integration-ebpf/src/task_storage.rs new file mode 100644 index 000000000..5d3000f15 --- /dev/null +++ b/test/integration-ebpf/src/task_storage.rs @@ -0,0 +1,59 @@ +#![no_std] +#![no_main] +#![expect(unused_crate_dependencies, reason = "used in other bins")] + +use aya_ebpf::{ + btf_maps::TaskStorage, + helpers::bpf_get_current_task_btf, + macros::{btf_map, kprobe}, + programs::ProbeContext, +}; + +#[cfg(not(test))] +extern crate ebpf_panic; + +#[btf_map] +static TASK_STORAGE: TaskStorage = TaskStorage::new(); + +#[kprobe] +fn task_storage_test(_ctx: ProbeContext) -> i64 { + let task = unsafe { bpf_get_current_task_btf() }; + + let mut initial_val: u64 = 42; + let ptr = unsafe { TASK_STORAGE.get_or_insert_ptr_mut(Some(task), Some(&mut initial_val)) }; + if ptr.is_null() { + return 0; + } + + let val = unsafe { *ptr }; + if val != 42 { + return 0; + } + + unsafe { + *ptr = 1337; + } + let ptr = unsafe { TASK_STORAGE.get_ptr_mut(Some(task)) }; + if ptr.is_null() { + return 0; + } + + let val = unsafe { *ptr }; + if val != 1337 { + return 0; + } + + let ret = unsafe { + TASK_STORAGE.delete(None /* uses current task */) + }; + if ret.is_err() { + return 0; + } + + let ptr = unsafe { TASK_STORAGE.get_ptr_mut(Some(task)) }; + if !ptr.is_null() { + return 0; + } + + 0 +} diff --git a/test/integration-test/src/lib.rs b/test/integration-test/src/lib.rs index 702c71908..75015fd4a 100644 --- a/test/integration-test/src/lib.rs +++ b/test/integration-test/src/lib.rs @@ -57,6 +57,9 @@ bpf_file!( SIMPLE_PROG => "simple_prog", SK_STORAGE => "sk_storage", STRNCMP => "strncmp", + // TODO(sblaisdo) - add integration test after userspace support is completed + // https://github.com/aya-rs/aya/pull/1161 + // TASK_STORAGE => "task_storage", TCX => "tcx", TEST => "test", TWO_PROGS => "two_progs", diff --git a/xtask/public-api/aya-ebpf.txt b/xtask/public-api/aya-ebpf.txt index a5d1de27b..20115588b 100644 --- a/xtask/public-api/aya-ebpf.txt +++ b/xtask/public-api/aya-ebpf.txt @@ -107,6 +107,39 @@ impl core::borrow::BorrowMut for aya_ebpf::btf_maps::sk_storage::SkStorage pub fn aya_ebpf::btf_maps::sk_storage::SkStorage::borrow_mut(&mut self) -> &mut T impl core::convert::From for aya_ebpf::btf_maps::sk_storage::SkStorage pub fn aya_ebpf::btf_maps::sk_storage::SkStorage::from(t: T) -> T +pub mod aya_ebpf::btf_maps::task_storage +#[repr(C)] pub struct aya_ebpf::btf_maps::task_storage::TaskStorage +impl aya_ebpf::btf_maps::task_storage::TaskStorage +pub unsafe fn aya_ebpf::btf_maps::task_storage::TaskStorage::delete(&self, task: core::option::Option<*mut aya_ebpf_bindings::x86_64::bindings::task_struct>) -> core::result::Result<(), i32> +pub unsafe fn aya_ebpf::btf_maps::task_storage::TaskStorage::get_or_insert_ptr_mut(&self, task: core::option::Option<*mut aya_ebpf_bindings::x86_64::bindings::task_struct>, value: core::option::Option<&mut T>) -> *mut T +pub unsafe fn aya_ebpf::btf_maps::task_storage::TaskStorage::get_ptr_mut(&self, task: core::option::Option<*mut aya_ebpf_bindings::x86_64::bindings::task_struct>) -> *mut T +impl aya_ebpf::btf_maps::task_storage::TaskStorage +pub const fn aya_ebpf::btf_maps::task_storage::TaskStorage::new() -> Self +impl core::default::Default for aya_ebpf::btf_maps::task_storage::TaskStorage +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::default() -> Self +impl core::marker::Sync for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::marker::Freeze for aya_ebpf::btf_maps::task_storage::TaskStorage +impl !core::marker::Send for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::marker::Unpin for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::marker::UnsafeUnpin for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::panic::unwind_safe::RefUnwindSafe for aya_ebpf::btf_maps::task_storage::TaskStorage where T: core::panic::unwind_safe::RefUnwindSafe +impl core::panic::unwind_safe::UnwindSafe for aya_ebpf::btf_maps::task_storage::TaskStorage where T: core::panic::unwind_safe::RefUnwindSafe +impl core::convert::Into for aya_ebpf::btf_maps::task_storage::TaskStorage where U: core::convert::From +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::into(self) -> U +impl core::convert::TryFrom for aya_ebpf::btf_maps::task_storage::TaskStorage where U: core::convert::Into +pub type aya_ebpf::btf_maps::task_storage::TaskStorage::Error = core::convert::Infallible +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::try_from(value: U) -> core::result::Result>::Error> +impl core::convert::TryInto for aya_ebpf::btf_maps::task_storage::TaskStorage where U: core::convert::TryFrom +pub type aya_ebpf::btf_maps::task_storage::TaskStorage::Error = >::Error +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::try_into(self) -> core::result::Result>::Error> +impl core::any::Any for aya_ebpf::btf_maps::task_storage::TaskStorage where T: 'static + ?core::marker::Sized +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::type_id(&self) -> core::any::TypeId +impl core::borrow::Borrow for aya_ebpf::btf_maps::task_storage::TaskStorage where T: ?core::marker::Sized +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::borrow(&self) -> &T +impl core::borrow::BorrowMut for aya_ebpf::btf_maps::task_storage::TaskStorage where T: ?core::marker::Sized +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::borrow_mut(&mut self) -> &mut T +impl core::convert::From for aya_ebpf::btf_maps::task_storage::TaskStorage +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::from(t: T) -> T #[repr(C)] pub struct aya_ebpf::btf_maps::Array impl aya_ebpf::btf_maps::array::Array pub fn aya_ebpf::btf_maps::array::Array::get(&self, index: u32) -> core::option::Option<&T> @@ -206,6 +239,38 @@ impl core::borrow::BorrowMut for aya_ebpf::btf_maps::sk_storage::SkStorage pub fn aya_ebpf::btf_maps::sk_storage::SkStorage::borrow_mut(&mut self) -> &mut T impl core::convert::From for aya_ebpf::btf_maps::sk_storage::SkStorage pub fn aya_ebpf::btf_maps::sk_storage::SkStorage::from(t: T) -> T +#[repr(C)] pub struct aya_ebpf::btf_maps::TaskStorage +impl aya_ebpf::btf_maps::task_storage::TaskStorage +pub unsafe fn aya_ebpf::btf_maps::task_storage::TaskStorage::delete(&self, task: core::option::Option<*mut aya_ebpf_bindings::x86_64::bindings::task_struct>) -> core::result::Result<(), i32> +pub unsafe fn aya_ebpf::btf_maps::task_storage::TaskStorage::get_or_insert_ptr_mut(&self, task: core::option::Option<*mut aya_ebpf_bindings::x86_64::bindings::task_struct>, value: core::option::Option<&mut T>) -> *mut T +pub unsafe fn aya_ebpf::btf_maps::task_storage::TaskStorage::get_ptr_mut(&self, task: core::option::Option<*mut aya_ebpf_bindings::x86_64::bindings::task_struct>) -> *mut T +impl aya_ebpf::btf_maps::task_storage::TaskStorage +pub const fn aya_ebpf::btf_maps::task_storage::TaskStorage::new() -> Self +impl core::default::Default for aya_ebpf::btf_maps::task_storage::TaskStorage +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::default() -> Self +impl core::marker::Sync for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::marker::Freeze for aya_ebpf::btf_maps::task_storage::TaskStorage +impl !core::marker::Send for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::marker::Unpin for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::marker::UnsafeUnpin for aya_ebpf::btf_maps::task_storage::TaskStorage +impl core::panic::unwind_safe::RefUnwindSafe for aya_ebpf::btf_maps::task_storage::TaskStorage where T: core::panic::unwind_safe::RefUnwindSafe +impl core::panic::unwind_safe::UnwindSafe for aya_ebpf::btf_maps::task_storage::TaskStorage where T: core::panic::unwind_safe::RefUnwindSafe +impl core::convert::Into for aya_ebpf::btf_maps::task_storage::TaskStorage where U: core::convert::From +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::into(self) -> U +impl core::convert::TryFrom for aya_ebpf::btf_maps::task_storage::TaskStorage where U: core::convert::Into +pub type aya_ebpf::btf_maps::task_storage::TaskStorage::Error = core::convert::Infallible +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::try_from(value: U) -> core::result::Result>::Error> +impl core::convert::TryInto for aya_ebpf::btf_maps::task_storage::TaskStorage where U: core::convert::TryFrom +pub type aya_ebpf::btf_maps::task_storage::TaskStorage::Error = >::Error +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::try_into(self) -> core::result::Result>::Error> +impl core::any::Any for aya_ebpf::btf_maps::task_storage::TaskStorage where T: 'static + ?core::marker::Sized +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::type_id(&self) -> core::any::TypeId +impl core::borrow::Borrow for aya_ebpf::btf_maps::task_storage::TaskStorage where T: ?core::marker::Sized +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::borrow(&self) -> &T +impl core::borrow::BorrowMut for aya_ebpf::btf_maps::task_storage::TaskStorage where T: ?core::marker::Sized +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::borrow_mut(&mut self) -> &mut T +impl core::convert::From for aya_ebpf::btf_maps::task_storage::TaskStorage +pub fn aya_ebpf::btf_maps::task_storage::TaskStorage::from(t: T) -> T pub mod aya_ebpf::helpers pub use aya_ebpf::helpers::generated pub macro aya_ebpf::helpers::bpf_printk!