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

ebpf/aya-ebpf/src/args.rs

@@ -27,14 +27,16 @@ 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()
     }
 }
 
@@ -42,10 +44,10 @@ unsafe impl<T> FromBtfArgument for *const T {
 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

@@ -31,7 +31,7 @@ impl FEntryContext {
     ///     Ok(0)
     /// }
     /// ```
-    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

@@ -31,7 +31,7 @@ impl FExitContext {
     ///     Ok(0)
     /// }
     /// ```
-    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

@@ -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

@@ -40,7 +40,7 @@ impl BtfTracePointContext {
     /// ```
     ///
     /// [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

@@ -17,6 +17,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

@@ -12,6 +12,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
@@ -1956,7 +1956,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
@@ -2046,7 +2046,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
@@ -2102,7 +2102,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
@@ -2130,7 +2130,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
@@ -2158,7 +2158,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