From b84ede10b9c4813f221fade16b60d5ced4ecdc58 Mon Sep 17 00:00:00 2001
From: tyrone-wu <wudevelops@gmail.com>
Date: Tue, 26 Mar 2024 02:21:13 +0000
Subject: [PATCH] aya-ebpf/macros: separate probe to probe ctx & retprobe to
 retprobe ctx

Added logic in expand function in both kprobe.rs and uprobe.rs for valid
macros. Now, kprobe & uprobe proc macros only accept ProbeContext, and
kretprobe & uretprobe only accept RetProbeContext.

Ref: #700
---
 aya-ebpf-macros/src/kprobe.rs | 10 ++++++++--
 aya-ebpf-macros/src/uprobe.rs | 10 ++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/aya-ebpf-macros/src/kprobe.rs b/aya-ebpf-macros/src/kprobe.rs
index 01633917..87eedb26 100644
--- a/aya-ebpf-macros/src/kprobe.rs
+++ b/aya-ebpf-macros/src/kprobe.rs
@@ -60,6 +60,12 @@ impl KProbe {
         } else {
             format!("{}", self.kind).into()
         };
+
+        let probe_type = if section_name.as_ref().starts_with("kprobe") {
+            quote! { ProbeContext }
+        } else {
+            quote! { RetProbeContext }
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
@@ -67,7 +73,7 @@ impl KProbe {
             #[no_mangle]
             #[link_section = #section_name]
             #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
-                let _ = #fn_name(::aya_ebpf::programs::ProbeContext::new(ctx));
+                let _ = #fn_name(::aya_ebpf::programs::#probe_type::new(ctx));
                 return 0;
 
                 #item
@@ -195,7 +201,7 @@ mod tests {
                 #[no_mangle]
                 #[link_section = "kretprobe"]
                 fn foo(ctx: *mut ::core::ffi::c_void) -> u32 {
-                    let _ = foo(::aya_ebpf::programs::ProbeContext::new(ctx));
+                    let _ = foo(::aya_ebpf::programs::RetProbeContext::new(ctx));
                     return 0;
 
                     fn foo(ctx: ProbeContext) -> u32 {
diff --git a/aya-ebpf-macros/src/uprobe.rs b/aya-ebpf-macros/src/uprobe.rs
index 812279ee..5ce01239 100644
--- a/aya-ebpf-macros/src/uprobe.rs
+++ b/aya-ebpf-macros/src/uprobe.rs
@@ -86,6 +86,12 @@ impl UProbe {
         } else {
             prefix.to_string().into()
         };
+
+        let probe_type = if section_name.as_ref().starts_with("uprobe") {
+            quote! { ProbeContext }
+        } else {
+            quote! { RetProbeContext }
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
@@ -93,7 +99,7 @@ impl UProbe {
             #[no_mangle]
             #[link_section = #section_name]
             #fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
-                let _ = #fn_name(::aya_ebpf::programs::ProbeContext::new(ctx));
+                let _ = #fn_name(::aya_ebpf::programs::#probe_type::new(ctx));
                 return 0;
 
                 #item
@@ -251,7 +257,7 @@ mod tests {
                 #[no_mangle]
                 #[link_section = "uretprobe"]
                 fn foo(ctx: *mut ::core::ffi::c_void) -> u32 {
-                    let _ = foo(::aya_ebpf::programs::ProbeContext::new(ctx));
+                    let _ = foo(::aya_ebpf::programs::RetProbeContext::new(ctx));
                     return 0;
 
                     fn foo(ctx: ProbeContext) -> u32 {