bpf: Handle raw tracepoint arguments

Provide an `arg()` method in `RawTracepointArgs` wrapper of
`bpf_raw_tracepoint_args` and also in `RawTracepointContext`, so
it's directly available in raw tracepoint programs.

The methods and traits implemented here are unsafe. There is no
way to reliably check the number of available arguments, so
requesting a non-existing one leads to undefined behavior.

bpf/aya-bpf/src/args.rs

@@ -5,7 +5,7 @@ use crate::bindings::pt_regs;
 use crate::bindings::user_pt_regs as pt_regs;
 #[cfg(bpf_target_arch = "riscv64")]
 use crate::bindings::user_regs_struct as pt_regs;
-use crate::{cty::c_void, helpers::bpf_probe_read};
+use crate::{bindings::bpf_raw_tracepoint_args, cty::c_void, helpers::bpf_probe_read};
 
 /// A trait that indicates a valid type for an argument which can be coerced from a BTF
 /// context.
@@ -323,3 +323,70 @@ impl_from_pt_regs!(i32);
 impl_from_pt_regs!(i64);
 impl_from_pt_regs!(usize);
 impl_from_pt_regs!(isize);
+
+/// A Rust wrapper on `bpf_raw_tracepoint_args`.
+pub struct RawTracepointArgs {
+    args: *mut bpf_raw_tracepoint_args,
+}
+
+impl RawTracepointArgs {
+    pub fn new(args: *mut bpf_raw_tracepoint_args) -> Self {
+        RawTracepointArgs { args }
+    }
+
+    /// Returns the n-th argument of the raw tracepoint.
+    ///
+    /// ## Safety
+    ///
+    /// This method is unsafe because it performs raw pointer conversion and makes assumptions
+    /// about the structure of the `bpf_raw_tracepoint_args` type. The tracepoint arguments are
+    /// represented as an array of `__u64` values (i.e., `__u64 args[0]`), and this method
+    /// provides a way to access these arguments conveniently in Rust using `as_slice`.
+    ///
+    /// However, the method does not check the total number of available arguments for a given
+    /// tracepoint and assumes that the slice has at least `n` elements, leading to undefined
+    /// behavior if this condition is not met.
+    pub unsafe fn arg<T: FromRawTracepointArgs>(&self, n: usize) -> *const T {
+        &T::from_argument(&*self.args, n)
+    }
+}
+
+pub unsafe trait FromRawTracepointArgs: Sized {
+    unsafe fn from_argument(ctx: &bpf_raw_tracepoint_args, n: usize) -> Self;
+}
+
+unsafe impl<T> FromRawTracepointArgs for *const T {
+    unsafe fn from_argument(ctx: &bpf_raw_tracepoint_args, n: usize) -> *const T {
+        // Raw tracepoint arguments are exposed as `__u64 args[0]`.
+        // https://elixir.bootlin.com/linux/v6.5.5/source/include/uapi/linux/bpf.h#L6829
+        //
+        // The most convenient way of accessing such type in Rust is to use
+        // `as_slice`.
+        //
+        // We don't know how many arguments are there for the given tracepoint,
+        // so we just assume that the slice has at least n elements. The whole
+        // assumntion and implementation is unsafe.
+        ctx.args.as_slice(n + 1)[n] as *const _
+    }
+}
+
+macro_rules! unsafe_impl_from_raw_tracepoint_args {
+    ($type:ident) => {
+        unsafe impl FromRawTracepointArgs for $type {
+            unsafe fn from_argument(ctx: &bpf_raw_tracepoint_args, n: usize) -> Self {
+                ctx.args.as_slice(n + 1)[n] as _
+            }
+        }
+    };
+}
+
+unsafe_impl_from_raw_tracepoint_args!(u8);
+unsafe_impl_from_raw_tracepoint_args!(u16);
+unsafe_impl_from_raw_tracepoint_args!(u32);
+unsafe_impl_from_raw_tracepoint_args!(u64);
+unsafe_impl_from_raw_tracepoint_args!(i8);
+unsafe_impl_from_raw_tracepoint_args!(i16);
+unsafe_impl_from_raw_tracepoint_args!(i32);
+unsafe_impl_from_raw_tracepoint_args!(i64);
+unsafe_impl_from_raw_tracepoint_args!(usize);
+unsafe_impl_from_raw_tracepoint_args!(isize);

bpf/aya-bpf/src/lib.rs

@@ -17,7 +17,7 @@
 pub use aya_bpf_bindings::bindings;
 
 mod args;
-pub use args::PtRegs;
+pub use args::{PtRegs, RawTracepointArgs};
 pub mod helpers;
 pub mod maps;
 pub mod programs;

bpf/aya-bpf/src/programs/raw_tracepoint.rs

