aya-bpf: implement argument coercion for pt_regs and BTF programs

Implements argument and return value coercion helpers for:
    - LSM programs
    - BTF tracepoints
    - [ku]{ret}probes

Signed-off-by: William Findlay <william@williamfindlay.com>

bpf/aya-bpf/build.rs

@@ -0,0 +1,9 @@
+use std::env;
+
+fn main() {
+    if env::var("CARGO_CFG_BPF_TARGET_ARCH").is_err() {
+        let arch = env::var("HOST").unwrap();
+        let arch = arch.split_once('-').map_or(&*arch, |x| x.0);
+        println!("cargo:rustc-cfg=bpf_target_arch=\"{}\"", arch);
+    }
+}

bpf/aya-bpf/src/args.rs

@@ -0,0 +1,186 @@
+use crate::cty::c_void;
+
+// aarch64 uses user_pt_regs instead of pt_regs
+#[cfg(not(bpf_target_arch = "aarch64"))]
+use crate::bindings::pt_regs;
+#[cfg(bpf_target_arch = "aarch64")]
+use crate::bindings::user_pt_regs as pt_regs;
+
+/// A trait that indicates a valid type for an argument which can be coerced from a BTF
+/// context.
+///
+/// Users should not implement this trait.
+///
+/// SAFETY: This trait is _only_ safe to implement on primitive types that can fit into
+/// a `u64`. For example, integers and raw pointers may be coerced from a BTF context.
+pub unsafe trait FromBtfArgument: Sized {
+    /// Coerces a `T` from the `n`th argument from a BTF context where `n` starts
+    /// at 0 and increases by 1 for each successive argument.
+    ///
+    /// SAFETY: This function is deeply unsafe, as we are reading raw pointers into kernel
+    /// 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 impl<T> FromBtfArgument for *const T {
+    unsafe fn from_argument(ctx: *const c_void, n: usize) -> *const T {
+        // 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 _
+    }
+}
+
+/// 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 {
+                // 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 _
+            }
+        }
+    };
+}
+
+unsafe_impl_from_btf_argument!(u8);
+unsafe_impl_from_btf_argument!(u16);
+unsafe_impl_from_btf_argument!(u32);
+unsafe_impl_from_btf_argument!(u64);
+unsafe_impl_from_btf_argument!(i8);
+unsafe_impl_from_btf_argument!(i16);
+unsafe_impl_from_btf_argument!(i32);
+unsafe_impl_from_btf_argument!(i64);
+unsafe_impl_from_btf_argument!(usize);
+unsafe_impl_from_btf_argument!(isize);
+
+/// A trait that indicates a valid type for an argument which can be coerced from
+/// a pt_regs context.
+///
+/// Any implementation of this trait is strictly architecture-specific and depends on the
+/// layout of the underlying pt_regs struct and the target processor's calling
+/// conventions. Users should not implement this trait.
+pub trait FromPtRegs: Sized {
+    /// Coerces a `T` from the `n`th argument of a pt_regs context where `n` starts
+    /// at 0 and increases by 1 for each successive argument.
+    fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self>;
+
+    /// Coerces a `T` from the return value of a pt_regs context.
+    fn from_retval(ctx: &pt_regs) -> Option<Self>;
+}
+
+#[cfg(bpf_target_arch = "x86_64")]
+impl<T> FromPtRegs for *const T {
+    fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        match n {
+            0 => ctx.rdi().map(|v| v as _),
+            1 => ctx.rsi().map(|v| v as _),
+            2 => ctx.rdx().map(|v| v as _),
+            3 => ctx.rcx().map(|v| v as _),
+            4 => ctx.r8().map(|v| v as _),
+            5 => ctx.r9().map(|v| v as _),
+            _ => None,
+        }
+    }
+
+    fn from_retval(ctx: &pt_regs) -> Option<Self> {
+        ctx.rax().map(|v| v as _)
+    }
+}
+
+#[cfg(bpf_target_arch = "armv7")]
+impl<T> FromPtRegs for *const T {
+    fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        if n <= 6 {
+            ctx.uregs().map(|regs| regs[n] as _)
+        } else {
+            None
+        }
+    }
+
+    fn from_retval(ctx: &pt_regs) -> Option<Self> {
+        ctx.uregs().map(|regs| regs[0] as _)
+    }
+}
+
+#[cfg(bpf_target_arch = "aarch64")]
+impl<T> FromPtRegs for *const T {
+    fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        if n <= 7 {
+            ctx.regs().map(|regs| regs[n] as _)
+        } else {
+            None
+        }
+    }
+
+    fn from_retval(ctx: &pt_regs) -> Option<Self> {
+        ctx.regs().map(|regs| regs[0] as _)
+    }
+}
+
+/// Helper macro to implement [`FromPtRegs`] for a primitive type.
+macro_rules! impl_from_pt_regs {
+    ($type:ident) => {
+        #[cfg(bpf_target_arch = "x86_64")]
+        impl FromPtRegs for $type {
+            fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+                match n {
+                    0 => ctx.rdi().map(|v| v as _),
+                    1 => ctx.rsi().map(|v| v as _),
+                    2 => ctx.rdx().map(|v| v as _),
+                    3 => ctx.rcx().map(|v| v as _),
+                    4 => ctx.r8().map(|v| v as _),
+                    5 => ctx.r9().map(|v| v as _),
+                    _ => None,
+                }
+            }
+
+            fn from_retval(ctx: &pt_regs) -> Option<Self> {
+                ctx.rax().map(|v| v as _)
+            }
+        }
+
+        #[cfg(bpf_target_arch = "armv7")]
+        impl FromPtRegs for $type {
+            fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+                if n <= 6 {
+                    ctx.uregs().map(|regs| regs[n] as _)
+                } else {
+                    None
+                }
+            }
+
+            fn from_retval(ctx: &pt_regs) -> Option<Self> {
+                ctx.uregs().map(|regs| regs[0] as _)
+            }
+        }
+
+        #[cfg(bpf_target_arch = "aarch64")]
+        impl FromPtRegs for $type {
+            fn from_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+                if n <= 7 {
+                    ctx.regs().map(|regs| regs[n] as _)
+                } else {
+                    None
+                }
+            }
+
+            fn from_retval(ctx: &pt_regs) -> Option<Self> {
+                ctx.regs().map(|regs| regs[0] as _)
+            }
+        }
+    };
+}
+
+impl_from_pt_regs!(u8);
+impl_from_pt_regs!(u16);
+impl_from_pt_regs!(u32);
+impl_from_pt_regs!(u64);
+impl_from_pt_regs!(i8);
+impl_from_pt_regs!(i16);
+impl_from_pt_regs!(i32);
+impl_from_pt_regs!(i64);
+impl_from_pt_regs!(usize);
+impl_from_pt_regs!(isize);

