diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs index 3b7a0630c..54548e603 100644 --- a/bpf/aya-bpf-macros/src/expand.rs +++ b/bpf/aya-bpf-macros/src/expand.rs @@ -686,6 +686,38 @@ impl FExit { } } +pub struct SkLookup { + item: ItemFn, + name: Option, +} + +impl SkLookup { + pub fn from_syn(mut args: Args, item: ItemFn) -> Result { + let name = name_arg(&mut args)?; + + Ok(SkLookup { item, name }) + } + + pub fn expand(&self) -> Result { + let section_name = if let Some(name) = &self.name { + format!("sk_lookup/{}", name) + } else { + "sk_lookup".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::bpf_sk_lookup) -> u32 { + return #fn_name(::aya_bpf::programs::SkLookupContext::new(ctx)); + + #item + } + }) + } +} + #[cfg(test)] mod tests { use syn::parse_quote; diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs index 45e54130b..7bd2378fd 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, CgroupSkb, CgroupSockAddr, CgroupSysctl, FEntry, FExit, Lsm, Map, - PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, + PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb, SkSkbKind, SockAddrArgs, SockOps, SocketFilter, TracePoint, Xdp, }; use proc_macro::TokenStream; @@ -455,3 +455,32 @@ pub fn fexit(attrs: TokenStream, item: TokenStream) -> TokenStream { .unwrap_or_else(|err| err.to_compile_error()) .into() } + +/// Marks a function as an eBPF Socket Lookup program that can be attached to +/// a network namespace. +/// +/// # Minimum kernel version +/// +/// The minimum kernel version required to use this feature is 5.9 +/// +/// # Examples +/// +/// ```no_run +/// use aya_bpf::{macros::sk_lookup, programs::SkLookupContext}; +/// +/// #[sk_lookup(name = "redirect")] +/// pub fn accept_all(_ctx: SkLookupContext) -> u32 { +/// // use sk_assign to redirect +/// return 0 +/// } +/// ``` +#[proc_macro_attribute] +pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream { + let args = parse_macro_input!(attrs as Args); + let item = parse_macro_input!(item as ItemFn); + + SkLookup::from_syn(args, item) + .and_then(|u| u.expand()) + .unwrap_or_else(|err| err.to_compile_error()) + .into() +} diff --git a/bpf/aya-bpf/src/maps/sock_hash.rs b/bpf/aya-bpf/src/maps/sock_hash.rs index 3148ee243..fe2348280 100644 --- a/bpf/aya-bpf/src/maps/sock_hash.rs +++ b/bpf/aya-bpf/src/maps/sock_hash.rs @@ -4,9 +4,9 @@ use aya_bpf_cty::c_void; use crate::{ bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKHASH, bpf_sock_ops}, - helpers::{bpf_msg_redirect_hash, bpf_sk_redirect_hash, bpf_sock_hash_update}, + helpers::{bpf_msg_redirect_hash, bpf_sk_redirect_hash, bpf_sock_hash_update, bpf_map_lookup_elem, bpf_sk_assign, bpf_sk_release}, maps::PinningType, - programs::{SkBuffContext, SkMsgContext}, + programs::{SkBuffContext, SkMsgContext, SkLookupContext}, BpfContext, }; @@ -85,4 +85,19 @@ impl SockHash { ) } } + + pub fn redirect_sk_lookup(&mut self, ctx: &SkLookupContext, key: K, flags: u64) -> Result<(),()> { + unsafe { + let sk = bpf_map_lookup_elem( + &mut self.def as *mut _ as *mut _, + &key as *const _ as *const c_void, + ); + if sk.is_null() { + return Err(()) + } + let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags); + bpf_sk_release(sk); + (ret >= 0).then(|| ()).ok_or(()) + } + } } diff --git a/bpf/aya-bpf/src/maps/sock_map.rs b/bpf/aya-bpf/src/maps/sock_map.rs index 7d546419c..3e3629ce7 100644 --- a/bpf/aya-bpf/src/maps/sock_map.rs +++ b/bpf/aya-bpf/src/maps/sock_map.rs @@ -4,9 +4,9 @@ use aya_bpf_cty::c_void; use crate::{ bindings::{bpf_map_def, bpf_map_type::BPF_MAP_TYPE_SOCKMAP, bpf_sock_ops}, - helpers::{bpf_msg_redirect_map, bpf_sk_redirect_map, bpf_sock_map_update}, + helpers::{bpf_msg_redirect_map, bpf_sk_redirect_map, bpf_sock_map_update, bpf_map_lookup_elem, bpf_sk_assign, bpf_sk_release}, maps::PinningType, - programs::{SkBuffContext, SkMsgContext}, + programs::{SkBuffContext, SkLookupContext, SkMsgContext}, BpfContext, }; @@ -80,4 +80,19 @@ impl SockMap { flags, ) } + + pub fn redirect_sk_lookup(&mut self, ctx: &SkLookupContext, index: u32, flags: u64) -> Result<(),()> { + unsafe { + let sk = bpf_map_lookup_elem( + &mut self.def as *mut _ as *mut _, + &index as *const _ as *const c_void, + ); + if sk.is_null() { + return Err(()) + } + let ret = bpf_sk_assign(ctx.as_ptr() as *mut _, sk, flags); + bpf_sk_release(sk); + (ret >= 0).then(|| ()).ok_or(()) + } + } } diff --git a/bpf/aya-bpf/src/programs/mod.rs b/bpf/aya-bpf/src/programs/mod.rs index ab238c518..68e3b8290 100644 --- a/bpf/aya-bpf/src/programs/mod.rs +++ b/bpf/aya-bpf/src/programs/mod.rs @@ -5,6 +5,7 @@ pub mod perf_event; pub mod probe; pub mod raw_tracepoint; pub mod sk_buff; +pub mod sk_lookup; pub mod sk_msg; pub mod sock_addr; pub mod sock_ops; @@ -20,6 +21,7 @@ pub use perf_event::PerfEventContext; pub use probe::ProbeContext; pub use raw_tracepoint::RawTracePointContext; pub use sk_buff::SkBuffContext; +pub use sk_lookup::SkLookupContext; pub use sk_msg::SkMsgContext; pub use sock_addr::SockAddrContext; pub use sock_ops::SockOpsContext; diff --git a/bpf/aya-bpf/src/programs/sk_lookup.rs b/bpf/aya-bpf/src/programs/sk_lookup.rs new file mode 100644 index 000000000..5c37066e8 --- /dev/null +++ b/bpf/aya-bpf/src/programs/sk_lookup.rs @@ -0,0 +1,19 @@ +use core::ffi::c_void; + +use crate::{bindings::bpf_sk_lookup, BpfContext}; + +pub struct SkLookupContext { + pub lookup: *mut bpf_sk_lookup, +} + +impl SkLookupContext { + pub fn new(lookup: *mut bpf_sk_lookup) -> SkLookupContext { + SkLookupContext { lookup } + } +} + +impl BpfContext for SkLookupContext { + fn as_ptr(&self) -> *mut c_void { + self.lookup as *mut _ + } +}