@@ -1,19 +1,25 @@
 use core::ffi::c_void;
 
-use crate::BpfContext;
+use crate::{args::FromRawTracepointArgs, bindings::bpf_raw_tracepoint_args, BpfContext};
 
 pub struct RawTracePointContext {
-    ctx: *mut c_void,
+    ctx: *mut bpf_raw_tracepoint_args,
 }
 
 impl RawTracePointContext {
     pub fn new(ctx: *mut c_void) -> RawTracePointContext {
-        RawTracePointContext { ctx }
+        RawTracePointContext {
+            ctx: ctx as *mut bpf_raw_tracepoint_args,
+        }
+    }
+
+    pub unsafe fn arg<T: FromRawTracepointArgs>(&self, n: usize) -> T {
+        T::from_argument(&*self.ctx, n)
     }
 }
 
 impl BpfContext for RawTracePointContext {
     fn as_ptr(&self) -> *mut c_void {
-        self.ctx
+        self.ctx as *mut c_void
     }
 }

test/integration-ebpf/Cargo.toml

@@ -51,3 +51,7 @@ path = "src/redirect.rs"
 [[bin]]
 name = "xdp_sec"
 path = "src/xdp_sec.rs"
+
+[[bin]]
+name = "raw_tracepoint"
+path = "src/raw_tracepoint.rs"

test/integration-ebpf/src/raw_tracepoint.rs

@@ -0,0 +1,50 @@
+#![no_std]
+#![no_main]
+
+use aya_bpf::{
+    macros::{map, raw_tracepoint},
+    maps::Array,
+    programs::RawTracePointContext,
+};
+
+#[map]
+static RESULT: Array<SysEnterEvent> = Array::with_max_entries(1, 0);
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct SysEnterEvent {
+    pub common_type: u16,
+    pub common_flags: u8,
+    _padding: u8,
+}
+
+impl SysEnterEvent {
+    pub fn new(common_type: u16, common_flags: u8) -> Self {
+        Self {
+            common_type,
+            common_flags,
+            _padding: 0,
+        }
+    }
+}
+
+#[raw_tracepoint(tracepoint = "sys_enter")]
+pub fn sys_enter(ctx: RawTracePointContext) -> i32 {
+    let common_type: u16 = unsafe { ctx.arg(0) };
+    let common_flags: u8 = unsafe { ctx.arg(1) };
+
+    if let Some(ptr) = RESULT.get_ptr_mut(0) {
+        unsafe {
+            (*ptr).common_type = common_type;
+            (*ptr).common_flags = common_flags;
+        }
+    }
+
+    0
+}
+
+#[cfg(not(test))]
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    loop {}
+}

test/integration-test/src/lib.rs

@@ -21,6 +21,8 @@ pub const BPF_PROBE_READ: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/bpf_probe_read"));
 pub const REDIRECT: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/redirect"));
 pub const XDP_SEC: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/xdp_sec"));
+pub const RAW_TRACEPOINT: &[u8] =
+    include_bytes_aligned!(concat!(env!("OUT_DIR"), "/raw_tracepoint"));
 
 #[cfg(test)]
 mod tests;

test/integration-test/src/tests.rs

@@ -3,6 +3,7 @@ mod btf_relocations;
 mod elf;
 mod load;
 mod log;
+mod raw_tracepoint;
 mod rbpf;
 mod relocations;
 mod smoke;

test/integration-test/src/tests/raw_tracepoint.rs

@@ -0,0 +1,25 @@
+use aya::{maps::Array, programs::RawTracePoint, Bpf};
+
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub struct SysEnterEvent {
+    pub common_type: u16,
+    pub common_flags: u8,
+    _padding: u8,
+}
+
+unsafe impl aya::Pod for SysEnterEvent {}
+
+#[test]
+fn raw_tracepoint() {
+    let mut bpf = Bpf::load(crate::RAW_TRACEPOINT).unwrap();
+    let prog: &mut RawTracePoint = bpf.program_mut("sys_enter").unwrap().try_into().unwrap();
+    prog.load().unwrap();
+    prog.attach("sys_enter").unwrap();
+
+    let map: Array<_, SysEnterEvent> = Array::try_from(bpf.map_mut("RESULT").unwrap()).unwrap();
+    let result = map.get(&0, 0).unwrap();
+
+    assert_ne!(result.common_type, 0);
+    assert_ne!(result.common_flags, 0);
+}

xtask/public-api/aya-bpf.txt

@@ -1402,6 +1402,7 @@ pub fn aya_bpf::programs::probe::ProbeContext::from(t: T) -> T
 pub mod aya_bpf::programs::raw_tracepoint
 pub struct aya_bpf::programs::raw_tracepoint::RawTracePointContext
 impl aya_bpf::programs::raw_tracepoint::RawTracePointContext
