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,14 +15,14 @@ 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"
             | "getsockname4" | "getsockname6" | "sendmsg4" | "sendmsg6" | "recvmsg4"
             | "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());
+    }
 }