aya-bpf-macros: Add sleepable for sleepable programs

This extends args.rs to permit boolean arguments, where a value
is true if an ident is present in the macro attributes.

This approach is used to implement sleepable for all sleepable
program types. For LSM, this was previously implemented as
`sleepable = true`.

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>

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| {
+        let args = Punctuated::<Arg, Token![,]>::parse_terminated_with(input, |input| {
-            Ok(NameValue {
+            let ident = input.parse::<Ident>()?;
-                name: input.parse()?,
+            let lookahead = input.lookahead1();
-                _eq: input.parse()?,
+            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,55 @@ impl Parse for Args {
     }
 }
 
-pub(crate) fn pop_arg(args: &mut Args, name: &str) -> Option<String> {
+pub(crate) fn pop_string_arg(args: &mut Args, name: &str) -> Option<String> {
-    match args.args.iter().position(|arg| arg.name == name) {
+    args.args
-        Some(index) => Some(args.args.remove(index).value.value()),
+        .iter()
-        None => None,
+        .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> {
+pub(crate) fn pop_bool_arg(args: &mut Args, name: &str) -> bool {
-    let value = match args.args.iter().position(|arg| arg.name == name) {
+    args.args
-        Some(index) => Some(args.args.remove(index).value.value()),
+        .iter()
-        None => None,
+        .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 pop_required_string_arg(args: &mut Args, name: &str) -> Result<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"),
+        })
+        .ok_or_else(|| {
+            let tokens = match args.args.first().unwrap() {
+                Arg::String(name_val) => &name_val.name,
+                Arg::Bool(ident) => ident,
             };
-    match value {
+            Error::new_spanned(tokens, "missing required argument")
-        Some(value) => Ok(value),
+        })
-        None => Err(Error::new_spanned(
-            args.args.first().unwrap().name.clone(),
-            format!("missing required argument `{}`", name),
-        )),
-    }
 }
 
 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>> {
+pub(crate) fn name_arg(args: &mut Args) -> Option<String> {
-    let name = pop_arg(args, "name");
+    pop_string_arg(args, "name")
-    Ok(name)
 }

aya-bpf-macros/src/btf_tracepoint.rs

@@ -2,43 +2,26 @@ 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_required_string_arg, Args};
 
 pub(crate) struct BtfTracePoint {
     item: ItemFn,
     function: String,
-    sleepable: bool,
 }
 
 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 function = pop_required_string_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'",
-                ));
-            }
-        }
         err_on_unknown_args(&args)?;
-        Ok(BtfTracePoint {
+        Ok(BtfTracePoint { item, function })
-            item,
-            function,
-            sleepable,
-        })
     }
 
     pub(crate) fn expand(&self) -> Result<TokenStream> {
-        let section_prefix = if self.sleepable { "tp_btf.s" } else { "tp_btf" };
+        let section_name: Cow<'_, _> = format!("tp_btf/{}", self.function).into();
-        let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into();
         let fn_vis = &self.item.vis;
         let fn_name = self.item.sig.ident.clone();
         let item = &self.item;

aya-bpf-macros/src/fentry.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_required_arg};
+use crate::args::{err_on_unknown_args, pop_bool_arg, pop_required_string_arg};
 
 pub(crate) struct FEntry {
     item: ItemFn,
@@ -20,12 +20,13 @@ impl FEntry {
         }
         let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let function = pop_required_arg(&mut args, "function")?;
+        let function = pop_required_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,
         })
     }
 
@@ -81,4 +82,34 @@ mod tests {
         };
         assert_eq!(expected.to_string(), expanded.to_string());
     }
+
+    #[test]
+    fn test_fentry_sleepable() {
+        let prog = FEntry::parse(
+            parse_quote! {
+                function = "sys_clone",
+                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/sys_clone"]
+            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

@@ -5,7 +5,7 @@ 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_required_string_arg};
 
 pub(crate) struct FExit {
     item: ItemFn,
@@ -20,12 +20,13 @@ impl FExit {
         }
         let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let function = pop_required_arg(&mut args, "function")?;
+        let function = pop_required_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,
         })
     }
 
@@ -81,4 +82,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)]
@@ -38,8 +38,8 @@ impl KProbe {
 
         if !attrs.is_empty() {
             let mut args = syn::parse2(attrs)?;
-            function = pop_arg(&mut args, "function");
+            function = pop_string_arg(&mut args, "function");
-            offset = pop_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
+            offset = pop_string_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
             err_on_unknown_args(&args)?;
         }
         let item = syn::parse2(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

@@ -5,7 +5,7 @@ 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_required_string_arg};
 
 pub(crate) struct Lsm {
     item: ItemFn,
@@ -20,12 +20,13 @@ impl Lsm {
         }
         let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let hook = crate::args::pop_required_arg(&mut args, "hook")?;
+        let hook = pop_required_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,
         })
     }
 
@@ -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

@@ -16,7 +16,7 @@ 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 name = name_arg(&mut args).unwrap_or_else(|| item.ident.to_string());
         Ok(Map { item, name })
     }
 

aya-bpf-macros/src/raw_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_required_arg};
+use crate::args::{err_on_unknown_args, pop_required_string_arg};
 
 pub(crate) struct RawTracePoint {
     item: ItemFn,
@@ -19,7 +19,7 @@ impl RawTracePoint {
         }
         let mut args = syn::parse2(attrs)?;
         let item = syn::parse2(item)?;
-        let tracepoint = pop_required_arg(&mut args, "tracepoint")?;
+        let tracepoint = pop_required_string_arg(&mut args, "tracepoint")?;
         err_on_unknown_args(&args)?;
         Ok(RawTracePoint { item, tracepoint })
     }

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,
@@ -17,8 +17,8 @@ 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 name = pop_string_arg(&mut args, "name");
-        let category = pop_arg(&mut args, "category");
+        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,6 +30,7 @@ pub(crate) struct UProbe {
     function: Option<String>,
     offset: Option<u64>,
     item: ItemFn,
+    sleepable: bool,
 }
 
 impl UProbe {
@@ -37,11 +38,13 @@ impl UProbe {
         let mut path = None;
         let mut function = None;
         let mut offset = None;
+        let mut sleepable = false;
         if !attrs.is_empty() {
             let mut args = syn::parse2(attrs)?;
-            path = pop_arg(&mut args, "path");
+            path = pop_string_arg(&mut args, "path");
-            function = pop_arg(&mut args, "function");
+            function = pop_string_arg(&mut args, "function");
-            offset = pop_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
+            offset = pop_string_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap());
+            sleepable = pop_bool_arg(&mut args, "sleepable");
             err_on_unknown_args(&args)?;
         }
 
@@ -52,10 +55,16 @@ impl UProbe {
             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 +75,7 @@ impl UProbe {
             }
             format!(
                 "{}/{}:{}+{}",
-                self.kind,
+                prefix,
                 path,
                 self.function.as_ref().unwrap(),
                 self.offset.unwrap()
@@ -80,9 +89,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 +144,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(