From a89fa33f4645bdda0052926aacd11f09e6cd513f Mon Sep 17 00:00:00 2001
From: pdliyan <pdliyan@didiglobal.com>
Date: Tue, 4 Jul 2023 20:13:16 +0800
Subject: [PATCH] try add sarg.

---
 bpf/aya-bpf/src/args.rs | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/bpf/aya-bpf/src/args.rs b/bpf/aya-bpf/src/args.rs
index fea024eb..48a4d53d 100644
--- a/bpf/aya-bpf/src/args.rs
+++ b/bpf/aya-bpf/src/args.rs
@@ -1,3 +1,6 @@
+use aya_bpf_bindings::bindings::__u64;
+use aya_bpf_cty::c_int;
+
 use crate::{cty::c_void, helpers::bpf_probe_read};
 
 // aarch64 uses user_pt_regs instead of pt_regs
@@ -93,6 +96,10 @@ pub trait FromPtRegs: Sized {
     /// 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 `n`th stack argument of a pt_regs context where `n`
+    /// starts at 0 and increases by 1 for each successive argument.
+    fn from_stack_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>;
 }
@@ -111,6 +118,11 @@ impl<T> FromPtRegs for *const T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        let addr: __u64 = &ctx.sp + 8 * (n + 1) as __u64;
+        unsafe { bpf_probe_read(addr as *const c_int).map(|v| v as *const _).ok() }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *const _).ok() }
     }
@@ -126,6 +138,11 @@ impl<T> FromPtRegs for *const T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        let addr: __u64 = &ctx.sp + 8 * (n + 1) as __u64;
+        unsafe { bpf_probe_read(addr as *const c_int).map(|v| v as *const _).ok() }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *const _).ok() }
     }
@@ -141,6 +158,11 @@ impl<T> FromPtRegs for *const T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        let addr: __u64 = &ctx.sp + 8 * (n + 1) as __u64;
+        unsafe { bpf_probe_read(addr as *const c_int).map(|v| v as *const _).ok() }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *const _).ok() }
     }
@@ -160,6 +182,11 @@ impl<T> FromPtRegs for *mut T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        let addr: __u64 = &ctx.sp + 8 * (n + 1) as __u64;
+        unsafe { bpf_probe_read(addr as *const c_int).map(|v| v as *mut _).ok() }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *mut _).ok() }
     }
@@ -175,6 +202,11 @@ impl<T> FromPtRegs for *mut T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        let addr: __u64 = &ctx.sp + 8 * (n + 1) as __u64;
+        unsafe { bpf_probe_read(addr as *const c_int).map(|v| v as *mut _).ok() }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *mut _).ok() }
     }
@@ -190,6 +222,11 @@ impl<T> FromPtRegs for *mut T {
         }
     }
 
+    fn from_stack_argument(ctx: &pt_regs, n: usize) -> Option<Self> {
+        let addr: __u64 = &ctx.sp + 8 * (n + 1) as __u64;
+        unsafe { bpf_probe_read(addr as *const c_int).map(|v| v as *mut _).ok() }
+    }
+
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *mut _).ok() }
     }