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.

Author: vadorovsky

Cargo.toml

@@ -7,6 +7,7 @@ members = [
     "aya-obj",
     "aya-tool",
     "init",
+    "test/integration-common",
     "test/integration-test",
     "xtask",
 
@@ -68,6 +69,7 @@ diff = { version = "0.1.13", default-features = false }
 env_logger = { version = "0.10", default-features = false }
 hashbrown = { version = "0.14", default-features = false }
 indoc = { version = "2.0", default-features = false }
+integration-common = { path = "test/integration-common", default-features = false }
 integration-ebpf = { path = "test/integration-ebpf", default-features = false }
 lazy_static = { version = "1", default-features = false }
 libc = { version = "0.2.105", default-features = false }

bpf/aya-bpf/src/args.rs

@@ -1,3 +1,5 @@
+use aya_bpf_bindings::bindings::bpf_raw_tracepoint_args;
+
 use crate::{cty::c_void, helpers::bpf_probe_read};
 
 // aarch64 uses user_pt_regs instead of pt_regs
@@ -326,3 +328,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-common/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "integration-common"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+default = []
+user = ["aya"]
+
+[dependencies]
+aya = { workspace = true, optional = true }

test/integration-common/src/lib.rs

@@ -0,0 +1,3 @@
+#![no_std]
+
+pub mod raw_tracepoint;

test/integration-common/src/raw_tracepoint.rs

@@ -0,0 +1,20 @@
+#[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,
+        }
+    }
+}
+
+#[cfg(feature = "user")]
+unsafe impl aya::Pod for SysEnterEvent {}

test/integration-ebpf/Cargo.toml

@@ -7,6 +7,7 @@ publish = false
 [dependencies]
 aya-bpf = { path = "../../bpf/aya-bpf" }
 aya-log-ebpf = { path = "../../bpf/aya-log-ebpf" }
+integration-common = { workspace = true }
 
 [build-dependencies]
 which = { workspace = true }
@@ -43,3 +44,7 @@ path = "src/bpf_probe_read.rs"
 [[bin]]
 name = "two_progs"
 path = "src/two_progs.rs"
+
+[[bin]]
+name = "raw_tracepoint"
+path = "src/raw_tracepoint.rs"

test/integration-ebpf/src/raw_tracepoint.rs

@@ -0,0 +1,34 @@
+#![no_std]
+#![no_main]
+
+use aya_bpf::{
+    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/Cargo.toml

@@ -10,6 +10,7 @@ assert_matches = { workspace = true }
 aya = { workspace = true }
 aya-log = { workspace = true }
 aya-obj = { workspace = true }
+integration-common = { workspace = true, features = ["user"] }
 libc = { workspace = true }
 log = { workspace = true }
 netns-rs = { workspace = true }

test/integration-test/src/lib.rs

@@ -19,6 +19,8 @@ pub const RELOCATIONS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "
 pub const TWO_PROGS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/two_progs"));
 pub const BPF_PROBE_READ: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/bpf_probe_read"));
+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,16 @@
+use aya::{maps::Array, programs::RawTracePoint, Bpf};
+use integration_common::raw_tracepoint::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);
+}