+pub unsafe fn aya_bpf::programs::raw_tracepoint::RawTracePointContext::arg<T: FromRawTracepointArgs>(&self, n: usize) -> T
 pub fn aya_bpf::programs::raw_tracepoint::RawTracePointContext::new(ctx: *mut core::ffi::c_void) -> aya_bpf::programs::raw_tracepoint::RawTracePointContext
 impl aya_bpf::BpfContext for aya_bpf::programs::raw_tracepoint::RawTracePointContext
 pub fn aya_bpf::programs::raw_tracepoint::RawTracePointContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2053,6 +2054,7 @@ impl<T> core::convert::From<T> for aya_bpf::programs::probe::ProbeContext
 pub fn aya_bpf::programs::probe::ProbeContext::from(t: T) -> T
 pub struct aya_bpf::programs::RawTracePointContext
 impl aya_bpf::programs::raw_tracepoint::RawTracePointContext
+pub unsafe fn aya_bpf::programs::raw_tracepoint::RawTracePointContext::arg<T: FromRawTracepointArgs>(&self, n: usize) -> T
 pub fn aya_bpf::programs::raw_tracepoint::RawTracePointContext::new(ctx: *mut core::ffi::c_void) -> aya_bpf::programs::raw_tracepoint::RawTracePointContext
 impl aya_bpf::BpfContext for aya_bpf::programs::raw_tracepoint::RawTracePointContext
 pub fn aya_bpf::programs::raw_tracepoint::RawTracePointContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2453,6 +2455,31 @@ impl<T> core::borrow::BorrowMut<T> for aya_bpf::PtRegs where T: core::marker::Si
 pub fn aya_bpf::PtRegs::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya_bpf::PtRegs
 pub fn aya_bpf::PtRegs::from(t: T) -> T
+pub struct aya_bpf::RawTracepointArgs
+impl aya_bpf::RawTracepointArgs
+pub unsafe fn aya_bpf::RawTracepointArgs::arg<T: FromRawTracepointArgs>(&self, n: usize) -> *const T
+pub fn aya_bpf::RawTracepointArgs::new(args: *mut aya_bpf_bindings::x86_64::bindings::bpf_raw_tracepoint_args) -> Self
+impl !core::marker::Send for aya_bpf::RawTracepointArgs
+impl !core::marker::Sync for aya_bpf::RawTracepointArgs
+impl core::marker::Unpin for aya_bpf::RawTracepointArgs
+impl core::panic::unwind_safe::RefUnwindSafe for aya_bpf::RawTracepointArgs
+impl core::panic::unwind_safe::UnwindSafe for aya_bpf::RawTracepointArgs
+impl<T, U> core::convert::Into<U> for aya_bpf::RawTracepointArgs where U: core::convert::From<T>
+pub fn aya_bpf::RawTracepointArgs::into(self) -> U
+impl<T, U> core::convert::TryFrom<U> for aya_bpf::RawTracepointArgs where U: core::convert::Into<T>
+pub type aya_bpf::RawTracepointArgs::Error = core::convert::Infallible
+pub fn aya_bpf::RawTracepointArgs::try_from(value: U) -> core::result::Result<T, <T as core::convert::TryFrom<U>>::Error>
+impl<T, U> core::convert::TryInto<U> for aya_bpf::RawTracepointArgs where U: core::convert::TryFrom<T>
+pub type aya_bpf::RawTracepointArgs::Error = <U as core::convert::TryFrom<T>>::Error
+pub fn aya_bpf::RawTracepointArgs::try_into(self) -> core::result::Result<U, <U as core::convert::TryFrom<T>>::Error>
+impl<T> core::any::Any for aya_bpf::RawTracepointArgs where T: 'static + core::marker::Sized
+pub fn aya_bpf::RawTracepointArgs::type_id(&self) -> core::any::TypeId
+impl<T> core::borrow::Borrow<T> for aya_bpf::RawTracepointArgs where T: core::marker::Sized
+pub fn aya_bpf::RawTracepointArgs::borrow(&self) -> &T
+impl<T> core::borrow::BorrowMut<T> for aya_bpf::RawTracepointArgs where T: core::marker::Sized
+pub fn aya_bpf::RawTracepointArgs::borrow_mut(&mut self) -> &mut T
+impl<T> core::convert::From<T> for aya_bpf::RawTracepointArgs
+pub fn aya_bpf::RawTracepointArgs::from(t: T) -> T
 pub const aya_bpf::TASK_COMM_LEN: usize = 16usize
 pub trait aya_bpf::BpfContext
 pub fn aya_bpf::BpfContext::as_ptr(&self) -> *mut core::ffi::c_void