bpf: Handle raw tracepoint arguments

ebpf/aya-ebpf/src/args.rs

@@ -10,7 +10,7 @@ use crate::bindings::user_pt_regs as pt_regs;
 // riscv64 uses user_regs_struct instead of 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.
@@ -418,3 +418,110 @@ 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 {
+    /// Creates a new instance of `RawTracepointArgs` from the given
+    /// `bpf_raw_tracepoint_args` raw pointer to allow easier access
+    /// to raw tracepoint argumetns.
+    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. To be precise, the wrapped
+    /// `bpf_raw_tracepoint_args` binding defines it as `__IncompleteArrayField<__u64>` and the
+    /// original C type as `__u64 args[0]`. This method provides a way to access these arguments
+    /// conveniently in Rust using `__IncompleteArrayField<T>::as_slice` to represent that array
+    /// as a slice of length n and then retrieve the n-th element of it.
+    ///
+    /// 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. Such check is impossible to do, because the
+    /// tracepoint context doesn't contain any information about number of arguments.
+    ///
+    /// This method also cannot guarantee that the requested type matches the actual value type.
+    /// Wrong assumptions about types can lead to undefined behavior. The tracepoint context
+    /// doesn't provide any type information.
+    ///
+    /// The caller is responsible for ensuring they have accurate knowledge of the arguments
+    /// and their respective types for the accessed tracepoint context.
+    pub unsafe fn arg<T: FromRawTracepointArgs>(&self, n: usize) -> *const T {
+        &T::from_argument(&*self.args, n)
+    }
+}
+
+pub unsafe trait FromRawTracepointArgs: Sized {
+    /// 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. To be precise, the wrapped
+    /// `bpf_raw_tracepoint_args` binding defines it as `__IncompleteArrayField<__u64>` and the
+    /// original C type as `__u64 args[0]`. This method provides a way to access these arguments
+    /// conveniently in Rust using `__IncompleteArrayField<T>::as_slice` to represent that array
+    /// as a slice of length n and then retrieve the n-th element of it.
+    ///
+    /// 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. Such check is impossible to do, because the
+    /// tracepoint context doesn't contain any information about number of arguments.
+    ///
+    /// This method also cannot guarantee that the requested type matches the actual value type.
+    /// Wrong assumptions about types can lead to undefined behavior. The tracepoint context
+    /// doesn't provide any type information.
+    ///
+    /// The caller is responsible for ensuring they have accurate knowledge of the arguments
+    /// and their respective types for the accessed tracepoint context.
+    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
+        // They are represented as `__IncompleteArrayField<T>` in the Rust
+        // wraapper.
+        //
+        // The most convenient way of accessing such type in Rust is to use
+        // `__IncompleteArrayField<T>::as_slice` to represent that array as a
+        // slice of length n and then retrieve the n-th element of it.
+        //
+        // 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);

ebpf/aya-ebpf/src/lib.rs

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

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

@@ -1,19 +1,25 @@
 use core::ffi::c_void;
 
-use crate::EbpfContext;
+use crate::{args::FromRawTracepointArgs, bindings::bpf_raw_tracepoint_args, EbpfContext};
 
 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 EbpfContext for RawTracePointContext {
     fn as_ptr(&self) -> *mut c_void {
-        self.ctx
+        self.ctx as *mut c_void
     }
 }

test/integration-common/src/lib.rs

@@ -14,6 +14,19 @@ pub mod bpf_probe_read {
     unsafe impl aya::Pod for TestResult {}
 }
 
+pub mod raw_tracepoint {
+    #[repr(C)]
+    #[derive(Clone, Copy)]
+    pub struct SysEnterEvent {
+        pub common_type: u16,
+        pub common_flags: u8,
+        _padding: u8, // Padding must be explicit to ensure zero-initialization.
+    }
+
+    #[cfg(feature = "user")]
+    unsafe impl aya::Pod for SysEnterEvent {}
+}
+
 pub mod ring_buf {
     // This structure's definition is duplicated in the probe.
     #[repr(C)]

test/integration-ebpf/Cargo.toml

@@ -42,6 +42,10 @@ path = "src/name_test.rs"
 name = "pass"
 path = "src/pass.rs"
 
+[[bin]]
+name = "raw_tracepoint"
+path = "src/raw_tracepoint.rs"
+
 [[bin]]
 name = "redirect"
 path = "src/redirect.rs"

test/integration-ebpf/src/raw_tracepoint.rs

@@ -0,0 +1,33 @@
+#![no_std]
+#![no_main]
+
+use aya_ebpf::{
+    macros::{map, raw_tracepoint},
+    maps::Array,
+    programs::RawTracePointContext,
+};
+use integration_common::raw_tracepoint::SysEnterEvent;
+
+#[map]
+static RESULT: Array<SysEnterEvent> = Array::with_max_entries(1, 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

@@ -20,6 +20,8 @@ pub const MAP_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/ma
 pub const MEMMOVE_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/memmove_test"));
 pub const NAME_TEST: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/name_test"));
 pub const PASS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/pass"));
+pub const RAW_TRACEPOINT: &[u8] =
+    include_bytes_aligned!(concat!(env!("OUT_DIR"), "/raw_tracepoint"));
 pub const REDIRECT: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/redirect"));
 pub const RELOCATIONS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/relocations"));
 pub const RING_BUF: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/ring_buf"));

test/integration-test/src/tests.rs

@@ -5,6 +5,7 @@ mod info;
 mod iter;
 mod load;
 mod log;
+mod raw_tracepoint;
 mod rbpf;
 mod relocations;
 mod ring_buf;

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

@@ -0,0 +1,19 @@
+use aya::{maps::Array, programs::RawTracePoint, Ebpf};
+use integration_common::raw_tracepoint::SysEnterEvent;
+
+#[test]
+fn raw_tracepoint() {
+    let mut bpf = Ebpf::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 SysEnterEvent {
+        common_type,
+        common_flags,
+        ..
+    } = map.get(&0, 0).unwrap();
+    assert_ne!(common_type, 0);
+    assert_ne!(common_flags, 0);
+}

xtask/public-api/aya-ebpf.txt

@@ -1540,6 +1540,7 @@ pub fn aya_ebpf::programs::probe::ProbeContext::from(t: T) -> T
 pub mod aya_ebpf::programs::raw_tracepoint
 pub struct aya_ebpf::programs::raw_tracepoint::RawTracePointContext
 impl aya_ebpf::programs::raw_tracepoint::RawTracePointContext
+pub unsafe fn aya_ebpf::programs::raw_tracepoint::RawTracePointContext::arg<T: FromRawTracepointArgs>(&self, n: usize) -> T
 pub fn aya_ebpf::programs::raw_tracepoint::RawTracePointContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::raw_tracepoint::RawTracePointContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::raw_tracepoint::RawTracePointContext
 pub fn aya_ebpf::programs::raw_tracepoint::RawTracePointContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2242,6 +2243,7 @@ impl<T> core::convert::From<T> for aya_ebpf::programs::probe::ProbeContext
 pub fn aya_ebpf::programs::probe::ProbeContext::from(t: T) -> T
 pub struct aya_ebpf::programs::RawTracePointContext
 impl aya_ebpf::programs::raw_tracepoint::RawTracePointContext
+pub unsafe fn aya_ebpf::programs::raw_tracepoint::RawTracePointContext::arg<T: FromRawTracepointArgs>(&self, n: usize) -> T
 pub fn aya_ebpf::programs::raw_tracepoint::RawTracePointContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::raw_tracepoint::RawTracePointContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::raw_tracepoint::RawTracePointContext
 pub fn aya_ebpf::programs::raw_tracepoint::RawTracePointContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2685,6 +2687,32 @@ impl<T> core::borrow::BorrowMut<T> for aya_ebpf::PtRegs where T: ?core::marker::
 pub fn aya_ebpf::PtRegs::borrow_mut(&mut self) -> &mut T
 impl<T> core::convert::From<T> for aya_ebpf::PtRegs
 pub fn aya_ebpf::PtRegs::from(t: T) -> T
+pub struct aya_ebpf::RawTracepointArgs
+impl aya_ebpf::RawTracepointArgs
+pub unsafe fn aya_ebpf::RawTracepointArgs::arg<T: FromRawTracepointArgs>(&self, n: usize) -> *const T
+pub fn aya_ebpf::RawTracepointArgs::new(args: *mut aya_ebpf_bindings::x86_64::bindings::bpf_raw_tracepoint_args) -> Self
+impl core::marker::Freeze for aya_ebpf::RawTracepointArgs
+impl !core::marker::Send for aya_ebpf::RawTracepointArgs
+impl !core::marker::Sync for aya_ebpf::RawTracepointArgs
+impl core::marker::Unpin for aya_ebpf::RawTracepointArgs
+impl core::panic::unwind_safe::RefUnwindSafe for aya_ebpf::RawTracepointArgs
+impl core::panic::unwind_safe::UnwindSafe for aya_ebpf::RawTracepointArgs
+impl<T, U> core::convert::Into<U> for aya_ebpf::RawTracepointArgs where U: core::convert::From<T>
+pub fn aya_ebpf::RawTracepointArgs::into(self) -> U
+impl<T, U> core::convert::TryFrom<U> for aya_ebpf::RawTracepointArgs where U: core::convert::Into<T>
+pub type aya_ebpf::RawTracepointArgs::Error = core::convert::Infallible
+pub fn aya_ebpf::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_ebpf::RawTracepointArgs where U: core::convert::TryFrom<T>
+pub type aya_ebpf::RawTracepointArgs::Error = <U as core::convert::TryFrom<T>>::Error
+pub fn aya_ebpf::RawTracepointArgs::try_into(self) -> core::result::Result<U, <U as core::convert::TryFrom<T>>::Error>
+impl<T> core::any::Any for aya_ebpf::RawTracepointArgs where T: 'static + ?core::marker::Sized
+pub fn aya_ebpf::RawTracepointArgs::type_id(&self) -> core::any::TypeId
+impl<T> core::borrow::Borrow<T> for aya_ebpf::RawTracepointArgs where T: ?core::marker::Sized
+pub fn aya_ebpf::RawTracepointArgs::borrow(&self) -> &T
+impl<T> core::borrow::BorrowMut<T> for aya_ebpf::RawTracepointArgs where T: ?core::marker::Sized
+pub fn aya_ebpf::RawTracepointArgs::borrow_mut(&mut self) -> &mut T
+impl<T> core::convert::From<T> for aya_ebpf::RawTracepointArgs
+pub fn aya_ebpf::RawTracepointArgs::from(t: T) -> T
 pub const aya_ebpf::TASK_COMM_LEN: usize
 pub trait aya_ebpf::EbpfContext
 pub fn aya_ebpf::EbpfContext::as_ptr(&self) -> *mut core::ffi::c_void