From 2d38b23b99cd259f7a249f4c63b12da909c67015 Mon Sep 17 00:00:00 2001
From: tyrone-wu <wudevelops@gmail.com>
Date: Tue, 26 Mar 2024 01:45:52 +0000
Subject: [PATCH] aya-ebpf/programs: moved ret from ProbeContext into new
 RetProbeContext

Created retprobe.rs to hold RetProbeContext and moved the ret from
ProbeContext in probe.rs into RetProbeContext. Now, only kprobe (which
uses ProbeContext) can access args, and kretprobe (which uses
RetProbeContext) can access ret.

Fixes: #700
---
 ebpf/aya-ebpf/src/programs/mod.rs      |  2 ++
 ebpf/aya-ebpf/src/programs/probe.rs    | 19 -----------
 ebpf/aya-ebpf/src/programs/retprobe.rs | 46 ++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 19 deletions(-)
 create mode 100644 ebpf/aya-ebpf/src/programs/retprobe.rs

diff --git a/ebpf/aya-ebpf/src/programs/mod.rs b/ebpf/aya-ebpf/src/programs/mod.rs
index 498b0d1c..557e688f 100644
--- a/ebpf/aya-ebpf/src/programs/mod.rs
+++ b/ebpf/aya-ebpf/src/programs/mod.rs
@@ -5,6 +5,7 @@ pub mod lsm;
 pub mod perf_event;
 pub mod probe;
 pub mod raw_tracepoint;
+pub mod retprobe;
 pub mod sk_buff;
 pub mod sk_lookup;
 pub mod sk_msg;
@@ -25,6 +26,7 @@ pub use lsm::LsmContext;
 pub use perf_event::PerfEventContext;
 pub use probe::ProbeContext;
 pub use raw_tracepoint::RawTracePointContext;
+pub use retprobe::RetProbeContext;
 pub use sk_buff::SkBuffContext;
 pub use sk_lookup::SkLookupContext;
 pub use sk_msg::SkMsgContext;
diff --git a/ebpf/aya-ebpf/src/programs/probe.rs b/ebpf/aya-ebpf/src/programs/probe.rs
index 330ac0ae..3f9b965b 100644
--- a/ebpf/aya-ebpf/src/programs/probe.rs
+++ b/ebpf/aya-ebpf/src/programs/probe.rs
@@ -43,25 +43,6 @@ impl ProbeContext {
     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_ebpf::{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 EbpfContext for ProbeContext {
diff --git a/ebpf/aya-ebpf/src/programs/retprobe.rs b/ebpf/aya-ebpf/src/programs/retprobe.rs
new file mode 100644
index 00000000..2cff9b55
--- /dev/null
+++ b/ebpf/aya-ebpf/src/programs/retprobe.rs
@@ -0,0 +1,46 @@
+use core::ffi::c_void;
+
+#[cfg(not(any(bpf_target_arch = "aarch64", bpf_target_arch = "riscv64")))]
+use crate::bindings::pt_regs;
+#[cfg(bpf_target_arch = "aarch64")]
+use crate::bindings::user_pt_regs as pt_regs;
+#[cfg(bpf_target_arch = "riscv64")]
+use crate::bindings::user_regs_struct as pt_regs;
+use crate::{args::FromPtRegs, EbpfContext};
+
+pub struct RetProbeContext {
+    pub regs: *mut pt_regs,
+}
+
+impl RetProbeContext {
+    pub fn new(ctx: *mut c_void) -> RetProbeContext {
+        RetProbeContext {
+            regs: ctx as *mut pt_regs,
+        }
+    }
+
+    /// Returns the return value of the probed function.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # #![allow(dead_code)]
+    /// # use aya_ebpf::{programs::RetProbeContext, cty::c_int};
+    /// unsafe fn try_kretprobe_try_to_wake_up(ctx: RetProbeContext) -> 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 EbpfContext for RetProbeContext {
+    fn as_ptr(&self) -> *mut c_void {
+        self.regs as *mut c_void
+    }
+}