diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs index 8a047558..d96f0ee6 100644 --- a/bpf/aya-bpf-macros/src/expand.rs +++ b/bpf/aya-bpf-macros/src/expand.rs @@ -410,6 +410,51 @@ impl CgroupSockAddr { } } +pub struct CgroupSock { + item: ItemFn, + attach_type: Option, + name: Option, +} + +impl CgroupSock { + pub fn from_syn(mut args: Args, item: ItemFn) -> Result { + let name = pop_arg(&mut args, "name"); + let attach_type = pop_arg(&mut args, "attach"); + err_on_unknown_args(&args)?; + + Ok(CgroupSock { + item, + attach_type, + name, + }) + } + + pub fn expand(&self) -> Result { + let section_name = if let Some(name) = &self.name { + if let Some(attach_type) = &self.attach_type { + format!("cgroup/{}/{}", attach_type, name) + } else { + format!("cgroup/sock/{}", name) + } + } else if let Some(attach_type) = &self.attach_type { + format!("cgroup/{}", attach_type) + } else { + "cgroup/sock".to_string() + }; + 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_sock) -> i32 { + return #fn_name(::aya_bpf::programs::SockContext::new(ctx)); + + #item + } + }) + } +} + fn pop_arg(args: &mut Args, name: &str) -> Option { match args.args.iter().position(|arg| arg.name == name) { Some(index) => Some(args.args.remove(index).value.value()), diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs index a22e9d25..be255952 100644 --- a/bpf/aya-bpf-macros/src/lib.rs +++ b/bpf/aya-bpf-macros/src/lib.rs @@ -1,9 +1,9 @@ mod expand; use expand::{ - Args, BtfTracePoint, CgroupSkb, CgroupSockAddr, CgroupSockopt, CgroupSysctl, FEntry, FExit, - Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb, - SkSkbKind, SockAddrArgs, SockOps, SocketFilter, SockoptArgs, TracePoint, Xdp, + Args, BtfTracePoint, CgroupSkb, CgroupSock, CgroupSockAddr, CgroupSockopt, CgroupSysctl, + FEntry, FExit, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkLookup, + SkMsg, SkSkb, SkSkbKind, SockAddrArgs, SockOps, SocketFilter, SockoptArgs, TracePoint, Xdp, }; use proc_macro::TokenStream; use syn::{parse_macro_input, ItemFn, ItemStatic}; @@ -129,6 +129,17 @@ pub fn cgroup_sock_addr(attrs: TokenStream, item: TokenStream) -> TokenStream { .into() } +#[proc_macro_attribute] +pub fn cgroup_sock(attrs: TokenStream, item: TokenStream) -> TokenStream { + let args = parse_macro_input!(attrs as Args); + let item = parse_macro_input!(item as ItemFn); + + CgroupSock::from_syn(args, item) + .and_then(|u| u.expand()) + .unwrap_or_else(|err| err.to_compile_error()) + .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); diff --git a/bpf/aya-bpf/src/programs/mod.rs b/bpf/aya-bpf/src/programs/mod.rs index c2c8034c..c16160cb 100644 --- a/bpf/aya-bpf/src/programs/mod.rs +++ b/bpf/aya-bpf/src/programs/mod.rs @@ -7,6 +7,7 @@ pub mod raw_tracepoint; pub mod sk_buff; pub mod sk_lookup; pub mod sk_msg; +pub mod sock; pub mod sock_addr; pub mod sock_ops; pub mod sockopt; @@ -24,6 +25,7 @@ pub use raw_tracepoint::RawTracePointContext; pub use sk_buff::SkBuffContext; pub use sk_lookup::SkLookupContext; pub use sk_msg::SkMsgContext; +pub use sock::SockContext; pub use sock_addr::SockAddrContext; pub use sock_ops::SockOpsContext; pub use sockopt::SockoptContext; diff --git a/bpf/aya-bpf/src/programs/sock.rs b/bpf/aya-bpf/src/programs/sock.rs new file mode 100644 index 00000000..0a04888c --- /dev/null +++ b/bpf/aya-bpf/src/programs/sock.rs @@ -0,0 +1,19 @@ +use core::ffi::c_void; + +use crate::{bindings::bpf_sock, BpfContext}; + +pub struct SockContext { + pub sock: *mut bpf_sock, +} + +impl SockContext { + pub fn new(sock: *mut bpf_sock) -> SockContext { + SockContext { sock } + } +} + +impl BpfContext for SockContext { + fn as_ptr(&self) -> *mut c_void { + self.sock as *mut _ + } +}