From 67623fda649f7d361dddeebbc541ab4fa56518b6 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Tue, 17 May 2022 10:02:24 +0200
Subject: [PATCH] bpf: Add PtRegs wrapper

This adds a portable wrapper around pt_regs and user_pt_regs.
It makes writing Raw Tracepoint or KProbe programs easier when the
arguments are one of these types while also ensuring code is portable
across architectures

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf/src/args.rs | 20 ++++++++++++++++++++
 bpf/aya-bpf/src/lib.rs  |  1 +
 2 files changed, 21 insertions(+)

diff --git a/bpf/aya-bpf/src/args.rs b/bpf/aya-bpf/src/args.rs
index 42bfad9f..c9042db0 100644
--- a/bpf/aya-bpf/src/args.rs
+++ b/bpf/aya-bpf/src/args.rs
@@ -56,6 +56,26 @@ unsafe_impl_from_btf_argument!(i64);
 unsafe_impl_from_btf_argument!(usize);
 unsafe_impl_from_btf_argument!(isize);
 
+pub struct PtRegs {
+    regs: *mut pt_regs,
+}
+
+/// PtRegs is a portable wrapper around
+/// pt_regs and user_pt_regs
+impl PtRegs {
+
+    pub fn new(ctx: *mut c_void) -> Self {
+        PtRegs {
+            regs: ctx as *mut _,
+        }
+    }
+
+    /// Returns the `n`th argument to passed to the probe function, starting from 0.
+    pub fn arg<T: FromPtRegs>(&self, n: usize) -> Option<T> {
+        T::from_argument( unsafe { &*self.regs } , n)
+    }
+}
+
 /// A trait that indicates a valid type for an argument which can be coerced from
 /// a pt_regs context.
 ///
diff --git a/bpf/aya-bpf/src/lib.rs b/bpf/aya-bpf/src/lib.rs
index 688708d1..e62f10bd 100644
--- a/bpf/aya-bpf/src/lib.rs
+++ b/bpf/aya-bpf/src/lib.rs
@@ -5,6 +5,7 @@
 pub use aya_bpf_bindings::bindings;
 
 mod args;
+pub use args::PtRegs;
 pub mod helpers;
 pub mod maps;
 pub mod programs;