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,17 @@ 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 _
+        // *(ctx as *const usize).add(n) as _
+        bpf_probe_read((ctx as *const usize).add(n))
+            .map(|v| v as *const _)
+            .ok()
     }
 }
 
@@ -42,10 +45,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

@@ -25,6 +25,10 @@ path = "src/bpf_probe_read.rs"
 name = "log"
 path = "src/log.rs"
 
+[[bin]]
+name = "args"
+path = "src/args.rs"
+
 [[bin]]
 name = "map_test"
 path = "src/map_test.rs"

test/integration-ebpf/src/args.rs

@@ -0,0 +1,33 @@
+#![no_std]
+#![no_main]
+
+use aya_bpf::{
+    cty::{c_long, c_longlong},
+    macros::{fentry, kprobe},
+    programs::{FEntryContext, ProbeContext},
+};
+
+#[kprobe]
+pub fn kprobe_vfs_write(ctx: ProbeContext) {
+    let _ = try_kprobe_vfs_write(ctx);
+}
+
+fn try_kprobe_vfs_write(ctx: ProbeContext) -> Result<(), c_long> {
+    let _pos: *const c_longlong = ctx.arg(3).ok_or(1)?;
+    Ok(())
+}
+
+#[fentry]
+pub fn fentry_vfs_write(ctx: FEntryContext) {
+    let _ = try_fentry_vfs_write(ctx);
+}
+
+fn try_fentry_vfs_write(ctx: FEntryContext) -> Result<(), c_long> {
+    let _pos: *const c_longlong = unsafe { ctx.arg(3).ok_or(1)? };
+    Ok(())
+}
+
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    unsafe { core::hint::unreachable_unchecked() }
+}

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,34 @@
+use aya::{
+    include_bytes_aligned,
+    programs::{FEntry, KProbe},
+    Bpf, Btf,
+};
+
+use super::{integration_test, IntegrationTest};
+
+#[integration_test]
+fn kprobe_args() {
+    let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/args");
+    let mut bpf = Bpf::load(bytes).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();
+}
+
+#[integration_test]
+fn fentry_args() {
+    let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/args");
+    let mut bpf = Bpf::load(bytes).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();
+}