bpf/aya-bpf/src/lib.rs

@@ -3,6 +3,7 @@
 
 pub use aya_bpf_bindings::bindings;
 
+mod args;
 pub mod helpers;
 pub mod maps;
 pub mod programs;

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

@@ -1,6 +1,6 @@
 use core::ffi::c_void;
 
-use crate::BpfContext;
+use crate::{args::FromBtfArgument, BpfContext};
 
 pub struct LsmContext {
     ctx: *mut c_void,
@@ -10,6 +10,49 @@ impl LsmContext {
     pub fn new(ctx: *mut c_void) -> LsmContext {
         LsmContext { ctx }
     }
+
+    /// Returns the `n`th argument passed to the LSM hook, starting from 0.
+    ///
+    /// You can refer to [the kernel's list of LSM hook definitions][1] to find the
+    /// appropriate argument list for your LSM hook, where the argument list starts
+    /// _after_ the third parameter to the kernel's `LSM_HOOK` macro.
+    ///
+    /// LSM probes specifically have access to an additional argument `retval: int`
+    /// which provides the return value of the previous LSM program that was called on
+    /// this code path, or 0 if this is the first LSM program to be called. This phony
+    /// argument is always last in the argument list.
+    ///
+    /// SAFETY: This function is deeply unsafe, as we are reading raw pointers into kernel memory.
+    /// In particular, the value of `n` must not exceed the number of function arguments.
+    /// Luckily, the BPF verifier will catch this for us.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # #![allow(dead_code)]
+    /// # use aya_bpf::{programs::LsmContext, cty::{c_int, c_ulong}};
+    /// 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);
+    ///
+    ///     // You can then do stuff with addr and retval down here.
+    ///
+    ///     // To practice good LSM hygiene, let's defer to a previous retval
+    ///     // if available:
+    ///     if (retval != 0) {
+    ///         return Ok(retval);
+    ///     }
+    ///
+    ///     Ok(0)
+    /// }
+    /// ```
+    ///
+    /// [1]: https://elixir.bootlin.com/linux/latest/source/include/linux/lsm_hook_defs.h
+    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> T {
+        T::from_argument(self.ctx as *const _, n)
+    }
 }
 
 impl BpfContext for LsmContext {

bpf/aya-bpf/src/programs/probe.rs

@@ -1,6 +1,12 @@
 use core::ffi::c_void;
 
-use crate::{bindings::pt_regs, BpfContext};
+use crate::{args::FromPtRegs, BpfContext};
+
+// aarch64 uses user_pt_regs instead of pt_regs
+#[cfg(not(bpf_target_arch = "aarch64"))]
+use crate::bindings::pt_regs;
+#[cfg(bpf_target_arch = "aarch64")]
+use crate::bindings::user_pt_regs as pt_regs;
 
 pub struct ProbeContext {
     pub regs: *mut pt_regs,
@@ -12,6 +18,50 @@ impl ProbeContext {
             regs: ctx as *mut pt_regs,
         }
     }
+
+    /// Returns the `n`th argument to passed to the probe function, starting from 0.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # #![allow(non_camel_case_types)]
+    /// # #![allow(dead_code)]
+    /// # use aya_bpf::{programs::ProbeContext, cty::c_int, helpers::bpf_probe_read};
+    /// # type pid_t = c_int;
+    /// # struct task_struct {
+    /// #     pid: pid_t,
+    /// # }
+    /// unsafe fn try_kprobe_try_to_wake_up(ctx: ProbeContext) -> Result<u32, u32> {
+    ///     let tp: *const task_struct = ctx.arg(0).ok_or(1u32)?;
+    ///     let pid = bpf_probe_read(&(*tp).pid as *const pid_t).map_err(|_| 1u32)?;
+    ///
+    ///     // Do something with pid or something else with tp
+    ///
+    ///     Ok(0)
+    /// }
+    /// ```
+    pub fn arg<T: FromPtRegs>(&self, n: usize) -> Option<T> {
+        T::from_argument(unsafe { &*self.regs }, n)
+    }
+
+    /// Returns the return value of the probed function.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # #![allow(dead_code)]
+    /// # use aya_bpf::{programs::ProbeContext, cty::c_int};
+    /// unsafe fn try_kretprobe_try_to_wake_up(ctx: ProbeContext) -> Result<u32, u32> {
+    ///     let retval: c_int = ctx.ret().ok_or(1u32)?;
+    ///
+    ///     // Do something with retval
+    ///
+    ///     Ok(0)
+    /// }
+    /// ```
+    pub fn ret<T: FromPtRegs>(&self) -> Option<T> {
+        T::from_retval(unsafe { &*self.regs })
+    }
 }
 
 impl BpfContext for ProbeContext {

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

@@ -1,6 +1,6 @@
 use core::ffi::c_void;
 
-use crate::BpfContext;
+use crate::{args::FromBtfArgument, BpfContext};
 
 pub struct BtfTracePointContext {
     ctx: *mut c_void,
@@ -10,6 +10,39 @@ impl BtfTracePointContext {
     pub fn new(ctx: *mut c_void) -> BtfTracePointContext {
         BtfTracePointContext { ctx }
     }
+
+    /// Returns the `n`th argument of the BTF tracepoint, starting from 0.
+    ///
+    /// You can use the tplist tool provided by bcc to get a list of tracepoints and their
+    /// arguments. TODO: document this better, possibly add a tplist alternative to aya.
+    ///
+    /// SAFETY: This function is deeply unsafe, as we are reading raw pointers into kernel memory.
+    /// In particular, the value of `n` must not exceed the number of function arguments.
+    /// Luckily, the BPF verifier will catch this for us.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # #![allow(dead_code)]
+    /// # use aya_bpf::{programs::BtfTracePointContext, cty::{c_int, c_ulong, c_char}};
+    /// unsafe fn try_tp_btf_sched_process_fork(ctx: BtfTracePointContext) -> Result<u32, u32> {
+    ///     // 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);
+    ///
+    ///     // You can then do stuff with parent_pidm parent_comm, child_pid, and
+    ///     // child_comm down here.
+    ///
+    ///     Ok(0)
+    /// }
+    /// ```
+    ///
+    /// [1]: https://elixir.bootlin.com/linux/latest/source/include/linux/lsm_hook_defs.h
+    pub unsafe fn arg<T: FromBtfArgument>(&self, n: usize) -> T {
+        T::from_argument(self.ctx as *const _, n)
+    }
 }
 
 impl BpfContext for BtfTracePointContext {