aya-bpf-macros: Refactor for ease of testing

The aya-bpf-macros needed refactoring for:

1. Ease of testing
2. To be consistent with when we use K/V args vs. idents
3. To deprecate the use of `name` to change the exported name of a
   function - we now use the symbol table.

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
pull/413/head
Dave Tucker 1 year ago
parent bf7fdff1ce
commit c72aab5f7b

@ -76,6 +76,7 @@ num_enum = { version = "0.6", default-features = false }
object = { version = "0.31", default-features = false } object = { version = "0.31", default-features = false }
parking_lot = { version = "0.12.0", default-features = false } parking_lot = { version = "0.12.0", default-features = false }
proc-macro2 = { version = "1", default-features = false } proc-macro2 = { version = "1", default-features = false }
proc-macro-error = { version = "1.0", default-features = false }
public-api = { version = "0.31.2", default-features = false } public-api = { version = "0.31.2", default-features = false }
quote = { version = "1", default-features = false } quote = { version = "1", default-features = false }
rbpf = { version = "0.2.0", default-features = false } rbpf = { version = "0.2.0", default-features = false }

@ -9,6 +9,7 @@ proc-macro = true
[dependencies] [dependencies]
proc-macro2 = { workspace = true } proc-macro2 = { workspace = true }
proc-macro-error = { workspace = true }
quote = { workspace = true } quote = { workspace = true }
syn = { workspace = true, default-features = true, features = ["full"] } syn = { workspace = true, default-features = true, features = ["full"] }

@ -0,0 +1,69 @@
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) struct Args {
pub(crate) args: Vec<NameValue>,
}
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()?,
})
})?
.into_pairs()
.map(|pair| match pair {
Pair::Punctuated(name_val, _) => name_val,
Pair::End(name_val) => name_val,
})
.collect();
Ok(Args { 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()),
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()),
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 err_on_unknown_args(args: &Args) -> Result<()> {
if let Some(arg) = args.args.get(0) {
return Err(Error::new_spanned(&arg.name, "invalid argument"));
}
Ok(())
}
pub(crate) fn name_arg(args: &mut Args) -> Result<Option<String>> {
let name = pop_arg(args, "name");
Ok(name)
}

