ebpf: Call bpf_probe_read on *const T BTF arguments

It's necessary to call `bpf_probe_read` not only for pointers retrieved
from `PtRegs`, but also from BTF arguments.

`bpf_probe_read` might return an error, so the return type of `.arg()`
methods in contexts handling BTF arguments changes from `T` to
`Option<T>`. `None` is returned when `bpf_probe_read` call is not
successful.

Fixes: #542

aya-ebpf-macros/src/lib.rs

@@ -485,8 +485,8 @@ pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 ///
 /// unsafe fn try_filename_lookup(ctx: FEntryContext) -> Result<i32, i32> {
-///     let _f: *const filename = ctx.arg(1);
-///     let _p: *const path = ctx.arg(3);
+///     let _f: *const filename = ctx.arg(1).ok_or(-1)?;
+///     let _p: *const path = ctx.arg(3).ok_or(-1)?;
 ///
 ///     Ok(0)
 /// }
@@ -527,8 +527,8 @@ pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 ///
 /// unsafe fn try_filename_lookup(ctx: FExitContext) -> Result<i32, i32> {
-///     let _f: *const filename = ctx.arg(1);
-///     let _p: *const path = ctx.arg(3);
+///     let _f: *const filename = ctx.arg(1).ok_or(-1)?;
+///     let _p: *const path = ctx.arg(3).ok_or(-1)?;
 ///
 ///     Ok(0)
 /// }

ebpf/aya-ebpf/src/args.rs

