diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs index 3d818129..6861bff7 100644 --- a/bpf/aya-bpf-macros/src/expand.rs +++ b/bpf/aya-bpf-macros/src/expand.rs @@ -491,3 +491,35 @@ impl SkSkb { }) } } + +pub struct SocketFilter { + item: ItemFn, + name: Option, +} + +impl SocketFilter { + pub fn from_syn(mut args: Args, item: ItemFn) -> Result { + let name = name_arg(&mut args)?; + + Ok(SocketFilter { item, name }) + } + + pub fn expand(&self) -> Result { + let section_name = if let Some(name) = &self.name { + format!("socket/{}", name) + } else { + "socket".to_owned() + }; + let fn_name = &self.item.sig.ident; + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + fn #fn_name(ctx: *mut ::aya_bpf::bindings::__sk_buff) -> i64 { + return #fn_name(::aya_bpf::programs::SkSkbContext::new(ctx)); + + #item + } + }) + } +} diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs index 55f013b6..bfa0b855 100644 --- a/bpf/aya-bpf-macros/src/lib.rs +++ b/bpf/aya-bpf-macros/src/lib.rs @@ -1,14 +1,12 @@ mod expand; use expand::{ - Args, BtfTracePoint, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, - SkMsg, SkSkb, SkSkbKind, SockOps, TracePoint, Xdp, + Args, BtfTracePoint, CgroupSkb, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, + SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint, Xdp, }; use proc_macro::TokenStream; use syn::{parse_macro_input, ItemFn, ItemStatic}; -use crate::expand::CgroupSkb; - #[proc_macro_attribute] pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream { let args = parse_macro_input!(attrs as Args); @@ -318,3 +316,31 @@ fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream .unwrap_or_else(|err| err.to_compile_error()) .into() } + +/// Marks a function as a eBPF Socket Filter program that can be attached to +/// a socket. +/// +/// # Minimum kernel version +/// +/// The minimum kernel version required to use this feature is 3.19 +/// +/// # Examples +/// +/// ```no_run +/// use aya_bpf::{macros::socket_filter, programs::SkSkbContext}; +/// +/// #[socket_filter(name = "accept_all")] +/// pub fn accept_all(_ctx: SkSkbContext) -> i64 { +/// return 0 +/// } +/// ``` +#[proc_macro_attribute] +pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream { + let args = parse_macro_input!(attrs as Args); + let item = parse_macro_input!(item as ItemFn); + + SocketFilter::from_syn(args, item) + .and_then(|u| u.expand()) + .unwrap_or_else(|err| err.to_compile_error()) + .into() +}