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 }) } diff --git a/aya-obj/src/obj.rs b/aya-obj/src/obj.rs index 06864abb..450b8e90 100644 --- a/aya-obj/src/obj.rs +++ b/aya-obj/src/obj.rs @@ -222,7 +222,6 @@ pub struct Function { /// Currently, the following section names are not supported yet: /// - `flow_dissector`: `BPF_PROG_TYPE_FLOW_DISSECTOR` /// - `ksyscall+` or `kretsyscall+` -/// - `uprobe.s+` or `uretprobe.s+` /// - `usdt+` /// - `kprobe.multi+` or `kretprobe.multi+`: `BPF_TRACE_KPROBE_MULTI` /// - `lsm_cgroup+` @@ -233,7 +232,6 @@ pub struct Function { /// - `syscall` /// - `struct_ops+` /// - `fmod_ret+`, `fmod_ret.s+` -/// - `fentry.s+`, `fexit.s+` /// - `iter+`, `iter.s+` /// - `xdp.frags/cpumap`, `xdp/cpumap` /// - `xdp.frags/devmap`, `xdp/devmap` @@ -248,9 +246,11 @@ pub enum ProgramSection { }, UProbe { name: String, + sleepable: bool, }, URetProbe { name: String, + sleepable: bool, }, TracePoint { name: String, @@ -315,9 +315,11 @@ pub enum ProgramSection { }, FEntry { name: String, + sleepable: bool, }, FExit { name: String, + sleepable: bool, }, Extension { name: String, @@ -340,8 +342,8 @@ impl ProgramSection { match self { ProgramSection::KRetProbe { name } => name, ProgramSection::KProbe { name } => name, - ProgramSection::UProbe { name } => name, - ProgramSection::URetProbe { name } => name, + ProgramSection::UProbe { name, .. } => name, + ProgramSection::URetProbe { name, .. } => name, ProgramSection::TracePoint { name } => name, ProgramSection::SocketFilter { name } => name, ProgramSection::Xdp { name, .. } => name, @@ -360,9 +362,9 @@ impl ProgramSection { ProgramSection::PerfEvent { name } => name, ProgramSection::RawTracePoint { name } => name, ProgramSection::Lsm { name, .. } => name, - ProgramSection::BtfTracePoint { name } => name, - ProgramSection::FEntry { name } => name, - ProgramSection::FExit { name } => name, + ProgramSection::BtfTracePoint { name, .. } => name, + ProgramSection::FEntry { name, .. } => name, + ProgramSection::FExit { name, .. } => name, ProgramSection::Extension { name } => name, ProgramSection::SkLookup { name } => name, ProgramSection::CgroupSock { name, .. } => name, @@ -388,8 +390,22 @@ impl FromStr for ProgramSection { Ok(match kind { "kprobe" => KProbe { name }, "kretprobe" => KRetProbe { name }, - "uprobe" => UProbe { name }, - "uretprobe" => URetProbe { name }, + "uprobe" => UProbe { + name, + sleepable: false, + }, + "uprobe.s" => UProbe { + name, + sleepable: true, + }, + "uretprobe" => URetProbe { + name, + sleepable: false, + }, + "uretprobe.s" => URetProbe { + name, + sleepable: true, + }, "xdp" => Xdp { name, frags: false }, "xdp.frags" => Xdp { name, frags: true }, "tp_btf" => BtfTracePoint { name }, @@ -551,8 +567,22 @@ impl FromStr for ProgramSection { name, sleepable: true, }, - "fentry" => FEntry { name }, - "fexit" => FExit { name }, + "fentry" => FEntry { + name, + sleepable: false, + }, + "fentry.s" => FEntry { + name, + sleepable: true, + }, + "fexit" => FExit { + name, + sleepable: false, + }, + "fexit.s" => FExit { + name, + sleepable: true, + }, "freplace" => Extension { name }, "sk_lookup" => SkLookup { name }, _ => { @@ -2018,6 +2048,81 @@ mod tests { ); } + #[test] + fn test_parse_section_uprobe_sleepable() { + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + + assert_matches!( + obj.parse_section(fake_section( + BpfSectionKind::Program, + "uprobe.s/foo", + bytes_of(&fake_ins()), + None + )), + Ok(()) + ); + assert_matches!( + obj.programs.get("foo"), + Some(Program { + section: ProgramSection::UProbe { + sleepable: true, + .. + }, + .. + }) + ); + } + + #[test] + fn test_parse_section_uretprobe() { + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + + assert_matches!( + obj.parse_section(fake_section( + BpfSectionKind::Program, + "uretprobe/foo", + bytes_of(&fake_ins()), + None + )), + Ok(()) + ); + assert_matches!( + obj.programs.get("foo"), + Some(Program { + section: ProgramSection::URetProbe { .. }, + .. + }) + ); + } + + #[test] + fn test_parse_section_uretprobe_sleepable() { + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + + assert_matches!( + obj.parse_section(fake_section( + BpfSectionKind::Program, + "uretprobe.s/foo", + bytes_of(&fake_ins()), + None + )), + Ok(()) + ); + assert_matches!( + obj.programs.get("foo"), + Some(Program { + section: ProgramSection::URetProbe { + sleepable: true, + .. + }, + .. + }) + ); + } + #[test] fn test_parse_section_trace_point() { let mut obj = fake_obj(); @@ -2313,6 +2418,32 @@ mod tests { ); } + #[test] + fn test_parse_section_fentry_sleepable() { + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + + assert_matches!( + obj.parse_section(fake_section( + BpfSectionKind::Program, + "fentry.s/foo", + bytes_of(&fake_ins()), + None + )), + Ok(()) + ); + assert_matches!( + obj.programs.get("foo"), + Some(Program { + section: ProgramSection::FEntry { + sleepable: true, + .. + }, + .. + }) + ); + } + #[test] fn test_parse_section_fexit() { let mut obj = fake_obj(); @@ -2336,6 +2467,32 @@ mod tests { ); } + #[test] + fn test_parse_section_fexit_sleepable() { + let mut obj = fake_obj(); + fake_sym(&mut obj, 0, 0, "foo", FAKE_INS_LEN); + + assert_matches!( + obj.parse_section(fake_section( + BpfSectionKind::Program, + "fexit.s/foo", + bytes_of(&fake_ins()), + None + )), + Ok(()) + ); + assert_matches!( + obj.programs.get("foo"), + Some(Program { + section: ProgramSection::FExit { + sleepable: true, + .. + }, + .. + }) + ); + } + #[test] fn test_parse_section_cgroup_skb_ingress_unnamed() { let mut obj = fake_obj(); diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index aea6d11e..ac3a21e3 100644 --- a/aya/src/bpf.rs +++ b/aya/src/bpf.rs @@ -566,14 +566,28 @@ impl<'a> BpfLoader<'a> { data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), kind: ProbeKind::KRetProbe, }), - ProgramSection::UProbe { .. } => Program::UProbe(UProbe { - data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), - kind: ProbeKind::UProbe, - }), - ProgramSection::URetProbe { .. } => Program::UProbe(UProbe { - data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), - kind: ProbeKind::URetProbe, - }), + ProgramSection::UProbe { sleepable, name: _ } => { + let mut data = + ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level); + if *sleepable { + data.flags = BPF_F_SLEEPABLE; + } + Program::UProbe(UProbe { + data, + kind: ProbeKind::UProbe, + }) + } + ProgramSection::URetProbe { sleepable, name: _ } => { + let mut data = + ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level); + if *sleepable { + data.flags = BPF_F_SLEEPABLE; + } + Program::UProbe(UProbe { + data, + kind: ProbeKind::URetProbe, + }) + } ProgramSection::TracePoint { .. } => Program::TracePoint(TracePoint { data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), }), @@ -653,7 +667,7 @@ impl<'a> BpfLoader<'a> { data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), }) } - ProgramSection::Lsm { sleepable, .. } => { + ProgramSection::Lsm { sleepable, name: _ } => { let mut data = ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level); if *sleepable { @@ -666,12 +680,22 @@ impl<'a> BpfLoader<'a> { data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), }) } - ProgramSection::FEntry { .. } => Program::FEntry(FEntry { - data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), - }), - ProgramSection::FExit { .. } => Program::FExit(FExit { - data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), - }), + ProgramSection::FEntry { sleepable, name: _ } => { + let mut data = + ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level); + if *sleepable { + data.flags = BPF_F_SLEEPABLE; + } + Program::FEntry(FEntry { data }) + } + ProgramSection::FExit { sleepable, name: _ } => { + let mut data = + ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level); + if *sleepable { + data.flags = BPF_F_SLEEPABLE; + } + Program::FExit(FExit { data }) + } ProgramSection::Extension { .. } => Program::Extension(Extension { data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level), }), diff --git a/xtask/public-api/aya-obj.txt b/xtask/public-api/aya-obj.txt index 8144037f..20527b54 100644 --- a/xtask/public-api/aya-obj.txt +++ b/xtask/public-api/aya-obj.txt @@ -5837,8 +5837,10 @@ pub aya_obj::obj::ProgramSection::Extension pub aya_obj::obj::ProgramSection::Extension::name: alloc::string::String pub aya_obj::obj::ProgramSection::FEntry pub aya_obj::obj::ProgramSection::FEntry::name: alloc::string::String +pub aya_obj::obj::ProgramSection::FEntry::sleepable: bool pub aya_obj::obj::ProgramSection::FExit pub aya_obj::obj::ProgramSection::FExit::name: alloc::string::String +pub aya_obj::obj::ProgramSection::FExit::sleepable: bool pub aya_obj::obj::ProgramSection::KProbe pub aya_obj::obj::ProgramSection::KProbe::name: alloc::string::String pub aya_obj::obj::ProgramSection::KRetProbe @@ -5870,8 +5872,10 @@ pub aya_obj::obj::ProgramSection::TracePoint pub aya_obj::obj::ProgramSection::TracePoint::name: alloc::string::String pub aya_obj::obj::ProgramSection::UProbe pub aya_obj::obj::ProgramSection::UProbe::name: alloc::string::String +pub aya_obj::obj::ProgramSection::UProbe::sleepable: bool pub aya_obj::obj::ProgramSection::URetProbe pub aya_obj::obj::ProgramSection::URetProbe::name: alloc::string::String +pub aya_obj::obj::ProgramSection::URetProbe::sleepable: bool pub aya_obj::obj::ProgramSection::Xdp pub aya_obj::obj::ProgramSection::Xdp::frags: bool pub aya_obj::obj::ProgramSection::Xdp::name: alloc::string::String @@ -6585,8 +6589,10 @@ pub aya_obj::ProgramSection::Extension pub aya_obj::ProgramSection::Extension::name: alloc::string::String pub aya_obj::ProgramSection::FEntry pub aya_obj::ProgramSection::FEntry::name: alloc::string::String +pub aya_obj::ProgramSection::FEntry::sleepable: bool pub aya_obj::ProgramSection::FExit pub aya_obj::ProgramSection::FExit::name: alloc::string::String +pub aya_obj::ProgramSection::FExit::sleepable: bool pub aya_obj::ProgramSection::KProbe pub aya_obj::ProgramSection::KProbe::name: alloc::string::String pub aya_obj::ProgramSection::KRetProbe @@ -6618,8 +6624,10 @@ pub aya_obj::ProgramSection::TracePoint pub aya_obj::ProgramSection::TracePoint::name: alloc::string::String pub aya_obj::ProgramSection::UProbe pub aya_obj::ProgramSection::UProbe::name: alloc::string::String +pub aya_obj::ProgramSection::UProbe::sleepable: bool pub aya_obj::ProgramSection::URetProbe pub aya_obj::ProgramSection::URetProbe::name: alloc::string::String +pub aya_obj::ProgramSection::URetProbe::sleepable: bool pub aya_obj::ProgramSection::Xdp pub aya_obj::ProgramSection::Xdp::frags: bool pub aya_obj::ProgramSection::Xdp::name: alloc::string::String