rex-macros: add macro support for uprobe

Signed-off-by: MinhPhan8803 <[email protected]>

Author: MinhPhan8803

rex-macros/src/kprobe.rs

@@ -1,9 +1,30 @@
+use std::fmt;
+
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::{parse2, ItemFn, Result};
 
 use crate::args::parse_string_args;
 
+#[allow(dead_code)]
+pub enum KprobeFlavor {
+    Kprobe,
+    Kretprobe,
+    Uprobe,
+    Uretprobe,
+}
+
+impl fmt::Display for KprobeFlavor {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            KprobeFlavor::Kprobe => write!(f, "kprobe"),
+            KprobeFlavor::Kretprobe => write!(f, "kretprobe"),
+            KprobeFlavor::Uprobe => write!(f, "uprobe"),
+            KprobeFlavor::Uretprobe => write!(f, "uretprobe"),
+        }
+    }
+}
+
 pub(crate) struct KProbe {
     function: Option<String>,
     item: ItemFn,
@@ -23,17 +44,17 @@ impl KProbe {
         Ok(KProbe { function, item })
     }
 
-    pub(crate) fn expand(&self) -> Result<TokenStream> {
+    pub(crate) fn expand(&self, flavor: KprobeFlavor) -> Result<TokenStream> {
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
         let function_name = format!("{}", fn_name);
         let prog_ident =
             format_ident!("PROG_{}", fn_name.to_string().to_uppercase());
 
         let attached_function = if self.function.is_some() {
-            format!("rex/kprobe/{}", self.function.as_ref().unwrap())
+            format!("rex/{}/{}", flavor, self.function.as_ref().unwrap())
         } else {
-            "rex/kprobe".to_string()
+            format!("rex/{}", flavor)
         };
 
         let function_body_tokens = quote! {

rex-macros/src/lib.rs

@@ -8,7 +8,7 @@ mod xdp;
 
 use std::borrow::Cow;
 
-use kprobe::KProbe;
+use kprobe::{KProbe, KprobeFlavor};
 use perf_event::PerfEvent;
 use proc_macro::TokenStream;
 use proc_macro_error::{abort, proc_macro_error};
@@ -47,7 +47,19 @@ pub fn rex_tc(attrs: TokenStream, item: TokenStream) -> TokenStream {
 pub fn rex_kprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
     match KProbe::parse(attrs.into(), item.into()) {
         Ok(prog) => prog
-            .expand()
+            .expand(KprobeFlavor::Kprobe)
+            .unwrap_or_else(|err| abort!(err.span(), "{}", err))
+            .into(),
+        Err(err) => abort!(err.span(), "{}", err),
+    }
+}
+
+#[proc_macro_error]
+#[proc_macro_attribute]
+pub fn rex_uprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    match KProbe::parse(attrs.into(), item.into()) {
+        Ok(prog) => prog
+            .expand(KprobeFlavor::Uprobe)
             .unwrap_or_else(|err| abort!(err.span(), "{}", err))
             .into(),
         Err(err) => abort!(err.span(), "{}", err),