aya-bpf-macros: Cleanup and remove required_args

This commit cleans up the recently refactored aya-bpf-macros
to be a little easier to maintain:

Cosmetic changes:

1. The flow of the parse function is the same in each file
2. Useless if !attrs.is_empty() guards were removed
3. The compile error when cgroup_sock programs have an invalid
   attach_type was in the wrong location

Functional changes:

1. Remove pop_required_arg. All args are optional and making them
   required was a mistake on my part. We may revisit this later.
2. Rename pop_arg to pop_string_arg
3. Implement pop_bool_arg to allow for single idents to be mixed with
   our key/value pair syntax in macro attributes. This is used for
   `sleepable` and `frags` in XDP programs.

Signed-off-by: Dave Tucker <[email protected]>

aya-bpf-macros/src/args.rs

@@ -1,28 +1,39 @@
 use syn::{
     parse::{Parse, ParseStream},
     punctuated::{Pair, Punctuated},
-    token::Eq,
     Error, Ident, LitStr, Result, Token,
 };
 
 pub(crate) struct NameValue {
     name: Ident,
-    _eq: Eq,
     value: LitStr,
 }
 
+pub(crate) enum Arg {
+    String(NameValue),
+    Bool(Ident),
+}
+
 pub(crate) struct Args {
-    pub(crate) args: Vec<NameValue>,
+    pub(crate) args: Vec<Arg>,
 }
 
 impl Parse for Args {
     fn parse(input: ParseStream) -> Result<Args> {
-        let args = Punctuated::<NameValue, Token![,]>::parse_terminated_with(input, |input| {
-            Ok(NameValue {
-                name: input.parse()?,
-                _eq: input.parse()?,
-                value: input.parse()?,
-            })
+        let args = Punctuated::<Arg, Token![,]>::parse_terminated_with(input, |input| {
+            let ident = input.parse::<Ident>()?;
+            let lookahead = input.lookahead1();
+            if input.is_empty() || lookahead.peek(Token![,]) {
+                Ok(Arg::Bool(ident))
+            } else if lookahead.peek(Token![=]) {
+                let _: Token![=] = input.parse()?;
+                Ok(Arg::String(NameValue {
+                    name: ident,
+                    value: input.parse()?,
+                }))
+            } else {
+                Err(lookahead.error())
+            }
         })?
         .into_pairs()
         .map(|pair| match pair {
@@ -35,35 +46,38 @@ impl Parse for Args {
     }
 }
 
-pub(crate) fn pop_arg(args: &mut Args, name: &str) -> Option<String> {
-    match args.args.iter().position(|arg| arg.name == name) {
-        Some(index) => Some(args.args.remove(index).value.value()),
-        None => None,
-    }
+pub(crate) fn pop_string_arg(args: &mut Args, name: &str) -> Option<String> {
+    args.args
+        .iter()
+        .position(|arg| matches!(arg, Arg::String(name_val) if name_val.name == name))
+        .map(|index| match args.args.remove(index) {
+            Arg::String(v) => v.value.value(),
+            _ => panic!("impossible variant"),
+        })
 }
 
-pub(crate) fn pop_required_arg(args: &mut Args, name: &str) -> Result<String> {
-    let value = match args.args.iter().position(|arg| arg.name == name) {
-        Some(index) => Some(args.args.remove(index).value.value()),
-        None => None,
-    };
-    match value {
-        Some(value) => Ok(value),
-        None => Err(Error::new_spanned(
-            args.args.first().unwrap().name.clone(),
-            format!("missing required argument `{}`", name),
-        )),
-    }
+pub(crate) fn pop_bool_arg(args: &mut Args, name: &str) -> bool {
+    args.args
+        .iter()
+        .position(|arg| matches!(arg, Arg::Bool(ident) if ident == name))
+        .map(|index| match args.args.remove(index) {
+            Arg::Bool(ident) => ident,
+            _ => panic!("impossible variant"),
+        })
+        .is_some()
 }
 
 pub(crate) fn err_on_unknown_args(args: &Args) -> Result<()> {
     if let Some(arg) = args.args.get(0) {
-        return Err(Error::new_spanned(&arg.name, "invalid argument"));
+        let tokens = match arg {
+            Arg::String(name_val) => name_val.name.clone(),
+            Arg::Bool(ident) => ident.clone(),
+        };
+        return Err(Error::new_spanned(tokens, "invalid argument"));
     }
     Ok(())
 }
 
-pub(crate) fn name_arg(args: &mut Args) -> Result<Option<String>> {
-    let name = pop_arg(args, "name");
-    Ok(name)
+pub(crate) fn name_arg(args: &mut Args) -> Option<String> {
+    pop_string_arg(args, "name")
 }

aya-bpf-macros/src/btf_tracepoint.rs

@@ -2,43 +2,31 @@ use std::borrow::Cow;
 
 use proc_macro2::TokenStream;
 use quote::quote;
-use syn::{Error, ItemFn, Result};
+use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_arg, pop_required_arg, Args};
+use crate::args::{err_on_unknown_args, pop_string_arg, Args};
 
 pub(crate) struct BtfTracePoint {
     item: ItemFn,
-    function: String,
-    sleepable: bool,
+    function: Option<String>,
 }
 
 impl BtfTracePoint {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
-        let mut args: Args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let function = pop_required_arg(&mut args, "function")?;
-        let mut sleepable = false;
-        if let Some(s) = pop_arg(&mut args, "sleepable") {
-            if let Ok(m) = s.parse() {
-                sleepable = m
-            } else {
-                return Err(Error::new_spanned(
-                    s,
-                    "invalid value. should be 'true' or 'false'",
-                ));
-            }
-        }
+        let mut args: Args = syn::parse2(attrs)?;
+        let function = pop_string_arg(&mut args, "function");
         err_on_unknown_args(&args)?;
-        Ok(BtfTracePoint {
-            item,
-            function,
-            sleepable,
-        })
+
+        Ok(BtfTracePoint { item, function })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
-        let section_prefix = if self.sleepable { "tp_btf.s" } else { "tp_btf" };
-        let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into();
+        let section_name: Cow<'_, _> = if let Some(function) = &self.function {
+            format!("tp_btf/{}", function).into()
+        } else {
+            "tp_btf".into()
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
@@ -62,6 +50,33 @@ mod tests {
 
     #[test]
     fn test_btf_tracepoint() {
+        let prog = BtfTracePoint::parse(
+            parse_quote!(),
+            parse_quote!(
+                fn foo(ctx: BtfTracePointContext) -> i32 {
+                    0
+                }
+            ),
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote!(
+            #[no_mangle]
+            #[link_section = "tp_btf"]
+            fn foo(ctx: *mut ::core::ffi::c_void) -> i32 {
+                let _ = foo(::aya_bpf::programs::BtfTracePointContext::new(ctx));
+                return 0;
+
+                fn foo(ctx: BtfTracePointContext) -> i32 {
+                    0
+                }
+            }
+        );
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+
+    #[test]
+    fn test_btf_tracepoint_with_function() {
         let prog = BtfTracePoint::parse(
             parse_quote!(function = "some_func"),
             parse_quote!(

aya-bpf-macros/src/cgroup_skb.rs

@@ -18,11 +18,10 @@ impl CgroupSkb {
             let ident: Ident = syn::parse2(attrs)?;
             match ident.to_string().as_str() {
                 "ingress" | "egress" => (),
-                _ => abort!(attach_type, "invalid attach type"),
+                _ => abort!(ident, "invalid attach type"),
             }
             attach_type = Some(ident.to_string());
         }
-
         Ok(CgroupSkb { item, attach_type })
     }
 

aya-bpf-macros/src/cgroup_sock.rs

@@ -12,6 +12,9 @@ pub(crate) struct CgroupSock {
 
 impl CgroupSock {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<CgroupSock> {
+        if attrs.is_empty() {
+            abort!(attrs, "missing attach type")
+        }
         let item: ItemFn = syn::parse2(item)?;
         let attach_type: Ident = syn::parse2(attrs)?;
         match attach_type.to_string().as_str() {

aya-bpf-macros/src/cgroup_sock_addr.rs

@@ -15,6 +15,7 @@ impl CgroupSockAddr {
         if attrs.is_empty() {
             abort!(attrs, "missing attach type")
         }
+        let item = syn::parse2(item)?;
         let attach_type: Ident = syn::parse2(attrs)?;
         match attach_type.to_string().as_str() {
             "connect4" | "connect6" | "bind4" | "bind6" | "getpeername4" | "getpeername6"
@@ -22,7 +23,6 @@ impl CgroupSockAddr {
             | "recvmsg6" => (),
             _ => abort!(attach_type, "invalid attach type"),
         }
-        let item = syn::parse2(item)?;
         Ok(CgroupSockAddr {
             item,
             attach_type: attach_type.to_string(),

aya-bpf-macros/src/cgroup_sockopt.rs

@@ -15,12 +15,12 @@ impl CgroupSockopt {
         if attrs.is_empty() {
             abort!(attrs, "expected attach type");
         }
+        let item = syn::parse2(item)?;
         let attach_type: Ident = syn::parse2(attrs)?;
         match attach_type.to_string().as_str() {
             "getsockopt" | "setsockopt" => (),
             _ => abort!(attach_type, "invalid attach type"),
         }
-        let item = syn::parse2(item)?;
         Ok(CgroupSockopt {
             item,
             attach_type: attach_type.to_string(),

aya-bpf-macros/src/fentry.rs

@@ -1,37 +1,38 @@
 use std::borrow::Cow;
 
 use proc_macro2::TokenStream;
-use proc_macro_error::abort;
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_required_arg};
+use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg};
 
 pub(crate) struct FEntry {
     item: ItemFn,
-    function: String,
+    function: Option<String>,
     sleepable: bool,
 }
 
 impl FEntry {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<FEntry> {
-        if attrs.is_empty() {
-            abort!(attrs, "missing function name");
-        }
-        let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let function = pop_required_arg(&mut args, "function")?;
+        let mut args = syn::parse2(attrs)?;
+        let function = pop_string_arg(&mut args, "function");
+        let sleepable = pop_bool_arg(&mut args, "sleepable");
         err_on_unknown_args(&args)?;
         Ok(FEntry {
             item,
             function,
-            sleepable: false,
+            sleepable,
         })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
         let section_prefix = if self.sleepable { "fentry.s" } else { "fentry" };
-        let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into();
+        let section_name: Cow<'_, _> = if let Some(function) = &self.function {
+            format!("{}/{}", section_prefix, function).into()
+        } else {
+            section_prefix.into()
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
@@ -55,6 +56,33 @@ mod tests {
 
     #[test]
     fn test_fentry() {
+        let prog = FEntry::parse(
+            parse_quote! {},
+            parse_quote! {
+                fn sys_clone(ctx: &mut aya_bpf::programs::FEntryContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote! {
+            #[no_mangle]
+            #[link_section = "fentry"]
+            fn sys_clone(ctx: *mut ::core::ffi::c_void) -> i32 {
+                let _ = sys_clone(::aya_bpf::programs::FEntryContext::new(ctx));
+                return 0;
+
+                fn sys_clone(ctx: &mut aya_bpf::programs::FEntryContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+
+    #[test]
+    fn test_fentry_with_function() {
         let prog = FEntry::parse(
             parse_quote! {
                 function = "sys_clone"
@@ -81,4 +109,33 @@ mod tests {
         };
         assert_eq!(expected.to_string(), expanded.to_string());
     }
+
+    #[test]
+    fn test_fentry_sleepable() {
+        let prog = FEntry::parse(
+            parse_quote! {
+                sleepable
+            },
+            parse_quote! {
+                fn sys_clone(ctx: &mut aya_bpf::programs::FEntryContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote! {
+            #[no_mangle]
+            #[link_section = "fentry.s"]
+            fn sys_clone(ctx: *mut ::core::ffi::c_void) -> i32 {
+                let _ = sys_clone(::aya_bpf::programs::FEntryContext::new(ctx));
+                return 0;
+
+                fn sys_clone(ctx: &mut aya_bpf::programs::FEntryContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
 }

aya-bpf-macros/src/fexit.rs

@@ -1,37 +1,38 @@
 use std::borrow::Cow;
 
 use proc_macro2::TokenStream;
-use proc_macro_error::abort;
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_required_arg};
+use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg};
 
 pub(crate) struct FExit {
     item: ItemFn,
-    function: String,
+    function: Option<String>,
     sleepable: bool,
 }
 
 impl FExit {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<FExit> {
-        if attrs.is_empty() {
-            abort!(attrs, "missing function name");
-        }
-        let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let function = pop_required_arg(&mut args, "function")?;
+        let mut args = syn::parse2(attrs)?;
+        let function = pop_string_arg(&mut args, "function");
+        let sleepable = pop_bool_arg(&mut args, "sleepable");
         err_on_unknown_args(&args)?;
         Ok(FExit {
             item,
             function,
-            sleepable: false,
+            sleepable,
         })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
         let section_prefix = if self.sleepable { "fexit.s" } else { "fexit" };
-        let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into();
+        let section_name: Cow<'_, _> = if let Some(function) = &self.function {
+            format!("{}/{}", section_prefix, function).into()
+        } else {
+            section_prefix.into()
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
@@ -55,6 +56,33 @@ mod tests {
 
     #[test]
     fn test_fexit() {
+        let prog = FExit::parse(
+            parse_quote! {},
+            parse_quote! {
+                fn sys_clone(ctx: &mut FExitContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote! {
+            #[no_mangle]
+            #[link_section = "fexit"]
+            fn sys_clone(ctx: *mut ::core::ffi::c_void) -> i32 {
+                let _ = sys_clone(::aya_bpf::programs::FExitContext::new(ctx));
+                return 0;
+
+                fn sys_clone(ctx: &mut FExitContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+
+    #[test]
+    fn test_fexit_with_function() {
         let prog = FExit::parse(
             parse_quote! {
                 function = "sys_clone"
@@ -81,4 +109,33 @@ mod tests {
         };
         assert_eq!(expected.to_string(), expanded.to_string());
     }
+
+    #[test]
+    fn test_fexit_sleepable() {
+        let prog = FExit::parse(
+            parse_quote! {
+                function = "sys_clone", sleepable
+            },
+            parse_quote! {
+                fn sys_clone(ctx: &mut FExitContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote! {
+            #[no_mangle]
+            #[link_section = "fexit.s/sys_clone"]
+            fn sys_clone(ctx: *mut ::core::ffi::c_void) -> i32 {
+                let _ = sys_clone(::aya_bpf::programs::FExitContext::new(ctx));
+                return 0;
+
+                fn sys_clone(ctx: &mut FExitContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
 }

aya-bpf-macros/src/kprobe.rs

@@ -5,7 +5,7 @@ use proc_macro2::TokenStream;
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_arg};
+use crate::args::{err_on_unknown_args, pop_string_arg};
 
 #[allow(clippy::enum_variant_names)]
 #[derive(Debug, Copy, Clone)]
@@ -33,16 +33,12 @@ pub(crate) struct KProbe {
 
 impl KProbe {
     pub(crate) fn parse(kind: KProbeKind, attrs: TokenStream, item: TokenStream) -> Result<KProbe> {
-        let mut function = None;
-        let mut offset = None;
-
-        if !attrs.is_empty() {
-            let mut args = syn::parse2(attrs)?;
-            function = pop_arg(&mut args, "function");
-            offset = pop_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
-            err_on_unknown_args(&args)?;
-        }
         let item = syn::parse2(item)?;
+        let mut args = syn::parse2(attrs)?;
+        let function = pop_string_arg(&mut args, "function");
+        let offset = pop_string_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
+        err_on_unknown_args(&args)?;
+
         Ok(KProbe {
             kind,
             item,

aya-bpf-macros/src/lib.rs

@@ -309,7 +309,7 @@ pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
 /// Used to implement security policy and audit logging.
 ///
 /// The hook name is the first argument to the macro.
-/// You may also provide `sleepable = true` to mark the program as sleepable.
+/// You may also provide `sleepable` to mark the program as sleepable.
 /// Arguments should be comma separated.
 ///
 /// LSM probes can be attached to the kernel's security hooks to implement mandatory

aya-bpf-macros/src/lsm.rs

@@ -1,37 +1,38 @@
 use std::borrow::Cow;
 
 use proc_macro2::TokenStream;
-use proc_macro_error::abort;
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::err_on_unknown_args;
+use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg};
 
 pub(crate) struct Lsm {
     item: ItemFn,
-    hook: String,
+    hook: Option<String>,
     sleepable: bool,
 }
 
 impl Lsm {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Lsm> {
-        if attrs.is_empty() {
-            abort!(attrs, "missing hook name");
-        }
-        let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let hook = crate::args::pop_required_arg(&mut args, "hook")?;
+        let mut args = syn::parse2(attrs)?;
+        let hook = pop_string_arg(&mut args, "hook");
+        let sleepable = pop_bool_arg(&mut args, "sleepable");
         err_on_unknown_args(&args)?;
         Ok(Lsm {
             item,
             hook,
-            sleepable: false,
+            sleepable,
         })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
         let section_prefix = if self.sleepable { "lsm.s" } else { "lsm" };
-        let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.hook).into();
+        let section_name: Cow<'_, _> = if let Some(hook) = &self.hook {
+            format!("{}/{}", section_prefix, hook).into()
+        } else {
+            section_prefix.into()
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;
@@ -55,6 +56,35 @@ mod tests {
     use super::*;
     use syn::parse_quote;
 
+    #[test]
+    fn test_lsm_sleepable() {
+        let prog = Lsm::parse(
+            parse_quote! {
+                sleepable,
+                hook = "bprm_committed_creds"
+            },
+            parse_quote! {
+                fn bprm_committed_creds(ctx: &mut ::aya_bpf::programs::LsmContext) -> i32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        let expanded = prog.expand().unwrap();
+        let expected = quote! {
+            #[no_mangle]
+            #[link_section = "lsm.s/bprm_committed_creds"]
+            fn bprm_committed_creds(ctx: *mut ::core::ffi::c_void) -> i32 {
+                return bprm_committed_creds(::aya_bpf::programs::LsmContext::new(ctx));
+
+                fn bprm_committed_creds(ctx: &mut ::aya_bpf::programs::LsmContext) -> i32 {
+                    0
+                }
+            }
+        };
+        assert_eq!(expected.to_string(), expanded.to_string());
+    }
+
     #[test]
     fn test_lsm() {
         let prog = Lsm::parse(

aya-bpf-macros/src/map.rs

@@ -14,9 +14,9 @@ pub(crate) struct Map {
 
 impl Map {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Map> {
-        let mut args = syn::parse2(attrs)?;
         let item: ItemStatic = syn::parse2(item)?;
-        let name = name_arg(&mut args)?.unwrap_or_else(|| item.ident.to_string());
+        let mut args = syn::parse2(attrs)?;
+        let name = name_arg(&mut args).unwrap_or_else(|| item.ident.to_string());
         Ok(Map { item, name })
     }
 

aya-bpf-macros/src/raw_tracepoint.rs

@@ -1,31 +1,32 @@
 use std::borrow::Cow;
 
 use proc_macro2::TokenStream;
-use proc_macro_error::abort;
+
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_required_arg};
+use crate::args::{err_on_unknown_args, pop_string_arg};
 
 pub(crate) struct RawTracePoint {
     item: ItemFn,
-    tracepoint: String,
+    tracepoint: Option<String>,
 }
 
 impl RawTracePoint {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<RawTracePoint> {
-        if attrs.is_empty() {
-            abort!(attrs, "expected `tracepoint` attribute")
-        }
-        let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let tracepoint = pop_required_arg(&mut args, "tracepoint")?;
+        let mut args = syn::parse2(attrs)?;
+        let tracepoint = pop_string_arg(&mut args, "tracepoint");
         err_on_unknown_args(&args)?;
         Ok(RawTracePoint { item, tracepoint })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
-        let section_name: Cow<'_, _> = format!("raw_tp/{}", self.tracepoint).into();
+        let section_name: Cow<'_, _> = if let Some(tracepoint) = &self.tracepoint {
+            format!("raw_tp/{}", tracepoint).into()
+        } else {
+            "raw_tp".into()
+        };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;

aya-bpf-macros/src/tracepoint.rs

@@ -5,7 +5,7 @@ use proc_macro_error::abort;
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_arg};
+use crate::args::{err_on_unknown_args, pop_string_arg};
 
 pub(crate) struct TracePoint {
     item: ItemFn,
@@ -15,10 +15,10 @@ pub(crate) struct TracePoint {
 
 impl TracePoint {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<TracePoint> {
-        let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let name = pop_arg(&mut args, "name");
-        let category = pop_arg(&mut args, "category");
+        let mut args = syn::parse2(attrs)?;
+        let name = pop_string_arg(&mut args, "name");
+        let category = pop_string_arg(&mut args, "category");
         err_on_unknown_args(&args)?;
         Ok(TracePoint {
             item,

aya-bpf-macros/src/uprobe.rs

@@ -5,7 +5,7 @@ use proc_macro_error::abort;
 use quote::quote;
 use syn::{ItemFn, Result};
 
-use crate::args::{err_on_unknown_args, pop_arg};
+use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg};
 
 #[allow(clippy::enum_variant_names)]
 #[derive(Debug, Copy, Clone)]
@@ -30,32 +30,34 @@ pub(crate) struct UProbe {
     function: Option<String>,
     offset: Option<u64>,
     item: ItemFn,
+    sleepable: bool,
 }
 
 impl UProbe {
     pub(crate) fn parse(kind: UProbeKind, attrs: TokenStream, item: TokenStream) -> Result<UProbe> {
-        let mut path = None;
-        let mut function = None;
-        let mut offset = None;
-        if !attrs.is_empty() {
-            let mut args = syn::parse2(attrs)?;
-            path = pop_arg(&mut args, "path");
-            function = pop_arg(&mut args, "function");
-            offset = pop_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
-            err_on_unknown_args(&args)?;
-        }
-
         let item = syn::parse2(item)?;
+        let mut args = syn::parse2(attrs)?;
+        let path = pop_string_arg(&mut args, "path");
+        let function = pop_string_arg(&mut args, "function");
+        let offset = pop_string_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
+        let sleepable = pop_bool_arg(&mut args, "sleepable");
+        err_on_unknown_args(&args)?;
         Ok(UProbe {
             kind,
             item,
             path,
             function,
             offset,
+            sleepable,
         })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
+        let prefix = if self.sleepable {
+            format!("{}.s", self.kind)
+        } else {
+            format!("{}", self.kind)
+        };
         let section_name: Cow<'_, _> = if self.path.is_some() && self.offset.is_some() {
             if self.function.is_none() {
                 abort!(self.item.sig.ident, "expected `function` attribute");
@@ -66,7 +68,7 @@ impl UProbe {
             }
             format!(
                 "{}/{}:{}+{}",
-                self.kind,
+                prefix,
                 path,
                 self.function.as_ref().unwrap(),
                 self.offset.unwrap()
@@ -80,9 +82,9 @@ impl UProbe {
             if path.starts_with('/') {
                 path.remove(0);
             }
-            format!("{}/{}:{}", self.kind, path, self.function.as_ref().unwrap()).into()
+            format!("{}/{}:{}", prefix, path, self.function.as_ref().unwrap()).into()
         } else {
-            format!("{}", self.kind).into()
+            prefix.to_string().into()
         };
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
@@ -135,6 +137,36 @@ mod tests {
         );
     }
 
+    #[test]
+    fn uprobe_sleepable() {
+        let uprobe = UProbe::parse(
+            UProbeKind::UProbe,
+            parse_quote! {sleepable},
+            parse_quote! {
+                fn foo(ctx: ProbeContext) -> u32 {
+                    0
+                }
+            },
+        )
+        .unwrap();
+        assert_eq!(
+            uprobe.expand().unwrap().to_string(),
+            quote! {
+                #[no_mangle]
+                #[link_section = "uprobe.s"]
+                fn foo(ctx: *mut ::core::ffi::c_void) -> u32 {
+                    let _ = foo(::aya_bpf::programs::ProbeContext::new(ctx));
+                    return 0;
+
+                    fn foo(ctx: ProbeContext) -> u32 {
+                        0
+                    }
+                }
+            }
+            .to_string()
+        );
+    }
+
     #[test]
     fn uprobe_with_path() {
         let uprobe = UProbe::parse(

aya-bpf-macros/src/xdp.rs

@@ -1,9 +1,10 @@
 use std::borrow::Cow;
 
 use proc_macro2::TokenStream;
-use proc_macro_error::abort;
 use quote::quote;
-use syn::{Ident, ItemFn, Result};
+use syn::{ItemFn, Result};
+
+use crate::args::{err_on_unknown_args, pop_bool_arg, Args};
 
 pub(crate) struct Xdp {
     item: ItemFn,
@@ -12,15 +13,10 @@ pub(crate) struct Xdp {
 
 impl Xdp {
     pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Xdp> {
-        let mut frags = false;
         let item = syn::parse2(item)?;
-        if !attrs.is_empty() {
-            let ident: Ident = syn::parse2(attrs)?;
-            if ident != "frags" {
-                abort!(ident, "unexpected attribute");
-            }
-            frags = true;
-        }
+        let mut args: Args = syn::parse2(attrs)?;
+        let frags = pop_bool_arg(&mut args, "frags");
+        err_on_unknown_args(&args)?;
         Ok(Xdp { item, frags })
     }