diff --git a/aya-bpf-macros/src/args.rs b/aya-bpf-macros/src/args.rs index 8a83cbbb..b695053a 100644 --- a/aya-bpf-macros/src/args.rs +++ b/aya-bpf-macros/src/args.rs @@ -67,23 +67,6 @@ pub(crate) fn pop_bool_arg(args: &mut Args, name: &str) -> bool { .is_some() } -pub(crate) fn pop_required_string_arg(args: &mut Args, name: &str) -> Result { - 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, - }; - Error::new_spanned(tokens, "missing required argument") - }) -} - pub(crate) fn err_on_unknown_args(args: &Args) -> Result<()> { if let Some(arg) = args.args.get(0) { let tokens = match arg { diff --git a/aya-bpf-macros/src/btf_tracepoint.rs b/aya-bpf-macros/src/btf_tracepoint.rs index 62c6ab5e..da1bc443 100644 --- a/aya-bpf-macros/src/btf_tracepoint.rs +++ b/aya-bpf-macros/src/btf_tracepoint.rs @@ -4,24 +4,29 @@ use proc_macro2::TokenStream; use quote::quote; use syn::{ItemFn, Result}; -use crate::args::{err_on_unknown_args, pop_required_string_arg, Args}; +use crate::args::{err_on_unknown_args, pop_string_arg, Args}; pub(crate) struct BtfTracePoint { item: ItemFn, - function: String, + function: Option, } impl BtfTracePoint { pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { - let mut args: Args = syn::parse2(attrs)?; let item = syn::parse2(item)?; - let function = pop_required_string_arg(&mut args, "function")?; + let mut args: Args = syn::parse2(attrs)?; + let function = pop_string_arg(&mut args, "function"); err_on_unknown_args(&args)?; + Ok(BtfTracePoint { item, function }) } pub(crate) fn expand(&self) -> Result { - let section_name: Cow<'_, _> = format!("tp_btf/{}", self.function).into(); + let section_name: Cow<'_, _> = if self.function.is_none() { + "tp_btf".into() + } else { + format!("tp_btf/{}", self.function.as_ref().unwrap()).into() + }; let fn_vis = &self.item.vis; let fn_name = self.item.sig.ident.clone(); let item = &self.item; @@ -45,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!( diff --git a/aya-bpf-macros/src/fentry.rs b/aya-bpf-macros/src/fentry.rs index d1a2ed7f..ad26f2ca 100644 --- a/aya-bpf-macros/src/fentry.rs +++ b/aya-bpf-macros/src/fentry.rs @@ -1,26 +1,22 @@ 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_bool_arg, pop_required_string_arg}; +use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg}; pub(crate) struct FEntry { item: ItemFn, - function: String, + function: Option, sleepable: bool, } impl FEntry { pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { - if attrs.is_empty() { - abort!(attrs, "missing function name"); - } - let mut args = syn::parse2(attrs)?; let item = syn::parse2(item)?; - let function = pop_required_string_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 { @@ -32,7 +28,11 @@ impl FEntry { pub(crate) fn expand(&self) -> Result { 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 self.function.is_none() { + section_prefix.into() + } else { + format!("{}/{}", section_prefix, self.function.as_ref().unwrap()).into() + }; let fn_vis = &self.item.vis; let fn_name = self.item.sig.ident.clone(); let item = &self.item; @@ -56,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" @@ -87,7 +114,6 @@ mod tests { fn test_fentry_sleepable() { let prog = FEntry::parse( parse_quote! { - function = "sys_clone", sleepable }, parse_quote! { @@ -100,7 +126,7 @@ mod tests { let expanded = prog.expand().unwrap(); let expected = quote! { #[no_mangle] - #[link_section = "fentry.s/sys_clone"] + #[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; diff --git a/aya-bpf-macros/src/fexit.rs b/aya-bpf-macros/src/fexit.rs index 80a63b96..ca73ec4f 100644 --- a/aya-bpf-macros/src/fexit.rs +++ b/aya-bpf-macros/src/fexit.rs @@ -1,26 +1,22 @@ 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_bool_arg, pop_required_string_arg}; +use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg}; pub(crate) struct FExit { item: ItemFn, - function: String, + function: Option, sleepable: bool, } impl FExit { pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { - if attrs.is_empty() { - abort!(attrs, "missing function name"); - } - let mut args = syn::parse2(attrs)?; let item = syn::parse2(item)?; - let function = pop_required_string_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 { @@ -32,7 +28,11 @@ impl FExit { pub(crate) fn expand(&self) -> Result { 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 self.function.is_none() { + section_prefix.into() + } else { + format!("{}/{}", section_prefix, self.function.as_ref().unwrap()).into() + }; let fn_vis = &self.item.vis; let fn_name = self.item.sig.ident.clone(); let item = &self.item; @@ -56,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" diff --git a/aya-bpf-macros/src/lsm.rs b/aya-bpf-macros/src/lsm.rs index fb358b0c..e4b1f409 100644 --- a/aya-bpf-macros/src/lsm.rs +++ b/aya-bpf-macros/src/lsm.rs @@ -1,26 +1,22 @@ 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_bool_arg, pop_required_string_arg}; +use crate::args::{err_on_unknown_args, pop_bool_arg, pop_string_arg}; pub(crate) struct Lsm { item: ItemFn, - hook: String, + hook: Option, sleepable: bool, } impl Lsm { pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { - if attrs.is_empty() { - abort!(attrs, "missing hook name"); - } - let mut args = syn::parse2(attrs)?; let item = syn::parse2(item)?; - let hook = pop_required_string_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 { @@ -32,7 +28,11 @@ impl Lsm { pub(crate) fn expand(&self) -> Result { 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 self.hook.is_none() { + section_prefix.into() + } else { + format!("{}/{}", section_prefix, self.hook.as_ref().unwrap()).into() + }; let fn_vis = &self.item.vis; let fn_name = self.item.sig.ident.clone(); let item = &self.item; diff --git a/aya-bpf-macros/src/raw_tracepoint.rs b/aya-bpf-macros/src/raw_tracepoint.rs index 4594b993..583477e1 100644 --- a/aya-bpf-macros/src/raw_tracepoint.rs +++ b/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_string_arg}; +use crate::args::{err_on_unknown_args, pop_string_arg}; pub(crate) struct RawTracePoint { item: ItemFn, - tracepoint: String, + tracepoint: Option, } impl RawTracePoint { pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result { - if attrs.is_empty() { - abort!(attrs, "expected `tracepoint` attribute") - } - let mut args = syn::parse2(attrs)?; let item = syn::parse2(item)?; - let tracepoint = pop_required_string_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 { - let section_name: Cow<'_, _> = format!("raw_tp/{}", self.tracepoint).into(); + let section_name: Cow<'_, _> = if self.tracepoint.is_none() { + "raw_tp".into() + } else { + format!("raw_tp/{}", self.tracepoint.as_ref().unwrap()).into() + }; let fn_vis = &self.item.vis; let fn_name = self.item.sig.ident.clone(); let item = &self.item;