diff --git a/aya-bpf-macros/src/args.rs b/aya-bpf-macros/src/args.rs index 44a0c2cd..5547dbaa 100644 --- a/aya-bpf-macros/src/args.rs +++ b/aya-bpf-macros/src/args.rs @@ -7,22 +7,32 @@ use syn::{ 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 lookahead.peek(Token![=]) { + let _ = input.parse::<Eq>()?; + Ok(Arg::String(NameValue { + name: ident, + value: input.parse()?, + })) + } else { + Ok(Arg::Bool(ident)) + } })? .into_pairs() .map(|pair| match pair { @@ -35,35 +45,66 @@ 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()), +pub(crate) fn pop_string_arg(args: &mut Args, name: &str) -> Option<String> { + let value = match args.args.iter().position(|arg| match arg { + Arg::String(name_val) => name_val.name == name, + Arg::Bool(_) => false, + }) { + Some(index) => Some(args.args.remove(index)), None => None, + }; + match value { + Some(Arg::String(value)) => Some(value.value.value()), + Some(Arg::Bool(_)) | None => None, } } -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()), +pub(crate) fn pop_bool_arg(args: &mut Args, name: &str) -> bool { + let value = match args.args.iter().position(|arg| match arg { + Arg::String(_) => false, + Arg::Bool(ident) => ident == name, + }) { + Some(index) => Some(args.args.remove(index)), + None => None, + }; + value.is_some() +} + +pub(crate) fn pop_required_string_arg(args: &mut Args, name: &str) -> Result<String> { + let value = match args.args.iter().position(|arg| match arg { + Arg::String(name_val) => name_val.name == name, + Arg::Bool(_) => false, + }) { + Some(index) => Some(args.args.remove(index)), None => None, }; match value { - Some(value) => Ok(value), - None => Err(Error::new_spanned( - args.args.first().unwrap().name.clone(), - format!("missing required argument `{}`", name), - )), + Some(Arg::String(value)) => Ok(value.value.value()), + Some(Arg::Bool(_)) => unreachable!("arg bool were filtered out"), + None => { + let tokens = match args.args.first().unwrap() { + Arg::String(name_val) => &name_val.name, + Arg::Bool(ident) => ident, + }; + Err(Error::new_spanned( + tokens, + 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>> { - let name = pop_arg(args, "name"); - Ok(name) +pub(crate) fn name_arg(args: &mut Args) -> Option<String> { + 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..62c6ab5e 100644 --- a/aya-bpf-macros/src/btf_tracepoint.rs +++ b/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 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 function = pop_required_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<'_, _> = format!("tp_btf/{}", self.function).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/fentry.rs b/aya-bpf-macros/src/fentry.rs index d63dfd8b..d1a2ed7f 100644 --- a/aya-bpf-macros/src/fentry.rs +++ b/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()); + } } diff --git a/aya-bpf-macros/src/fexit.rs b/aya-bpf-macros/src/fexit.rs index b81be64a..80a63b96 100644 --- a/aya-bpf-macros/src/fexit.rs +++ b/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()); + } } diff --git a/aya-bpf-macros/src/kprobe.rs b/aya-bpf-macros/src/kprobe.rs index 5393187d..7fab3077 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)] @@ -38,8 +38,8 @@ impl KProbe { 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()); + function = pop_string_arg(&mut args, "function"); + offset = pop_string_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap()); err_on_unknown_args(&args)?; } let item = syn::parse2(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..fb358b0c 100644 --- a/aya-bpf-macros/src/lsm.rs +++ b/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( diff --git a/aya-bpf-macros/src/map.rs b/aya-bpf-macros/src/map.rs index 044d9395..16d59077 100644 --- a/aya-bpf-macros/src/map.rs +++ b/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 }) } diff --git a/aya-bpf-macros/src/raw_tracepoint.rs b/aya-bpf-macros/src/raw_tracepoint.rs index a220251d..4594b993 100644 --- a/aya-bpf-macros/src/raw_tracepoint.rs +++ b/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 }) } diff --git a/aya-bpf-macros/src/tracepoint.rs b/aya-bpf-macros/src/tracepoint.rs index 4a1b5f50..85bce4eb 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, @@ -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 category = pop_arg(&mut args, "category"); + 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..06140cb3 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,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"); - function = pop_arg(&mut args, "function"); - offset = pop_arg(&mut args, "offset").map(|v| v.parse::<u64>().unwrap()); + path = pop_string_arg(&mut args, "path"); + function = pop_string_arg(&mut args, "function"); + 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(