From 4acd996cb83d6f447f2717b1adcbb8c4ea69121f 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 | 26 ++++++++++++++++++++++++++
 bpf/aya-bpf/src/lib.rs  |  1 +
 2 files changed, 27 insertions(+)

diff --git a/bpf/aya-bpf/src/args.rs b/bpf/aya-bpf/src/args.rs
index 42bfad9f..c919576a 100644
--- a/bpf/aya-bpf/src/args.rs
+++ b/bpf/aya-bpf/src/args.rs
@@ -56,6 +56,32 @@ 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,
+}
+
+/// A portable wrapper around pt_regs and user_pt_regs.
+impl PtRegs {
+    pub fn new(regs: *mut pt_regs) -> Self {
+        PtRegs { regs }
+    }
+
+    /// Returns the value of the register used to pass arg `n`.
+    pub fn arg<T: FromPtRegs>(&self, n: usize) -> Option<T> {
+        T::from_argument(unsafe { &*self.regs }, n)
+    }
+
+    /// Returns the value of the register used to pass the return value.
+    pub fn ret<T: FromPtRegs>(&self) -> Option<T> {
+        T::from_retval(unsafe { &*self.regs })
+    }
+
+    /// Returns a pointer to the wrapped value.
+    pub fn as_ptr(&self) -> *mut pt_regs {
+        self.regs
+    }
+}
+
 /// 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;