From 54377b6140620dadf17a2c1f9a27334d44b67ab1 Mon Sep 17 00:00:00 2001
From: Michal Rostecki <mrostecki@opensuse.org>
Date: Tue, 23 Nov 2021 19:49:01 +0100
Subject: [PATCH] aya-bpf: Allow to convert probe arguments to mutable pointers

Before this change, arguments fetched with `arg` from `ProbeContext`
could be only fetched as const pointers. This change allows to get mut
pointers as well.

Signed-off-by: Michal Rostecki <mrostecki@opensuse.org>
---
 bpf/aya-bpf/src/args.rs | 49 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/bpf/aya-bpf/src/args.rs b/bpf/aya-bpf/src/args.rs
index 310a9611..42bfad9f 100644
--- a/bpf/aya-bpf/src/args.rs
+++ b/bpf/aya-bpf/src/args.rs
@@ -120,6 +120,55 @@ impl<T> FromPtRegs for *const T {
     }
 }
 
+#[cfg(bpf_target_arch = "x86_64")]
+impl<T> FromPtRegs for *mut 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 *mut 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 *mut 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) => {