Browse Source

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
tyrone-wu 11 months ago
parent
commit
b84ede10b9
2 changed files with 16 additions and 4 deletions
  1. 8 2
      aya-ebpf-macros/src/kprobe.rs
  2. 8 2
      aya-ebpf-macros/src/uprobe.rs

+ 8 - 2
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 {

+ 8 - 2
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 {