Merge pull request #711 from dave-tucker/sleepable

Tidy up aya-bpf-macros (again)

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 })
     }
 

aya-obj/src/obj.rs

@@ -222,7 +222,6 @@ pub struct Function {
 /// Currently, the following section names are not supported yet:
 /// - `flow_dissector`: `BPF_PROG_TYPE_FLOW_DISSECTOR`
 /// - `ksyscall+` or `kretsyscall+`
-/// - `uprobe.s+` or `uretprobe.s+`
 /// - `usdt+`
 /// - `kprobe.multi+` or `kretprobe.multi+`: `BPF_TRACE_KPROBE_MULTI`
 /// - `lsm_cgroup+`
@@ -233,7 +232,6 @@ pub struct Function {
 /// - `syscall`
 /// - `struct_ops+`
 /// - `fmod_ret+`, `fmod_ret.s+`
-/// - `fentry.s+`, `fexit.s+`
 /// - `iter+`, `iter.s+`
 /// - `xdp.frags/cpumap`, `xdp/cpumap`
 /// - `xdp.frags/devmap`, `xdp/devmap`
@@ -248,9 +246,11 @@ pub enum ProgramSection {
     },
     UProbe {
         name: String,
+        sleepable: bool,
     },
     URetProbe {
         name: String,
+        sleepable: bool,
     },
     TracePoint {
         name: String,
@@ -315,9 +315,11 @@ pub enum ProgramSection {
     },
     FEntry {
         name: String,
+        sleepable: bool,
     },
     FExit {
         name: String,
+        sleepable: bool,
     },
     Extension {
         name: String,
@@ -340,8 +342,8 @@ impl ProgramSection {
         match self {
             ProgramSection::KRetProbe { name } => name,
             ProgramSection::KProbe { name } => name,
-            ProgramSection::UProbe { name } => name,
-            ProgramSection::URetProbe { name } => name,
+            ProgramSection::UProbe { name, .. } => name,
+            ProgramSection::URetProbe { name, .. } => name,
             ProgramSection::TracePoint { name } => name,
             ProgramSection::SocketFilter { name } => name,
             ProgramSection::Xdp { name, .. } => name,
@@ -360,9 +362,9 @@ impl ProgramSection {
             ProgramSection::PerfEvent { name } => name,
             ProgramSection::RawTracePoint { name } => name,
             ProgramSection::Lsm { name, .. } => name,
-            ProgramSection::BtfTracePoint { name } => name,
-            ProgramSection::FEntry { name } => name,
-            ProgramSection::FExit { name } => name,
+            ProgramSection::BtfTracePoint { name, .. } => name,
+            ProgramSection::FEntry { name, .. } => name,
+            ProgramSection::FExit { name, .. } => name,
             ProgramSection::Extension { name } => name,
             ProgramSection::SkLookup { name } => name,
             ProgramSection::CgroupSock { name, .. } => name,
@@ -388,8 +390,22 @@ impl FromStr for ProgramSection {
         Ok(match kind {
             "kprobe" => KProbe { name },
             "kretprobe" => KRetProbe { name },
-            "uprobe" => UProbe { name },
-            "uretprobe" => URetProbe { name },
+            "uprobe" => UProbe {
+                name,
+                sleepable: false,
+            },
+            "uprobe.s" => UProbe {
+                name,
+                sleepable: true,
+            },
+            "uretprobe" => URetProbe {
+                name,
+                sleepable: false,
+            },
+            "uretprobe.s" => URetProbe {
+                name,
+                sleepable: true,
+            },
             "xdp" => Xdp { name, frags: false },
             "xdp.frags" => Xdp { name, frags: true },
             "tp_btf" => BtfTracePoint { name },
@@ -551,8 +567,22 @@ impl FromStr for ProgramSection {
                 name,
                 sleepable: true,
             },
-            "fentry" => FEntry { name },
-            "fexit" => FExit { name },
+            "fentry" => FEntry {
+                name,
+                sleepable: false,
+            },
+            "fentry.s" => FEntry {
+                name,
+                sleepable: true,
+            },
+            "fexit" => FExit {
+                name,
+                sleepable: false,
+            },
+            "fexit.s" => FExit {
+                name,
+                sleepable: true,
+            },
             "freplace" => Extension { name },
             "sk_lookup" => SkLookup { name },
             _ => {
@@ -2018,6 +2048,81 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_parse_section_uprobe_sleepable() {
+        let mut obj = fake_obj();
+        fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);
+
+        assert_matches!(
+            obj.parse_section(fake_section(
+                BpfSectionKind::Program,
+                "uprobe.s/foo",
+                bytes_of(&fake_ins()),
+                None
+            )),
+            Ok(())
+        );
+        assert_matches!(
+            obj.programs.get("foo"),
+            Some(Program {
+                section: ProgramSection::UProbe {
+                    sleepable: true,
+                    ..
+                },
+                ..
+            })
+        );
+    }
+
+    #[test]
+    fn test_parse_section_uretprobe() {
+        let mut obj = fake_obj();
+        fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);
+
+        assert_matches!(
+            obj.parse_section(fake_section(
+                BpfSectionKind::Program,
+                "uretprobe/foo",
+                bytes_of(&fake_ins()),
+                None
+            )),
+            Ok(())
+        );
+        assert_matches!(
+            obj.programs.get("foo"),
+            Some(Program {
+                section: ProgramSection::URetProbe { .. },
+                ..
+            })
+        );
+    }
+
+    #[test]
+    fn test_parse_section_uretprobe_sleepable() {
+        let mut obj = fake_obj();
+        fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);
+
+        assert_matches!(
+            obj.parse_section(fake_section(
+                BpfSectionKind::Program,
+                "uretprobe.s/foo",
+                bytes_of(&fake_ins()),
+                None
+            )),
+            Ok(())
+        );
+        assert_matches!(
+            obj.programs.get("foo"),
+            Some(Program {
+                section: ProgramSection::URetProbe {
+                    sleepable: true,
+                    ..
+                },
+                ..
+            })
+        );
+    }
+
     #[test]
     fn test_parse_section_trace_point() {
         let mut obj = fake_obj();
@@ -2313,6 +2418,32 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_parse_section_fentry_sleepable() {
+        let mut obj = fake_obj();
+        fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);
+
+        assert_matches!(
+            obj.parse_section(fake_section(
+                BpfSectionKind::Program,
+                "fentry.s/foo",
+                bytes_of(&fake_ins()),
+                None
+            )),
+            Ok(())
+        );
+        assert_matches!(
+            obj.programs.get("foo"),
+            Some(Program {
+                section: ProgramSection::FEntry {
+                    sleepable: true,
+                    ..
+                },
+                ..
+            })
+        );
+    }
+
     #[test]
     fn test_parse_section_fexit() {
         let mut obj = fake_obj();
@@ -2336,6 +2467,32 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_parse_section_fexit_sleepable() {
+        let mut obj = fake_obj();
+        fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN);
+
+        assert_matches!(
+            obj.parse_section(fake_section(
+                BpfSectionKind::Program,
+                "fexit.s/foo",
+                bytes_of(&fake_ins()),
+                None
+            )),
+            Ok(())
+        );
+        assert_matches!(
+            obj.programs.get("foo"),
+            Some(Program {
+                section: ProgramSection::FExit {
+                    sleepable: true,
+                    ..
+                },
+                ..
+            })
+        );
+    }
+
     #[test]
     fn test_parse_section_cgroup_skb_ingress_unnamed() {
         let mut obj = fake_obj();

aya/src/bpf.rs

@@ -566,14 +566,28 @@ impl<'a> BpfLoader<'a> {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                             kind: ProbeKind::KRetProbe,
                         }),
-                        ProgramSection::UProbe { .. } => Program::UProbe(UProbe {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            kind: ProbeKind::UProbe,
-                        }),
-                        ProgramSection::URetProbe { .. } => Program::UProbe(UProbe {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            kind: ProbeKind::URetProbe,
-                        }),
+                        ProgramSection::UProbe { sleepable, name: _ } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::UProbe(UProbe {
+                                data,
+                                kind: ProbeKind::UProbe,
+                            })
+                        }
+                        ProgramSection::URetProbe { sleepable, name: _ } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::UProbe(UProbe {
+                                data,
+                                kind: ProbeKind::URetProbe,
+                            })
+                        }
                         ProgramSection::TracePoint { .. } => Program::TracePoint(TracePoint {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                         }),
@@ -653,7 +667,7 @@ impl<'a> BpfLoader<'a> {
                                 data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                             })
                         }
-                        ProgramSection::Lsm { sleepable, .. } => {
+                        ProgramSection::Lsm { sleepable, name: _ } => {
                             let mut data =
                                 ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
                             if *sleepable {
@@ -666,12 +680,22 @@ impl<'a> BpfLoader<'a> {
                                 data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                             })
                         }
-                        ProgramSection::FEntry { .. } => Program::FEntry(FEntry {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                        }),
-                        ProgramSection::FExit { .. } => Program::FExit(FExit {
-                            data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                        }),
+                        ProgramSection::FEntry { sleepable, name: _ } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::FEntry(FEntry { data })
+                        }
+                        ProgramSection::FExit { sleepable, name: _ } => {
+                            let mut data =
+                                ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level);
+                            if *sleepable {
+                                data.flags = BPF_F_SLEEPABLE;
+                            }
+                            Program::FExit(FExit { data })
+                        }
                         ProgramSection::Extension { .. } => Program::Extension(Extension {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
                         }),

xtask/public-api/aya-obj.txt

@@ -5837,8 +5837,10 @@ pub aya_obj::obj::ProgramSection::Extension
 pub aya_obj::obj::ProgramSection::Extension::name: alloc::string::String
 pub aya_obj::obj::ProgramSection::FEntry
 pub aya_obj::obj::ProgramSection::FEntry::name: alloc::string::String
+pub aya_obj::obj::ProgramSection::FEntry::sleepable: bool
 pub aya_obj::obj::ProgramSection::FExit
 pub aya_obj::obj::ProgramSection::FExit::name: alloc::string::String
+pub aya_obj::obj::ProgramSection::FExit::sleepable: bool
 pub aya_obj::obj::ProgramSection::KProbe
 pub aya_obj::obj::ProgramSection::KProbe::name: alloc::string::String
 pub aya_obj::obj::ProgramSection::KRetProbe
@@ -5870,8 +5872,10 @@ pub aya_obj::obj::ProgramSection::TracePoint
 pub aya_obj::obj::ProgramSection::TracePoint::name: alloc::string::String
 pub aya_obj::obj::ProgramSection::UProbe
 pub aya_obj::obj::ProgramSection::UProbe::name: alloc::string::String
+pub aya_obj::obj::ProgramSection::UProbe::sleepable: bool
 pub aya_obj::obj::ProgramSection::URetProbe
 pub aya_obj::obj::ProgramSection::URetProbe::name: alloc::string::String
+pub aya_obj::obj::ProgramSection::URetProbe::sleepable: bool
 pub aya_obj::obj::ProgramSection::Xdp
 pub aya_obj::obj::ProgramSection::Xdp::frags: bool
 pub aya_obj::obj::ProgramSection::Xdp::name: alloc::string::String
@@ -6585,8 +6589,10 @@ pub aya_obj::ProgramSection::Extension
 pub aya_obj::ProgramSection::Extension::name: alloc::string::String
 pub aya_obj::ProgramSection::FEntry
 pub aya_obj::ProgramSection::FEntry::name: alloc::string::String
+pub aya_obj::ProgramSection::FEntry::sleepable: bool
 pub aya_obj::ProgramSection::FExit
 pub aya_obj::ProgramSection::FExit::name: alloc::string::String
+pub aya_obj::ProgramSection::FExit::sleepable: bool
 pub aya_obj::ProgramSection::KProbe
 pub aya_obj::ProgramSection::KProbe::name: alloc::string::String
 pub aya_obj::ProgramSection::KRetProbe
@@ -6618,8 +6624,10 @@ pub aya_obj::ProgramSection::TracePoint
 pub aya_obj::ProgramSection::TracePoint::name: alloc::string::String
 pub aya_obj::ProgramSection::UProbe
 pub aya_obj::ProgramSection::UProbe::name: alloc::string::String
+pub aya_obj::ProgramSection::UProbe::sleepable: bool
 pub aya_obj::ProgramSection::URetProbe
 pub aya_obj::ProgramSection::URetProbe::name: alloc::string::String
+pub aya_obj::ProgramSection::URetProbe::sleepable: bool
 pub aya_obj::ProgramSection::Xdp
 pub aya_obj::ProgramSection::Xdp::frags: bool
 pub aya_obj::ProgramSection::Xdp::name: alloc::string::String