diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs index 3ac63028..3d818129 100644 --- a/bpf/aya-bpf-macros/src/expand.rs +++ b/bpf/aya-bpf-macros/src/expand.rs @@ -441,3 +441,53 @@ impl BtfTracePoint { }) } } + +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Copy, Clone)] +pub 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 struct SkSkb { + kind: SkSkbKind, + item: ItemFn, + name: Option, +} + +impl SkSkb { + pub fn from_syn(kind: SkSkbKind, mut args: Args, item: ItemFn) -> Result { + let name = pop_arg(&mut args, "name"); + Ok(SkSkb { item, kind, name }) + } + + pub fn expand(&self) -> Result { + let kind = &self.kind; + let section_name = if let Some(name) = &self.name { + format!("sk_skb/{}/{}", kind, name) + } else { + format!("sk_skb/{}", kind) + }; + 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) -> u32 { + 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 59eeaceb..55f013b6 100644 --- a/bpf/aya-bpf-macros/src/lib.rs +++ b/bpf/aya-bpf-macros/src/lib.rs @@ -2,7 +2,7 @@ mod expand; use expand::{ Args, BtfTracePoint, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, - SkMsg, SockOps, TracePoint, Xdp, + SkMsg, SkSkb, SkSkbKind, SockOps, TracePoint, Xdp, }; use proc_macro::TokenStream; use syn::{parse_macro_input, ItemFn, ItemStatic}; @@ -248,3 +248,73 @@ pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream { .unwrap_or_else(|err| err.to_compile_error()) .into() } + +/// Marks a function as a SK_SKB Stream Parser eBPF program that can be attached +/// to a SockMap +/// +/// # Minimum kernel version +/// +/// The minimum kernel version required to use this feature is 4.14 +/// +/// # Examples +/// +/// ```no_run +/// use aya_bpf::{macros::stream_parser, programs::SkSkbContext}; +/// +/// +///#[stream_parser] +///fn stream_parser(ctx: SkSkbContext) -> u32 { +/// match { try_stream_parser(ctx) } { +/// Ok(ret) => ret, +/// Err(ret) => ret, +/// } +///} +/// +///fn try_stream_parser(ctx: SkSkbContext) -> Result { +/// Ok(ctx.len()) +///} +/// ``` +#[proc_macro_attribute] +pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream { + sk_skb(SkSkbKind::StreamParser, attrs, item) +} + +/// Marks a function as a SK_SKB Stream Verdict eBPF program that can be attached +/// to a SockMap +/// +/// # Minimum kernel version +/// +/// The minimum kernel version required to use this feature is 4.14 +/// +/// # Examples +/// +/// ```no_run +/// use aya_bpf::{macros::stream_verdict, programs::SkSkbContext, bindings::sk_action}; +/// +/// +///#[stream_verdict] +///fn stream_verdict(ctx: SkSkbContext) -> u32 { +/// match { try_stream_verdict(ctx) } { +/// Ok(ret) => ret, +/// Err(ret) => ret, +/// } +///} +/// +///fn try_stream_verdict(_ctx: SkSkbContext) -> Result { +/// Ok(sk_action::SK_PASS) +///} +/// ``` +#[proc_macro_attribute] +pub fn stream_verdict(attrs: TokenStream, item: TokenStream) -> TokenStream { + sk_skb(SkSkbKind::StreamVerdict, attrs, item) +} + +fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream { + let args = parse_macro_input!(attrs as Args); + let item = parse_macro_input!(item as ItemFn); + + SkSkb::from_syn(kind, args, item) + .and_then(|u| u.expand()) + .unwrap_or_else(|err| err.to_compile_error()) + .into() +}