@@ -27,25 +27,27 @@ pub unsafe trait FromBtfArgument: Sized {
     /// memory. In particular, the value of `n` must not exceed the number of function
     /// arguments. Moreover, `ctx` must be a valid pointer to a BTF context, and `T` must
     /// be the right type for the given argument.
-    unsafe fn from_argument(ctx: *const c_void, n: usize) -> Self;
+    unsafe fn from_argument(ctx: *const c_void, n: usize) -> Option<Self>;
 }
 
 unsafe impl<T> FromBtfArgument for *const T {
-    unsafe fn from_argument(ctx: *const c_void, n: usize) -> *const T {
+    unsafe fn from_argument(ctx: *const c_void, n: usize) -> Option<Self> {
         // BTF arguments are exposed as an array of `usize` where `usize` can
         // either be treated as a pointer or a primitive type
-        *(ctx as *const usize).add(n) as _
+        bpf_probe_read((ctx as *const usize).add(n))
+            .map(|v| v as *const _)
+            .ok()
     }
 }
 
 /// Helper macro to implement [`FromBtfArgument`] for a primitive type.
 macro_rules! unsafe_impl_from_btf_argument {
     ($type:ident) => {
         unsafe impl FromBtfArgument for $type {
-            unsafe fn from_argument(ctx: *const c_void, n: usize) -> Self {
+            unsafe fn from_argument(ctx: *const c_void, n: usize) -> Option<Self> {
                 // BTF arguments are exposed as an array of `usize` where `usize` can
                 // either be treated as a pointer or a primitive type
-                *(ctx as *const usize).add(n) as _
+                Some(*(ctx as *const usize).add(n) as _)
             }
         }
     };

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

@@ -23,15 +23,15 @@ impl FEntryContext {
     /// # struct task_struct {
     /// #     pid: pid_t,
     /// # }
-    /// unsafe fn try_fentry_try_to_wake_up(ctx: FEntryContext) -> Result<u32, u32> {
-    ///     let tp: *const task_struct = ctx.arg(0);
+    /// unsafe fn try_fentry_try_to_wake_up(ctx: FEntryContext) -> Result<(), i32> {
+    ///     let tp: *const task_struct = ctx.arg(0).ok_or(-1)?;
     ///
     ///     // Do something with tp
     ///
-    ///     Ok(0)
+    ///     Ok(())
     /// }
     /// ```
-    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> T {
+    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> Option<T> {
         T::from_argument(self.ctx as *const _, n)
     }
 }

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

@@ -23,15 +23,15 @@ impl FExitContext {
     /// # struct task_struct {
     /// #     pid: pid_t,
     /// # }
-    /// unsafe fn try_filename_lookup(ctx: FExitContext) -> Result<u32, u32> {
-    ///     let tp: *const task_struct = ctx.arg(0);
+    /// unsafe fn try_filename_lookup(ctx: FExitContext) -> Result<(), i32> {
+    ///     let tp: *const task_struct = ctx.arg(0).ok_or(-1)?;
     ///
     ///     // Do something with tp
     ///
-    ///     Ok(0)
+    ///     Ok(())
     /// }
     /// ```
-    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> T {
+    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> Option<T> {
         T::from_argument(self.ctx as *const _, n)
     }
 }

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

@@ -34,8 +34,8 @@ impl LsmContext {
     /// unsafe fn try_lsm_mmap_addr(ctx: LsmContext) -> Result<i32, i32> {
     ///     // In the kernel, this hook is defined as:
     ///     //   LSM_HOOK(int, 0, mmap_addr, unsigned long addr)
-    ///     let addr: c_ulong = ctx.arg(0);
-    ///     let retval: c_int = ctx.arg(1);
+    ///     let addr: c_ulong = ctx.arg(0).ok_or(-1)?;
+    ///     let retval: c_int = ctx.arg(1).ok_or(-1)?;
     ///
     ///     // You can then do stuff with addr and retval down here.
     ///
@@ -50,7 +50,7 @@ impl LsmContext {
     /// ```
     ///
     /// [1]: https://elixir.bootlin.com/linux/latest/source/include/linux/lsm_hook_defs.h
-    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> T {
+    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> Option<T> {
         T::from_argument(self.ctx as *const _, n)
     }
 }

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

@@ -25,22 +25,22 @@ impl BtfTracePointContext {
     /// ```no_run
     /// # #![allow(dead_code)]
     /// # use aya_ebpf::{programs::BtfTracePointContext, cty::{c_int, c_ulong, c_char}};
-    /// unsafe fn try_tp_btf_sched_process_fork(ctx: BtfTracePointContext) -> Result<u32, u32> {
+    /// unsafe fn try_tp_btf_sched_process_fork(ctx: BtfTracePointContext) -> Result<(), i32> {
     ///     // Grab arguments
-    ///     let parent_comm: *const c_char = ctx.arg(0);
-    ///     let parent_pid: c_int = ctx.arg(1);
-    ///     let child_comm: *const c_char = ctx.arg(2);
-    ///     let child_pid: c_int = ctx.arg(3);
+    ///     let parent_comm: *const c_char = ctx.arg(0).ok_or(-1)?;
+    ///     let parent_pid: c_int = ctx.arg(1).ok_or(-1)?;
+    ///     let child_comm: *const c_char = ctx.arg(2).ok_or(-1)?;
+    ///     let child_pid: c_int = ctx.arg(3).ok_or(-1)?;
     ///
     ///     // You can then do stuff with parent_pidm parent_comm, child_pid, and
     ///     // child_comm down here.
     ///
-    ///     Ok(0)
+    ///     Ok(())
     /// }
     /// ```
     ///
     /// [1]: https://elixir.bootlin.com/linux/latest/source/include/linux/lsm_hook_defs.h
-    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> T {
+    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> Option<T> {
         T::from_argument(self.ctx as *const _, n)
     }
 }

test/integration-ebpf/Cargo.toml

@@ -18,6 +18,10 @@ network-types = "0.0.7"
 which = { workspace = true }
 xtask = { path = "../../xtask" }
 
+[[bin]]
+name = "args"
+path = "src/args.rs"
+
 [[bin]]
 name = "bpf_probe_read"
 path = "src/bpf_probe_read.rs"

test/integration-ebpf/src/args.rs

@@ -0,0 +1,23 @@
+#![no_std]
+#![no_main]
+
+use aya_ebpf::{
+    macros::{fentry, kprobe},
+    programs::{FEntryContext, ProbeContext},
+};
+
+#[kprobe]
+pub fn kprobe_vfs_write(ctx: ProbeContext) {
+    let _: Option<usize> = ctx.arg(3);
+}
+
+#[fentry]
+pub fn fentry_vfs_write(ctx: FEntryContext) {
+    let _: Option<usize> = unsafe { ctx.arg(3) };
+}
+
+#[cfg(not(test))]
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    loop {}
+}

test/integration-test/src/lib.rs

@@ -13,6 +13,7 @@ pub const TEXT_64_64_RELOC: &[u8] =
 pub const VARIABLES_RELOC: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/variables_reloc.bpf.o"));
 
+pub const ARGS: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/args"));
 pub const BPF_PROBE_READ: &[u8] =
     include_bytes_aligned!(concat!(env!("OUT_DIR"), "/bpf_probe_read"));
 pub const LOG: &[u8] = include_bytes_aligned!(concat!(env!("OUT_DIR"), "/log"));

test/integration-test/src/tests.rs

@@ -1,3 +1,4 @@
+mod args;
 mod bpf_probe_read;
 mod btf_relocations;
 mod elf;

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

@@ -0,0 +1,29 @@
+use aya::{
+    programs::{FEntry, KProbe},
+    Btf, Ebpf,
+};
+
+#[test]
+fn kprobe_args() {
+    let mut bpf = Ebpf::load(crate::ARGS).unwrap();
+    let kprobe_vfs_write: &mut KProbe = bpf
+        .program_mut("kprobe_vfs_write")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    kprobe_vfs_write.load().unwrap();
+    kprobe_vfs_write.attach("vfs_write", 0).unwrap();
+}
+
+#[test]
+fn fentry_args() {
+    let mut bpf = Ebpf::load(crate::ARGS).unwrap();
+    let fentry_vfs_write: &mut FEntry = bpf
+        .program_mut("fentry_vfs_write")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    let btf = Btf::from_sys_fs().unwrap();
+    fentry_vfs_write.load("vfs_write", &btf).unwrap();
+    fentry_vfs_write.attach().unwrap();
+}

xtask/public-api/aya-ebpf.txt

@@ -1395,7 +1395,7 @@ pub fn aya_ebpf::programs::device::DeviceContext::from(t: T) -> T
 pub mod aya_ebpf::programs::fentry
 pub struct aya_ebpf::programs::fentry::FEntryContext
 impl aya_ebpf::programs::fentry::FEntryContext
-pub unsafe fn aya_ebpf::programs::fentry::FEntryContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::fentry::FEntryContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::fentry::FEntryContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::fentry::FEntryContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::fentry::FEntryContext
 pub fn aya_ebpf::programs::fentry::FEntryContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -1424,7 +1424,7 @@ pub fn aya_ebpf::programs::fentry::FEntryContext::from(t: T) -> T
 pub mod aya_ebpf::programs::fexit
 pub struct aya_ebpf::programs::fexit::FExitContext
 impl aya_ebpf::programs::fexit::FExitContext
-pub unsafe fn aya_ebpf::programs::fexit::FExitContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::fexit::FExitContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::fexit::FExitContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::fexit::FExitContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::fexit::FExitContext
 pub fn aya_ebpf::programs::fexit::FExitContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -1453,7 +1453,7 @@ pub fn aya_ebpf::programs::fexit::FExitContext::from(t: T) -> T
 pub mod aya_ebpf::programs::lsm
 pub struct aya_ebpf::programs::lsm::LsmContext
 impl aya_ebpf::programs::lsm::LsmContext
-pub unsafe fn aya_ebpf::programs::lsm::LsmContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::lsm::LsmContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::lsm::LsmContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::lsm::LsmContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::lsm::LsmContext
 pub fn aya_ebpf::programs::lsm::LsmContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -1957,7 +1957,7 @@ pub fn aya_ebpf::programs::tc::TcContext::from(t: T) -> T
 pub mod aya_ebpf::programs::tp_btf
 pub struct aya_ebpf::programs::tp_btf::BtfTracePointContext
 impl aya_ebpf::programs::tp_btf::BtfTracePointContext
-pub unsafe fn aya_ebpf::programs::tp_btf::BtfTracePointContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::tp_btf::BtfTracePointContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::tp_btf::BtfTracePointContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::tp_btf::BtfTracePointContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::tp_btf::BtfTracePointContext
 pub fn aya_ebpf::programs::tp_btf::BtfTracePointContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2047,7 +2047,7 @@ impl<T> core::convert::From<T> for aya_ebpf::programs::xdp::XdpContext
 pub fn aya_ebpf::programs::xdp::XdpContext::from(t: T) -> T
 pub struct aya_ebpf::programs::BtfTracePointContext
 impl aya_ebpf::programs::tp_btf::BtfTracePointContext
-pub unsafe fn aya_ebpf::programs::tp_btf::BtfTracePointContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::tp_btf::BtfTracePointContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::tp_btf::BtfTracePointContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::tp_btf::BtfTracePointContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::tp_btf::BtfTracePointContext
 pub fn aya_ebpf::programs::tp_btf::BtfTracePointContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2103,7 +2103,7 @@ impl<T> core::convert::From<T> for aya_ebpf::programs::device::DeviceContext
 pub fn aya_ebpf::programs::device::DeviceContext::from(t: T) -> T
 pub struct aya_ebpf::programs::FEntryContext
 impl aya_ebpf::programs::fentry::FEntryContext
-pub unsafe fn aya_ebpf::programs::fentry::FEntryContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::fentry::FEntryContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::fentry::FEntryContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::fentry::FEntryContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::fentry::FEntryContext
 pub fn aya_ebpf::programs::fentry::FEntryContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2131,7 +2131,7 @@ impl<T> core::convert::From<T> for aya_ebpf::programs::fentry::FEntryContext
 pub fn aya_ebpf::programs::fentry::FEntryContext::from(t: T) -> T
 pub struct aya_ebpf::programs::FExitContext
 impl aya_ebpf::programs::fexit::FExitContext
-pub unsafe fn aya_ebpf::programs::fexit::FExitContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::fexit::FExitContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::fexit::FExitContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::fexit::FExitContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::fexit::FExitContext
 pub fn aya_ebpf::programs::fexit::FExitContext::as_ptr(&self) -> *mut core::ffi::c_void
@@ -2159,7 +2159,7 @@ impl<T> core::convert::From<T> for aya_ebpf::programs::fexit::FExitContext
 pub fn aya_ebpf::programs::fexit::FExitContext::from(t: T) -> T
 pub struct aya_ebpf::programs::LsmContext
 impl aya_ebpf::programs::lsm::LsmContext
-pub unsafe fn aya_ebpf::programs::lsm::LsmContext::arg<T: FromBtfArgument>(&self, n: usize) -> T
+pub unsafe fn aya_ebpf::programs::lsm::LsmContext::arg<T: FromBtfArgument>(&self, n: usize) -> core::option::Option<T>
 pub fn aya_ebpf::programs::lsm::LsmContext::new(ctx: *mut core::ffi::c_void) -> aya_ebpf::programs::lsm::LsmContext
 impl aya_ebpf::EbpfContext for aya_ebpf::programs::lsm::LsmContext
 pub fn aya_ebpf::programs::lsm::LsmContext::as_ptr(&self) -> *mut core::ffi::c_void