From 8f4b609a36bda59e0ccd1c5d732aec65d0b7f936 Mon Sep 17 00:00:00 2001 From: Dave Tucker Date: Tue, 1 Aug 2023 17:11:26 +0100 Subject: [PATCH] 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 --- aya-bpf-macros/src/args.rs | 87 ++++++++++++++++++++-------- aya-bpf-macros/src/btf_tracepoint.rs | 27 ++------- aya-bpf-macros/src/fentry.rs | 37 +++++++++++- aya-bpf-macros/src/fexit.rs | 36 +++++++++++- aya-bpf-macros/src/kprobe.rs | 6 +- aya-bpf-macros/src/lib.rs | 2 +- aya-bpf-macros/src/lsm.rs | 36 +++++++++++- aya-bpf-macros/src/map.rs | 2 +- aya-bpf-macros/src/raw_tracepoint.rs | 4 +- aya-bpf-macros/src/tracepoint.rs | 6 +- aya-bpf-macros/src/uprobe.rs | 53 ++++++++++++++--- 11 files changed, 225 insertions(+), 71 deletions(-) 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, + 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 lookahead.peek(Token![=]) { + let _ = input.parse::()?; + 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 { - 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 { + 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 { - 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 { + 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> { - 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..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 { 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 { - 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::().unwrap()); + function = pop_string_arg(&mut args, "function"); + offset = pop_string_arg(&mut args, "offset").map(|v| v.parse::().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 { 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 { 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, offset: Option, 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::().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::().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 { + 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(