diff --git a/aya-bpf-macros/src/args.rs b/aya-bpf-macros/src/args.rs index 44a0c2cd..b695053a 100644 --- a/aya-bpf-macros/src/args.rs +++ b/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, + pub(crate) args: Vec, } impl Parse for Args { fn parse(input: ParseStream) -> Result { - let args = Punctuated::::parse_terminated_with(input, |input| { - Ok(NameValue { - name: input.parse()?, - _eq: input.parse()?, - value: input.parse()?, - }) + let args = Punctuated::::parse_terminated_with(input, |input| { + let ident = input.parse::()?; + 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 { - 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 { + 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 { - 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> { - let name = pop_arg(args, "name"); - Ok(name) +pub(crate) fn name_arg(args: &mut Args) -> Option { + pop_string_arg(args, "name") } diff --git a/aya-bpf-macros/src/btf_tracepoint.rs b/aya-bpf-macros/src/btf_tracepoint.rs index f05b561e..d6148abd 100644 --- a/aya-bpf-macros/src/btf_tracepoint.rs +++ b/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, } 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_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 { - 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!( diff --git a/aya-bpf-macros/src/cgroup_skb.rs b/aya-bpf-macros/src/cgroup_skb.rs index 94488382..606eb7b5 100644 --- a/aya-bpf-macros/src/cgroup_skb.rs +++ b/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 }) } diff --git a/aya-bpf-macros/src/cgroup_sock.rs b/aya-bpf-macros/src/cgroup_sock.rs index dfc24561..8facb9c1 100644 --- a/aya-bpf-macros/src/cgroup_sock.rs +++ b/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 { + 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() { diff --git a/aya-bpf-macros/src/cgroup_sock_addr.rs b/aya-bpf-macros/src/cgroup_sock_addr.rs index 6ca0f20e..43870c54 100644 --- a/aya-bpf-macros/src/cgroup_sock_addr.rs +++ b/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(), diff --git a/aya-bpf-macros/src/cgroup_sockopt.rs b/aya-bpf-macros/src/cgroup_sockopt.rs index 11440f12..b5bba0dd 100644 --- a/aya-bpf-macros/src/cgroup_sockopt.rs +++ b/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(), diff --git a/aya-bpf-macros/src/fentry.rs b/aya-bpf-macros/src/fentry.rs index d63dfd8b..c8158dc8 100644 --- a/aya-bpf-macros/src/fentry.rs +++ b/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, 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_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 { 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()); + } } diff --git a/aya-bpf-macros/src/fexit.rs b/aya-bpf-macros/src/fexit.rs index b81be64a..7c7b17e1 100644 --- a/aya-bpf-macros/src/fexit.rs +++ b/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, 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_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 { 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()); + } } diff --git a/aya-bpf-macros/src/kprobe.rs b/aya-bpf-macros/src/kprobe.rs index 5393187d..4fb9895f 100644 --- a/aya-bpf-macros/src/kprobe.rs +++ b/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 { - 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::().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::().unwrap()); + err_on_unknown_args(&args)?; + Ok(KProbe { kind, item, diff --git a/aya-bpf-macros/src/lib.rs b/aya-bpf-macros/src/lib.rs index 247b636a..5696b82d 100644 --- a/aya-bpf-macros/src/lib.rs +++ b/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 diff --git a/aya-bpf-macros/src/lsm.rs b/aya-bpf-macros/src/lsm.rs index 546d9385..88a66b9b 100644 --- a/aya-bpf-macros/src/lsm.rs +++ b/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, 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 = 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 { 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( diff --git a/aya-bpf-macros/src/map.rs b/aya-bpf-macros/src/map.rs index 044d9395..21f9a8fa 100644 --- a/aya-bpf-macros/src/map.rs +++ b/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 { - 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 }) } diff --git a/aya-bpf-macros/src/raw_tracepoint.rs b/aya-bpf-macros/src/raw_tracepoint.rs index a220251d..9bb7b658 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_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_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 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; diff --git a/aya-bpf-macros/src/tracepoint.rs b/aya-bpf-macros/src/tracepoint.rs index 4a1b5f50..0e2a7cf6 100644 --- a/aya-bpf-macros/src/tracepoint.rs +++ b/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 { - 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, diff --git a/aya-bpf-macros/src/uprobe.rs b/aya-bpf-macros/src/uprobe.rs index 7f06d0d3..7144f30f 100644 --- a/aya-bpf-macros/src/uprobe.rs +++ b/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, offset: Option, item: ItemFn, + sleepable: bool, } impl UProbe { pub(crate) fn parse(kind: UProbeKind, attrs: TokenStream, item: TokenStream) -> Result { - 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::().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::().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 { + 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( diff --git a/aya-bpf-macros/src/xdp.rs b/aya-bpf-macros/src/xdp.rs index f9b01529..9a5ff36a 100644 --- a/aya-bpf-macros/src/xdp.rs +++ b/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 { - 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 }) }