@ -0,0 +1,89 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use quote::quote;
use syn::{Error, ItemFn, Result};
use crate::args::{err_on_unknown_args, pop_arg, pop_required_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'",
));
}
}
err_on_unknown_args(&args)?;
Ok(BtfTracePoint {
item,
function,
sleepable,
})
}
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 fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 {
let _ = #fn_name(::aya_bpf::programs::BtfTracePointContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_btf_tracepoint() {
let prog = BtfTracePoint::parse(
parse_quote!(function = "some_func"),
parse_quote!(
fn foo(ctx: BtfTracePointContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote!(
#[no_mangle]
#[link_section = "tp_btf/some_func"]
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());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct CgroupDevice {
item: ItemFn,
}
impl CgroupDevice {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(CgroupDevice { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "cgroup/dev"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_cgroup_dev_ctx) -> i32 {
return #fn_name(::aya_bpf::programs::DeviceContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_cgroup_device() {
let prog = CgroupDevice::parse(
parse_quote!(),
parse_quote!(
fn foo(ctx: DeviceContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/dev"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_cgroup_dev_ctx) -> i32 {
return foo(::aya_bpf::programs::DeviceContext::new(ctx));
fn foo(ctx: DeviceContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,210 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{Ident, ItemFn, Result};
pub(crate) struct CgroupSkb {
item: ItemFn,
attach_type: Option<String>,
}
impl CgroupSkb {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<CgroupSkb> {
let item: ItemFn = syn::parse2(item)?;
let mut attach_type = None;
if !attrs.is_empty() {
let ident: Ident = syn::parse2(attrs)?;
match ident.to_string().as_str() {
"ingress" | "egress" => (),
_ => abort!(attach_type, "invalid attach type"),
}
attach_type = Some(ident.to_string());
}
Ok(CgroupSkb { item, attach_type })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = if self.attach_type.is_some() {
format!("cgroup_skb/{}", self.attach_type.as_ref().unwrap()).into()
} else {
"cgroup_skb".into()
};
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn cgroup_skb() {
let prog = CgroupSkb::parse(
parse_quote!(),
parse_quote!(
fn foo(ctx: SkBuffContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup_skb"]
fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return foo(::aya_bpf::programs::SkBuffContext::new(ctx));
fn foo(ctx: SkBuffContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_skb_egress() {
let prog = CgroupSkb::parse(
parse_quote!(egress),
parse_quote!(
fn foo(ctx: SkBuffContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup_skb/egress"]
fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return foo(::aya_bpf::programs::SkBuffContext::new(ctx));
fn foo(ctx: SkBuffContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_skb_ingress() {
let prog = CgroupSkb::parse(
parse_quote!(ingress),
parse_quote!(
fn foo(ctx: SkBuffContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup_skb/ingress"]
fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return foo(::aya_bpf::programs::SkBuffContext::new(ctx));
fn foo(ctx: SkBuffContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn priv_function() {
let prog = CgroupSkb::parse(
parse_quote!(egress),
parse_quote!(
fn foo(ctx: SkBuffContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup_skb/egress"]
fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return foo(::aya_bpf::programs::SkBuffContext::new(ctx));
fn foo(ctx: SkBuffContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn pub_function() {
let prog = CgroupSkb::parse(
parse_quote!(egress),
parse_quote!(
pub fn foo(ctx: SkBuffContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup_skb/egress"]
pub fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return foo(::aya_bpf::programs::SkBuffContext::new(ctx));
pub fn foo(ctx: SkBuffContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn pub_crate_function() {
let prog = CgroupSkb::parse(
parse_quote!(egress),
parse_quote!(
pub(crate) fn foo(ctx: SkBuffContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup_skb/egress"]
pub(crate) fn foo(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return foo(::aya_bpf::programs::SkBuffContext::new(ctx));
pub(crate) fn foo(ctx: SkBuffContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,150 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{Ident, ItemFn, Result};
pub(crate) struct CgroupSock {
item: ItemFn,
attach_type: String,
}
impl CgroupSock {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<CgroupSock> {
let item: ItemFn = syn::parse2(item)?;
let attach_type: Ident = syn::parse2(attrs)?;
match attach_type.to_string().as_str() {
"post_bind4" | "post_bind6" | "sock_create" | "sock_release" => (),
_ => abort!(attach_type, "invalid attach type"),
}
Ok(CgroupSock {
item,
attach_type: attach_type.to_string(),
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 {
return #fn_name(::aya_bpf::programs::SockContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn cgroup_sock_post_bind4() {
let prog = CgroupSock::parse(
parse_quote!(post_bind4),
parse_quote!(
fn foo(ctx: SockContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/post_bind4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 {
return foo(::aya_bpf::programs::SockContext::new(ctx));
fn foo(ctx: SockContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_post_bind6() {
let prog = CgroupSock::parse(
parse_quote!(post_bind6),
parse_quote!(
fn foo(ctx: SockContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/post_bind6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 {
return foo(::aya_bpf::programs::SockContext::new(ctx));
fn foo(ctx: SockContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_sock_create() {
let prog = CgroupSock::parse(
parse_quote!(sock_create),
parse_quote!(
fn foo(ctx: SockContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/sock_create"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 {
return foo(::aya_bpf::programs::SockContext::new(ctx));
fn foo(ctx: SockContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_sock_release() {
let prog = CgroupSock::parse(
parse_quote!(sock_release),
parse_quote!(
fn foo(ctx: SockContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/sock_release"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock) -> i32 {
return foo(::aya_bpf::programs::SockContext::new(ctx));
fn foo(ctx: SockContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,365 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{Ident, ItemFn, Result};
pub(crate) struct CgroupSockAddr {
item: ItemFn,
attach_type: String,
}
impl CgroupSockAddr {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if attrs.is_empty() {
abort!(attrs, "missing attach type")
}
let attach_type: Ident = syn::parse2(attrs)?;
match attach_type.to_string().as_str() {
"connect4" | "connect6" | "bind4" | "bind6" | "getpeername4" | "getpeername6"
| "getsockname4" | "getsockname6" | "sendmsg4" | "sendmsg6" | "recvmsg4"
| "recvmsg6" => (),
_ => abort!(attach_type, "invalid attach type"),
}
let item = syn::parse2(item)?;
Ok(CgroupSockAddr {
item,
attach_type: attach_type.to_string(),
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return #fn_name(::aya_bpf::programs::SockAddrContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn cgroup_sock_addr_connect4() {
let prog = CgroupSockAddr::parse(
parse_quote!(connect4),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/connect4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_connect6() {
let prog = CgroupSockAddr::parse(
parse_quote!(connect6),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/connect6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_bind4() {
let prog = CgroupSockAddr::parse(
parse_quote!(bind4),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/bind4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_bind6() {
let prog = CgroupSockAddr::parse(
parse_quote!(bind6),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/bind6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_getpeername4() {
let prog = CgroupSockAddr::parse(
parse_quote!(getpeername4),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/getpeername4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_getpeername6() {
let prog = CgroupSockAddr::parse(
parse_quote!(getpeername6),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/getpeername6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_getsockname4() {
let prog = CgroupSockAddr::parse(
parse_quote!(getsockname4),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/getsockname4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_getsockname6() {
let prog = CgroupSockAddr::parse(
parse_quote!(getsockname6),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/getsockname6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_sendmsg4() {
let prog = CgroupSockAddr::parse(
parse_quote!(sendmsg4),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/sendmsg4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_sendmsg6() {
let prog = CgroupSockAddr::parse(
parse_quote!(sendmsg6),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/sendmsg6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_recvmsg4() {
let prog = CgroupSockAddr::parse(
parse_quote!(recvmsg4),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/recvmsg4"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sock_addr_recvmsg6() {
let prog = CgroupSockAddr::parse(
parse_quote!(recvmsg6),
parse_quote!(
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/recvmsg6"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sock_addr) -> i32 {
return foo(::aya_bpf::programs::SockAddrContext::new(ctx));
fn foo(ctx: CgroupSockAddrContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,103 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{Ident, ItemFn, Result};
pub(crate) struct CgroupSockopt {
item: ItemFn,
attach_type: String,
}
impl CgroupSockopt {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<CgroupSockopt> {
if attrs.is_empty() {
abort!(attrs, "expected attach type");
}
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(),
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = format!("cgroup/{}", self.attach_type).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 {
return #fn_name(::aya_bpf::programs::SockoptContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn cgroup_sockopt_getsockopt() {
let prog = CgroupSockopt::parse(
parse_quote!(getsockopt),
parse_quote!(
fn foo(ctx: SockoptContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote!(
#[no_mangle]
#[link_section = "cgroup/getsockopt"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 {
return foo(::aya_bpf::programs::SockoptContext::new(ctx));
fn foo(ctx: SockoptContext) -> i32 {
0
}
}
);
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn cgroup_sockopt_setsockopt() {
let prog = CgroupSockopt::parse(
parse_quote!(setsockopt),
parse_quote!(
fn foo(ctx: SockoptContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote!(
#[no_mangle]
#[link_section = "cgroup/setsockopt"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sockopt) -> i32 {
return foo(::aya_bpf::programs::SockoptContext::new(ctx));
fn foo(ctx: SockoptContext) -> i32 {
0
}
}
);
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct CgroupSysctl {
item: ItemFn,
}
impl CgroupSysctl {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(CgroupSysctl { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "cgroup/sysctl"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sysctl) -> i32 {
return #fn_name(::aya_bpf::programs::SysctlContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_cgroup_sysctl() {
let prog = CgroupSysctl::parse(
parse_quote!(),
parse_quote!(
fn foo(ctx: SysctlContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "cgroup/sysctl"]
fn foo(ctx: *mut ::aya_bpf::bindings::bpf_sysctl) -> i32 {
return foo(::aya_bpf::programs::SysctlContext::new(ctx));
fn foo(ctx: SysctlContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

File diff suppressed because it is too large Load Diff

@ -0,0 +1,84 @@
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};
pub(crate) struct FEntry {
item: ItemFn,
function: String,
sleepable: bool,
}
impl FEntry {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<FEntry> {
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")?;
err_on_unknown_args(&args)?;
Ok(FEntry {
item,
function,
sleepable: false,
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_prefix = if self.sleepable { "fentry.s" } else { "fentry" };
let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 {
let _ = #fn_name(::aya_bpf::programs::FEntryContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_fentry() {
let prog = FEntry::parse(
parse_quote! {
function = "sys_clone"
},
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/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());
}
}

@ -0,0 +1,84 @@
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};
pub(crate) struct FExit {
item: ItemFn,
function: String,
sleepable: bool,
}
impl FExit {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<FExit> {
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")?;
err_on_unknown_args(&args)?;
Ok(FExit {
item,
function,
sleepable: false,
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_prefix = if self.sleepable { "fexit.s" } else { "fexit" };
let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.function).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 {
let _ = #fn_name(::aya_bpf::programs::FExitContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_fexit() {
let prog = FExit::parse(
parse_quote! {
function = "sys_clone"
},
parse_quote! {
fn sys_clone(ctx: &mut FExitContext) -> i32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "fexit/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());
}
}

@ -0,0 +1,213 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use quote::quote;
use syn::{ItemFn, Result};
use crate::args::{err_on_unknown_args, pop_arg};
#[allow(clippy::enum_variant_names)]
#[derive(Debug, Copy, Clone)]
pub(crate) enum KProbeKind {
KProbe,
KRetProbe,
}
impl std::fmt::Display for KProbeKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
use KProbeKind::*;
match self {
KProbe => write!(f, "kprobe"),
KRetProbe => write!(f, "kretprobe"),
}
}
}
pub(crate) struct KProbe {
kind: KProbeKind,
function: Option<String>,
offset: Option<u64>,
item: ItemFn,
}
impl KProbe {
pub(crate) fn parse(kind: KProbeKind, attrs: TokenStream, item: TokenStream) -> Result<KProbe> {
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::<u64>().unwrap());
err_on_unknown_args(&args)?;
}
let item = syn::parse2(item)?;
Ok(KProbe {
kind,
item,
function,
offset,
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = if self.function.is_some() && self.offset.is_some() {
format!(
"{}/{}+{}",
self.kind,
self.function.as_ref().unwrap(),
self.offset.unwrap()
)
.into()
} else if self.function.is_some() {
format!("{}/{}", self.kind, self.function.as_ref().unwrap()).into()
} else {
format!("{}", self.kind).into()
};
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = #fn_name(::aya_bpf::programs::ProbeContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_kprobe() {
let kprobe = KProbe::parse(
KProbeKind::KProbe,
parse_quote! {},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
kprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "kprobe"]
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 test_kprobe_with_function() {
let kprobe = KProbe::parse(
KProbeKind::KProbe,
parse_quote! {
function = "fib_lookup"
},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
kprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "kprobe/fib_lookup"]
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 test_kprobe_with_function_and_offset() {
let kprobe = KProbe::parse(
KProbeKind::KProbe,
parse_quote! {
function = "fib_lookup",
offset = "10"
},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
kprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "kprobe/fib_lookup+10"]
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 test_kretprobe() {
let kprobe = KProbe::parse(
KProbeKind::KRetProbe,
parse_quote! {},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
kprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "kretprobe"]
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()
);
}
}

@ -1,121 +1,187 @@
mod expand; pub(crate) mod args;
mod btf_tracepoint;
use expand::{ mod cgroup_device;
Args, BtfTracePoint, CgroupDevice, CgroupSkb, CgroupSock, CgroupSockAddr, CgroupSockopt, mod cgroup_skb;
CgroupSysctl, FEntry, FExit, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, mod cgroup_sock;
SchedClassifier, SkLookup, SkMsg, SkSkb, SkSkbKind, SockAddrArgs, SockOps, SocketFilter, mod cgroup_sock_addr;
SockoptArgs, TracePoint, Xdp, mod cgroup_sockopt;
}; mod cgroup_sysctl;
mod fentry;
mod fexit;
mod kprobe;
mod lsm;
mod map;
mod perf_event;
mod raw_tracepoint;
mod sk_lookup;
mod sk_msg;
mod sk_skb;
mod sock_ops;
mod socket_filter;
mod tc;
mod tracepoint;
mod uprobe;
mod xdp;
use btf_tracepoint::BtfTracePoint;
use cgroup_device::CgroupDevice;
use cgroup_skb::CgroupSkb;
use cgroup_sock::CgroupSock;
use cgroup_sock_addr::CgroupSockAddr;
use cgroup_sockopt::CgroupSockopt;
use cgroup_sysctl::CgroupSysctl;
use fentry::FEntry;
use fexit::FExit;
use kprobe::{KProbe, KProbeKind};
use lsm::Lsm;
use map::Map;
use perf_event::PerfEvent;
use proc_macro::TokenStream; use proc_macro::TokenStream;
use syn::{parse_macro_input, ItemFn, ItemStatic}; use proc_macro_error::{abort, proc_macro_error};
use raw_tracepoint::RawTracePoint;
use sk_lookup::SkLookup;
use sk_msg::SkMsg;
use sk_skb::{SkSkb, SkSkbKind};
use sock_ops::SockOps;
use socket_filter::SocketFilter;
use uprobe::{UProbe, UProbeKind};
use tc::SchedClassifier;
use tracepoint::TracePoint;
use xdp::Xdp;
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match Map::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemStatic); Ok(prog) => prog
.expand()
Map::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn kprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn kprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
probe(ProbeKind::KProbe, attrs, item) match KProbe::parse(KProbeKind::KProbe, attrs.into(), item.into()) {
Ok(prog) => prog
.expand()
.unwrap_or_else(|err| abort!(err.span(), "{}", err))
.into(),
Err(err) => abort!(err.span(), "{}", err),
}
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn kretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn kretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
probe(ProbeKind::KRetProbe, attrs, item) match KProbe::parse(KProbeKind::KRetProbe, attrs.into(), item.into()) {
Ok(prog) => prog
.expand()
.unwrap_or_else(|err| abort!(err.span(), "{}", err))
.into(),
Err(err) => abort!(err.span(), "{}", err),
}
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn uprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn uprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
probe(ProbeKind::UProbe, attrs, item) match UProbe::parse(UProbeKind::UProbe, attrs.into(), item.into()) {
Ok(prog) => prog
.expand()
.unwrap_or_else(|err| abort!(err.span(), "{}", err))
.into(),
Err(err) => abort!(err.span(), "{}", err),
}
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn uretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn uretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
probe(ProbeKind::URetProbe, attrs, item) match UProbe::parse(UProbeKind::URetProbe, attrs.into(), item.into()) {
Ok(prog) => prog
.expand()
.unwrap_or_else(|err| abort!(err.span(), "{}", err))
.into(),
Err(err) => abort!(err.span(), "{}", err),
}
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn sock_ops(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn sock_ops(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match SockOps::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
SockOps::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn sk_msg(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn sk_msg(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match SkMsg::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
SkMsg::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn xdp(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn xdp(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match Xdp::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
Xdp::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn classifier(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn classifier(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match SchedClassifier::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
SchedClassifier::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn cgroup_sysctl(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn cgroup_sysctl(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match CgroupSysctl::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
CgroupSysctl::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn cgroup_sockopt(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn cgroup_sockopt(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as SockoptArgs); match CgroupSockopt::parse(attrs.into(), item.into()) {
let attach_type = args.attach_type.to_string(); Ok(prog) => prog
let item = parse_macro_input!(item as ItemFn); .expand()
.unwrap_or_else(|err| abort!(err.span(), "{}", err))
CgroupSockopt::from_syn(args.args, item, attach_type) .into(),
.and_then(|u| u.expand()) Err(err) => abort!(err.span(), "{}", err),
.unwrap_or_else(|err| err.to_compile_error()) }
.into()
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match CgroupSkb::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
CgroupSkb::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a [`CgroupSockAddr`] eBPF program. /// Marks a function as a [`CgroupSockAddr`] eBPF program.
@ -152,59 +218,51 @@ pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(0) /// Ok(0)
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn cgroup_sock_addr(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn cgroup_sock_addr(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as SockAddrArgs); match CgroupSockAddr::parse(attrs.into(), item.into()) {
let attach_type = args.attach_type.to_string(); Ok(prog) => prog
let item = parse_macro_input!(item as ItemFn); .expand()
.unwrap_or_else(|err| abort!(err.span(), "{}", err))
CgroupSockAddr::from_syn(args.args, item, attach_type) .into(),
.and_then(|u| u.expand()) Err(err) => abort!(err.span(), "{}", err),
.unwrap_or_else(|err| err.to_compile_error()) }
.into()
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn cgroup_sock(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn cgroup_sock(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match CgroupSock::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
CgroupSock::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
}
fn probe(kind: ProbeKind, attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args);
let item = parse_macro_input!(item as ItemFn);
Probe::from_syn(kind, args, item)
.and_then(|u| u.expand())
.unwrap_or_else(|err| err.to_compile_error())
.into()
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match TracePoint::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
TracePoint::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match PerfEvent::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
PerfEvent::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a raw tracepoint eBPF program that can be attached at a /// Marks a function as a raw tracepoint eBPF program that can be attached at a
@ -223,7 +281,7 @@ pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// ```no_run /// ```no_run
/// use aya_bpf::{macros::raw_tracepoint, programs::RawTracePointContext}; /// use aya_bpf::{macros::raw_tracepoint, programs::RawTracePointContext};
/// ///
/// #[raw_tracepoint(name = "sys_enter")] /// #[raw_tracepoint(tracepoint = "sys_enter")]
/// pub fn sys_enter(ctx: RawTracePointContext) -> i32 { /// pub fn sys_enter(ctx: RawTracePointContext) -> i32 {
/// match unsafe { try_sys_enter(ctx) } { /// match unsafe { try_sys_enter(ctx) } {
/// Ok(ret) => ret, /// Ok(ret) => ret,
@ -235,20 +293,25 @@ pub fn perf_event(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(0) /// Ok(0)
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match RawTracePoint::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
RawTracePoint::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as an LSM program that can be attached to Linux LSM hooks. /// Marks a function as an LSM program that can be attached to Linux LSM hooks.
/// Used to implement security policy and audit logging. /// 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.
/// Arguments should be comma separated.
///
/// LSM probes can be attached to the kernel's security hooks to implement mandatory /// LSM probes can be attached to the kernel's security hooks to implement mandatory
/// access control policy and security auditing. /// access control policy and security auditing.
/// ///
@ -265,7 +328,7 @@ pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// ```no_run /// ```no_run
/// use aya_bpf::{macros::lsm, programs::LsmContext}; /// use aya_bpf::{macros::lsm, programs::LsmContext};
/// ///
/// #[lsm(name = "file_open")] /// #[lsm(hook = "file_open")]
/// pub fn file_open(ctx: LsmContext) -> i32 { /// pub fn file_open(ctx: LsmContext) -> i32 {
/// match unsafe { try_file_open(ctx) } { /// match unsafe { try_file_open(ctx) } {
/// Ok(ret) => ret, /// Ok(ret) => ret,
@ -277,15 +340,16 @@ pub fn raw_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(0) /// Ok(0)
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match Lsm::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
Lsm::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a [BTF-enabled raw tracepoint][1] eBPF program that can be attached at /// Marks a function as a [BTF-enabled raw tracepoint][1] eBPF program that can be attached at
@ -304,7 +368,7 @@ pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// ```no_run /// ```no_run
/// use aya_bpf::{macros::btf_tracepoint, programs::BtfTracePointContext}; /// use aya_bpf::{macros::btf_tracepoint, programs::BtfTracePointContext};
/// ///
/// #[btf_tracepoint(name = "sched_process_fork")] /// #[btf_tracepoint(function = "sched_process_fork")]
/// pub fn sched_process_fork(ctx: BtfTracePointContext) -> u32 { /// pub fn sched_process_fork(ctx: BtfTracePointContext) -> u32 {
/// match unsafe { try_sched_process_fork(ctx) } { /// match unsafe { try_sched_process_fork(ctx) } {
/// Ok(ret) => ret, /// Ok(ret) => ret,
@ -318,15 +382,16 @@ pub fn lsm(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// ``` /// ```
/// ///
/// [1]: https://github.com/torvalds/linux/commit/9e15db66136a14cde3f35691f1d839d950118826 /// [1]: https://github.com/torvalds/linux/commit/9e15db66136a14cde3f35691f1d839d950118826
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match BtfTracePoint::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
BtfTracePoint::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a SK_SKB Stream Parser eBPF program that can be attached /// Marks a function as a SK_SKB Stream Parser eBPF program that can be attached
@ -354,6 +419,7 @@ pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(ctx.len()) /// Ok(ctx.len())
///} ///}
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream {
sk_skb(SkSkbKind::StreamParser, attrs, item) sk_skb(SkSkbKind::StreamParser, attrs, item)
@ -384,19 +450,20 @@ pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(sk_action::SK_PASS) /// Ok(sk_action::SK_PASS)
///} ///}
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn stream_verdict(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn stream_verdict(attrs: TokenStream, item: TokenStream) -> TokenStream {
sk_skb(SkSkbKind::StreamVerdict, attrs, item) sk_skb(SkSkbKind::StreamVerdict, attrs, item)
} }
fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream { fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match SkSkb::parse(kind, attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
SkSkb::from_syn(kind, args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a eBPF Socket Filter program that can be attached to /// Marks a function as a eBPF Socket Filter program that can be attached to
@ -411,20 +478,21 @@ fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream
/// ```no_run /// ```no_run
/// use aya_bpf::{macros::socket_filter, programs::SkBuffContext}; /// use aya_bpf::{macros::socket_filter, programs::SkBuffContext};
/// ///
/// #[socket_filter(name = "accept_all")] /// #[socket_filter]
/// pub fn accept_all(_ctx: SkBuffContext) -> i64 { /// pub fn accept_all(_ctx: SkBuffContext) -> i64 {
/// return 0 /// return 0
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match SocketFilter::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
SocketFilter::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a fentry eBPF program that can be attached to almost /// Marks a function as a fentry eBPF program that can be attached to almost
@ -444,7 +512,7 @@ pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// # type filename = u32; /// # type filename = u32;
/// # type path = u32; /// # type path = u32;
/// ///
/// #[fentry(name = "filename_lookup")] /// #[fentry(function = "filename_lookup")]
/// fn filename_lookup(ctx: FEntryContext) -> i32 { /// fn filename_lookup(ctx: FEntryContext) -> i32 {
/// match unsafe { try_filename_lookup(ctx) } { /// match unsafe { try_filename_lookup(ctx) } {
/// Ok(ret) => ret, /// Ok(ret) => ret,
@ -459,15 +527,16 @@ pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(0) /// Ok(0)
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match FEntry::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
FEntry::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a fexit eBPF program that can be attached to almost /// Marks a function as a fexit eBPF program that can be attached to almost
@ -488,7 +557,7 @@ pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// # type filename = u32; /// # type filename = u32;
/// # type path = u32; /// # type path = u32;
/// ///
/// #[fexit(name = "filename_lookup")] /// #[fexit(function = "filename_lookup")]
/// fn filename_lookup(ctx: FExitContext) -> i32 { /// fn filename_lookup(ctx: FExitContext) -> i32 {
/// match unsafe { try_filename_lookup(ctx) } { /// match unsafe { try_filename_lookup(ctx) } {
/// Ok(ret) => ret, /// Ok(ret) => ret,
@ -503,15 +572,16 @@ pub fn fentry(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// Ok(0) /// Ok(0)
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match FExit::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
FExit::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as an eBPF Socket Lookup program that can be attached to /// Marks a function as an eBPF Socket Lookup program that can be attached to
@ -526,21 +596,22 @@ pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// ```no_run /// ```no_run
/// use aya_bpf::{macros::sk_lookup, programs::SkLookupContext}; /// use aya_bpf::{macros::sk_lookup, programs::SkLookupContext};
/// ///
/// #[sk_lookup(name = "redirect")] /// #[sk_lookup]
/// pub fn accept_all(_ctx: SkLookupContext) -> u32 { /// pub fn accept_all(_ctx: SkLookupContext) -> u32 {
/// // use sk_assign to redirect /// // use sk_assign to redirect
/// return 0 /// return 0
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match SkLookup::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
SkLookup::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }
/// Marks a function as a cgroup device eBPF program that can be attached to a /// Marks a function as a cgroup device eBPF program that can be attached to a
@ -558,19 +629,20 @@ pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream {
/// programs::DeviceContext, /// programs::DeviceContext,
/// }; /// };
/// ///
/// #[cgroup_device(name="cgroup_dev")] /// #[cgroup_device]
/// pub fn cgroup_dev(ctx: DeviceContext) -> i32 { /// pub fn cgroup_dev(ctx: DeviceContext) -> i32 {
/// // Reject all device access /// // Reject all device access
/// return 0; /// return 0;
/// } /// }
/// ``` /// ```
#[proc_macro_error]
#[proc_macro_attribute] #[proc_macro_attribute]
pub fn cgroup_device(attrs: TokenStream, item: TokenStream) -> TokenStream { pub fn cgroup_device(attrs: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(attrs as Args); match CgroupDevice::parse(attrs.into(), item.into()) {
let item = parse_macro_input!(item as ItemFn); Ok(prog) => prog
.expand()
CgroupDevice::from_syn(args, item) .unwrap_or_else(|err| abort!(err.span(), "{}", err))
.and_then(|u| u.expand()) .into(),
.unwrap_or_else(|err| err.to_compile_error()) Err(err) => abort!(err.span(), "{}", err),
.into() }
} }

@ -0,0 +1,85 @@
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;
pub(crate) struct Lsm {
item: ItemFn,
hook: String,
sleepable: bool,
}
impl Lsm {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Lsm> {
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")?;
err_on_unknown_args(&args)?;
Ok(Lsm {
item,
hook,
sleepable: false,
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_prefix = if self.sleepable { "lsm.s" } else { "lsm" };
let section_name: Cow<'_, _> = format!("{}/{}", section_prefix, self.hook).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
// LSM probes need to return an integer corresponding to the correct
// policy decision. Therefore we do not simply default to a return value
// of 0 as in other program types.
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> i32 {
return #fn_name(::aya_bpf::programs::LsmContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_lsm() {
let prog = Lsm::parse(
parse_quote! {
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/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());
}
}

@ -0,0 +1,75 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use quote::quote;
use syn::Result;
use crate::args::name_arg;
use syn::ItemStatic;
pub(crate) struct Map {
item: ItemStatic,
name: String,
}
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());
Ok(Map { item, name })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = "maps".to_string().into();
let name = &self.name;
let item = &self.item;
Ok(quote! {
#[link_section = #section_name]
#[export_name = #name]
#item
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_map_with_name() {
let map = Map::parse(
parse_quote!(name = "foo"),
parse_quote!(
static BAR: HashMap<&'static str, u32> = HashMap::new();
),
)
.unwrap();
let expanded = map.expand().unwrap();
let expected = quote!(
#[link_section = "maps"]
#[export_name = "foo"]
static BAR: HashMap<&'static str, u32> = HashMap::new();
);
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn test_map_no_name() {
let map = Map::parse(
parse_quote!(),
parse_quote!(
static BAR: HashMap<&'static str, u32> = HashMap::new();
),
)
.unwrap();
let expanded = map.expand().unwrap();
let expected = quote!(
#[link_section = "maps"]
#[export_name = "BAR"]
static BAR: HashMap<&'static str, u32> = HashMap::new();
);
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,67 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct PerfEvent {
item: ItemFn,
}
impl PerfEvent {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(PerfEvent { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "perf_event"]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = #fn_name(::aya_bpf::programs::PerfEventContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_perf_event() {
let prog = PerfEvent::parse(
parse_quote!(),
parse_quote!(
fn foo(ctx: PerfEventContext) -> i32 {
0
}
),
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "perf_event"]
fn foo(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = foo(::aya_bpf::programs::PerfEventContext::new(ctx));
return 0;
fn foo(ctx: PerfEventContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,76 @@
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};
pub(crate) struct RawTracePoint {
item: ItemFn,
tracepoint: String,
}
impl RawTracePoint {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<RawTracePoint> {
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")?;
err_on_unknown_args(&args)?;
Ok(RawTracePoint { item, tracepoint })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = format!("raw_tp/{}", self.tracepoint).into();
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = #fn_name(::aya_bpf::programs::RawTracePointContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_raw_tracepoint() {
let prog = RawTracePoint::parse(
parse_quote! { tracepoint = "sys_enter" },
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::RawTracePointContext) -> i32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "raw_tp/sys_enter"]
fn prog(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = prog(::aya_bpf::programs::RawTracePointContext::new(ctx));
return 0;
fn prog(ctx: &mut ::aya_bpf::programs::RawTracePointContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct SkLookup {
item: ItemFn,
}
impl SkLookup {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(SkLookup { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_name = self.item.sig.ident.clone();
let fn_vis = &self.item.vis;
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "sk_lookup"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sk_lookup) -> u32 {
return #fn_name(::aya_bpf::programs::SkLookupContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_sk_lookup() {
let prog = SkLookup::parse(
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::SkLookupContext) -> u32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "sk_lookup"]
fn prog(ctx: *mut ::aya_bpf::bindings::bpf_sk_lookup) -> u32 {
return prog(::aya_bpf::programs::SkLookupContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::SkLookupContext) -> u32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct SkMsg {
item: ItemFn,
}
impl SkMsg {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(SkMsg { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "sk_msg"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::sk_msg_md) -> u32 {
return #fn_name(::aya_bpf::programs::SkMsgContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_sk_msg() {
let prog = SkMsg::parse(
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::SkMsgContext) -> u32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "sk_msg"]
fn prog(ctx: *mut ::aya_bpf::bindings:: sk_msg_md) -> u32 {
return prog(::aya_bpf::programs::SkMsgContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::SkMsgContext) -> u32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,115 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
#[allow(clippy::enum_variant_names)]
#[derive(Debug, Copy, Clone)]
pub(crate) enum SkSkbKind {
StreamVerdict,
StreamParser,
}
impl std::fmt::Display for SkSkbKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
use SkSkbKind::*;
match self {
StreamVerdict => write!(f, "stream_verdict"),
StreamParser => write!(f, "stream_parser"),
}
}
}
pub(crate) struct SkSkb {
kind: SkSkbKind,
item: ItemFn,
}
impl SkSkb {
pub(crate) fn parse(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> Result<SkSkb> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute");
}
let item = syn::parse2(item)?;
Ok(SkSkb { item, kind })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let kind = &self.kind;
let section_name: Cow<'_, _> = format!("sk_skb/{kind}").into();
let fn_name = self.item.sig.ident.clone();
let fn_vis = &self.item.vis;
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 {
return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_stream_parser() {
let prog = SkSkb::parse(
SkSkbKind::StreamParser,
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "sk_skb/stream_parser"]
fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 {
return prog(::aya_bpf::programs::SkBuffContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn test_stream_verdict() {
let prog = SkSkb::parse(
SkSkbKind::StreamVerdict,
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "sk_skb/stream_verdict"]
fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> u32 {
return prog(::aya_bpf::programs::SkBuffContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> u32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct SockOps {
item: ItemFn,
}
impl SockOps {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(SockOps { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "sockops"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::bpf_sock_ops) -> u32 {
return #fn_name(::aya_bpf::programs::SockOpsContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_sock_ops() {
let prog = SockOps::parse(
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::SockOpsContext) -> u32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "sockops"]
fn prog(ctx: *mut ::aya_bpf::bindings::bpf_sock_ops) -> u32 {
return prog(::aya_bpf::programs::SockOpsContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::SockOpsContext) -> u32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct SocketFilter {
item: ItemFn,
}
impl SocketFilter {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(SocketFilter { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_name = self.item.sig.ident.clone();
let fn_vis = &self.item.vis;
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "socket"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i64 {
return #fn_name(::aya_bpf::programs::SkBuffContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_socket_filter() {
let prog = SocketFilter::parse(
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> i64 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "socket"]
fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i64 {
return prog(::aya_bpf::programs::SkBuffContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::SkBuffContext) -> i64 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,65 @@
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{ItemFn, Result};
pub(crate) struct SchedClassifier {
item: ItemFn,
}
impl SchedClassifier {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Self> {
if !attrs.is_empty() {
abort!(attrs, "unexpected attribute")
}
let item = syn::parse2(item)?;
Ok(SchedClassifier { item })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = "classifier"]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return #fn_name(::aya_bpf::programs::TcContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_sched_classifier() {
let prog = SchedClassifier::parse(
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::TcContext) -> i32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "classifier"]
fn prog(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i32 {
return prog(::aya_bpf::programs::TcContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::TcContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,84 @@
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_arg};
pub(crate) struct TracePoint {
item: ItemFn,
category: Option<String>,
name: Option<String>,
}
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");
err_on_unknown_args(&args)?;
Ok(TracePoint {
item,
category,
name,
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = match (&self.category, &self.name) {
(Some(category), Some(name)) => format!("tracepoint/{}/{}", category, name).into(),
(Some(_), None) => abort!(self.item, "expected `name` and `category` arguments"),
(None, Some(_)) => abort!(self.item, "expected `name` and `category` arguments"),
_ => "tracepoint".into(),
};
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = #fn_name(::aya_bpf::programs::TracePointContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_tracepoint() {
let prog = TracePoint::parse(
parse_quote! { name = "sys_enter_bind", category = "syscalls" },
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::TracePointContext) -> i32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "tracepoint/syscalls/sys_enter_bind"]
fn prog(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = prog(::aya_bpf::programs::TracePointContext::new(ctx));
return 0;
fn prog(ctx: &mut ::aya_bpf::programs::TracePointContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -0,0 +1,232 @@
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_arg};
#[allow(clippy::enum_variant_names)]
#[derive(Debug, Copy, Clone)]
pub(crate) enum UProbeKind {
UProbe,
URetProbe,
}
impl std::fmt::Display for UProbeKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
use UProbeKind::*;
match self {
UProbe => write!(f, "uprobe"),
URetProbe => write!(f, "uretprobe"),
}
}
}
pub(crate) struct UProbe {
kind: UProbeKind,
path: Option<String>,
function: Option<String>,
offset: Option<u64>,
item: ItemFn,
}
impl UProbe {
pub(crate) fn parse(kind: UProbeKind, attrs: TokenStream, item: TokenStream) -> Result<UProbe> {
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::<u64>().unwrap());
err_on_unknown_args(&args)?;
}
let item = syn::parse2(item)?;
Ok(UProbe {
kind,
item,
path,
function,
offset,
})
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
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");
}
let mut path = self.path.as_ref().unwrap().clone();
if path.starts_with('/') {
path.remove(0);
}
format!(
"{}/{}:{}+{}",
self.kind,
path,
self.function.as_ref().unwrap(),
self.offset.unwrap()
)
.into()
} else if self.path.is_some() {
if self.function.is_none() {
abort!(self.item.sig.ident, "expected `function` attribute");
}
let mut path = self.path.as_ref().unwrap().clone();
if path.starts_with('/') {
path.remove(0);
}
format!("{}/{}:{}", self.kind, path, self.function.as_ref().unwrap()).into()
} else {
format!("{}", self.kind).into()
};
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 {
let _ = #fn_name(::aya_bpf::programs::ProbeContext::new(ctx));
return 0;
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn uprobe() {
let uprobe = UProbe::parse(
UProbeKind::UProbe,
parse_quote! {},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
uprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "uprobe"]
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(
UProbeKind::UProbe,
parse_quote! {
path = "/self/proc/exe",
function = "trigger_uprobe"
},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
uprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "uprobe/self/proc/exe:trigger_uprobe"]
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 test_uprobe_with_path_and_offset() {
let uprobe = UProbe::parse(
UProbeKind::UProbe,
parse_quote! {
path = "/self/proc/exe", function = "foo", offset = "123"
},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
uprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "uprobe/self/proc/exe:foo+123"]
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 test_uretprobe() {
let uprobe = UProbe::parse(
UProbeKind::URetProbe,
parse_quote! {},
parse_quote! {
fn foo(ctx: ProbeContext) -> u32 {
0
}
},
)
.unwrap();
assert_eq!(
uprobe.expand().unwrap().to_string(),
quote! {
#[no_mangle]
#[link_section = "uretprobe"]
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()
);
}
}

@ -0,0 +1,104 @@
use std::borrow::Cow;
use proc_macro2::TokenStream;
use proc_macro_error::abort;
use quote::quote;
use syn::{Ident, ItemFn, Result};
pub(crate) struct Xdp {
item: ItemFn,
frags: bool,
}
impl Xdp {
pub(crate) fn parse(attrs: TokenStream, item: TokenStream) -> Result<Xdp> {
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;
}
Ok(Xdp { item, frags })
}
pub(crate) fn expand(&self) -> Result<TokenStream> {
let section_name: Cow<'_, _> = if self.frags {
"xdp.frags".into()
} else {
"xdp".into()
};
let fn_vis = &self.item.vis;
let fn_name = self.item.sig.ident.clone();
let item = &self.item;
Ok(quote! {
#[no_mangle]
#[link_section = #section_name]
#fn_vis fn #fn_name(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 {
return #fn_name(::aya_bpf::programs::XdpContext::new(ctx));
#item
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use syn::parse_quote;
#[test]
fn test_xdp() {
let prog = Xdp::parse(
parse_quote! {},
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "xdp"]
fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 {
return prog(::aya_bpf::programs::XdpContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
#[test]
fn test_xdp_frags() {
let prog = Xdp::parse(
parse_quote! { frags },
parse_quote! {
fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 {
0
}
},
)
.unwrap();
let expanded = prog.expand().unwrap();
let expected = quote! {
#[no_mangle]
#[link_section = "xdp.frags"]
fn prog(ctx: *mut ::aya_bpf::bindings::xdp_md) -> u32 {
return prog(::aya_bpf::programs::XdpContext::new(ctx));
fn prog(ctx: &mut ::aya_bpf::programs::XdpContext) -> i32 {
0
}
}
};
assert_eq!(expected.to_string(), expanded.to_string());
}
}

@ -14,7 +14,7 @@ static FOO: Array<u32> = Array::<u32>::with_max_entries(10, 0);
#[map(name = "BAR")] #[map(name = "BAR")]
static BAZ: Array<u32> = Array::<u32>::with_max_entries(10, 0); static BAZ: Array<u32> = Array::<u32>::with_max_entries(10, 0);
#[xdp(frags = "true")] #[xdp(frags)]
pub fn pass(ctx: XdpContext) -> u32 { pub fn pass(ctx: XdpContext) -> u32 {
match unsafe { try_pass(ctx) } { match unsafe { try_pass(ctx) } {
Ok(ret) => ret, Ok(ret) => ret,

@ -5,7 +5,7 @@ use aya_bpf::{bindings::xdp_action, macros::xdp, programs::XdpContext};
// Note: the `frags` attribute causes this probe to be incompatible with kernel versions < 5.18.0. // Note: the `frags` attribute causes this probe to be incompatible with kernel versions < 5.18.0.
// See https://github.com/torvalds/linux/commit/c2f2cdb. // See https://github.com/torvalds/linux/commit/c2f2cdb.
#[xdp(name = "pass", frags = "true")] #[xdp(frags)]
pub fn pass(ctx: XdpContext) -> u32 { pub fn pass(ctx: XdpContext) -> u32 {
match unsafe { try_pass(ctx) } { match unsafe { try_pass(ctx) } {
Ok(ret) => ret, Ok(ret) => ret,

@ -7,7 +7,7 @@ use aya_bpf::{
programs::{ProbeContext, TracePointContext, XdpContext}, programs::{ProbeContext, TracePointContext, XdpContext},
}; };
#[xdp(name = "test_xdp")] #[xdp]
pub fn pass(ctx: XdpContext) -> u32 { pub fn pass(ctx: XdpContext) -> u32 {
match unsafe { try_pass(ctx) } { match unsafe { try_pass(ctx) } {
Ok(ret) => ret, Ok(ret) => ret,

@ -1,6 +1,7 @@
use core::{mem::size_of, ptr::null_mut, slice::from_raw_parts}; use core::{mem::size_of, ptr::null_mut, slice::from_raw_parts};
use std::collections::HashMap; use std::collections::HashMap;
use assert_matches::assert_matches;
use aya_obj::{generated::bpf_insn, Object, ProgramSection}; use aya_obj::{generated::bpf_insn, Object, ProgramSection};
#[test] #[test]
@ -8,8 +9,8 @@ fn run_with_rbpf() {
let object = Object::parse(crate::PASS).unwrap(); let object = Object::parse(crate::PASS).unwrap();
assert_eq!(object.programs.len(), 1); assert_eq!(object.programs.len(), 1);
assert_matches::assert_matches!(object.programs["pass"].section, ProgramSection::Xdp { .. }); assert_matches!(object.programs["pass"].section, ProgramSection::Xdp { .. });
assert_eq!(object.programs["pass"].section.name(), "pass"); assert_eq!(object.programs["pass"].section.name(), "xdp.frags");
let instructions = &object let instructions = &object
.functions .functions
@ -35,7 +36,7 @@ fn use_map_with_rbpf() {
let mut object = Object::parse(crate::MULTIMAP_BTF).unwrap(); let mut object = Object::parse(crate::MULTIMAP_BTF).unwrap();
assert_eq!(object.programs.len(), 1); assert_eq!(object.programs.len(), 1);
assert_matches::assert_matches!( assert_matches!(
object.programs["bpf_prog"].section, object.programs["bpf_prog"].section,
ProgramSection::TracePoint { .. } ProgramSection::TracePoint { .. }
); );

Loading…
Cancel
Save