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: aya-rs#700
astoycos · Apr 24, 2024 · b84ede1 · b84ede1
1 parent 2d38b23
commit b84ede1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/aya-ebpf-macros/src/kprobe.rs b/aya-ebpf-macros/src/kprobe.rs
@@ -60,14 +60,20 @@ 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;
         Ok(quote! {
             #[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
@@ -86,14 +86,20 @@ 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;
         Ok(quote! {
             #[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 {