From 2e5ba768374514ec95a326ad7cddc75077d92371 Mon Sep 17 00:00:00 2001 From: Dave Tucker Date: Wed, 18 May 2022 18:37:33 +0100 Subject: [PATCH] bpf: Add support for BPF_PROG_TYPE_CGROUP_SOCK_ADDR Signed-off-by: Dave Tucker --- bpf/aya-bpf-macros/src/expand.rs | 65 +++++++++++++++++++++++++++ bpf/aya-bpf-macros/src/lib.rs | 18 ++++++-- bpf/aya-bpf/src/programs/mod.rs | 2 + bpf/aya-bpf/src/programs/sock_addr.rs | 19 ++++++++ 4 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 bpf/aya-bpf/src/programs/sock_addr.rs diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs index 7e3947acd..3b7a0630c 100644 --- a/bpf/aya-bpf-macros/src/expand.rs +++ b/bpf/aya-bpf-macros/src/expand.rs @@ -37,6 +37,29 @@ impl Parse for Args { } } +pub struct SockAddrArgs { + pub(crate) attach_type: Ident, + pub(crate) args: Args, +} + +impl Parse for SockAddrArgs { + fn parse(input: ParseStream) -> Result { + let attach_type: Ident = input.parse()?; + match attach_type.to_string().as_str() { + "connect4" | "connect6" | "bind4" | "bind6" | "getpeername4" | "getpeername6" + | "getsockname4" | "getsockname6" | "sendmsg4" | "sendmsg6" | "recvmsg4" + | "recvmsg6" => (), + _ => return Err(input.error("invalid attach type")), + } + let args = if input.parse::().is_ok() { + Args::parse(input)? + } else { + Args { args: vec![] } + }; + Ok(SockAddrArgs { attach_type, args }) + } +} + pub struct Map { item: ItemStatic, name: String, @@ -257,6 +280,7 @@ impl CgroupSkb { pub fn from_syn(mut args: Args, item: ItemFn) -> Result { let name = pop_arg(&mut args, "name"); let expected_attach_type = pop_arg(&mut args, "attach"); + err_on_unknown_args(&args)?; Ok(CgroupSkb { item, @@ -291,6 +315,44 @@ impl CgroupSkb { } } +pub struct CgroupSockAddr { + item: ItemFn, + attach_type: String, + name: Option, +} + +impl CgroupSockAddr { + pub fn from_syn(mut args: Args, item: ItemFn, attach_type: String) -> Result { + let name = pop_arg(&mut args, "name"); + err_on_unknown_args(&args)?; + + Ok(CgroupSockAddr { + item, + attach_type, + name, + }) + } + + pub fn expand(&self) -> Result { + let section_name = if let Some(name) = &self.name { + format!("cgroup/{}/{}", self.attach_type, name) + } else { + format!("cgroup/{}", self.attach_type) + }; + 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_addr) -> i32 { + return #fn_name(::aya_bpf::programs::SockAddrContext::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()), @@ -507,6 +569,8 @@ pub struct SkSkb { impl SkSkb { pub fn from_syn(kind: SkSkbKind, mut args: Args, item: ItemFn) -> Result { let name = pop_arg(&mut args, "name"); + err_on_unknown_args(&args)?; + Ok(SkSkb { item, kind, name }) } @@ -539,6 +603,7 @@ pub struct SocketFilter { impl SocketFilter { pub fn from_syn(mut args: Args, item: ItemFn) -> Result { let name = name_arg(&mut args)?; + err_on_unknown_args(&args)?; Ok(SocketFilter { item, name }) } diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs index 1512842d2..45e54130b 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, CgroupSysctl, FEntry, FExit, Lsm, Map, PerfEvent, Probe, - ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, - TracePoint, Xdp, + Args, BtfTracePoint, CgroupSkb, CgroupSockAddr, CgroupSysctl, FEntry, FExit, Lsm, Map, + PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, + SockAddrArgs, SockOps, SocketFilter, TracePoint, Xdp, }; use proc_macro::TokenStream; use syn::{parse_macro_input, ItemFn, ItemStatic}; @@ -105,6 +105,18 @@ pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream { .into() } +#[proc_macro_attribute] +pub fn cgroup_sock_addr(attrs: TokenStream, item: TokenStream) -> TokenStream { + let args = parse_macro_input!(attrs as SockAddrArgs); + let attach_type = args.attach_type.to_string(); + let item = parse_macro_input!(item as ItemFn); + + CgroupSockAddr::from_syn(args.args, item, attach_type) + .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 3c4348b6e..ab238c518 100644 --- a/bpf/aya-bpf/src/programs/mod.rs +++ b/bpf/aya-bpf/src/programs/mod.rs @@ -6,6 +6,7 @@ pub mod probe; pub mod raw_tracepoint; pub mod sk_buff; pub mod sk_msg; +pub mod sock_addr; pub mod sock_ops; pub mod sysctl; pub mod tp_btf; @@ -20,6 +21,7 @@ pub use probe::ProbeContext; pub use raw_tracepoint::RawTracePointContext; pub use sk_buff::SkBuffContext; pub use sk_msg::SkMsgContext; +pub use sock_addr::SockAddrContext; pub use sock_ops::SockOpsContext; pub use sysctl::SysctlContext; pub use tp_btf::BtfTracePointContext; diff --git a/bpf/aya-bpf/src/programs/sock_addr.rs b/bpf/aya-bpf/src/programs/sock_addr.rs new file mode 100644 index 000000000..294ff6360 --- /dev/null +++ b/bpf/aya-bpf/src/programs/sock_addr.rs @@ -0,0 +1,19 @@ +use core::ffi::c_void; + +use crate::{bindings::bpf_sock_addr, BpfContext}; + +pub struct SockAddrContext { + pub sock_addr: *mut bpf_sock_addr, +} + +impl SockAddrContext { + pub fn new(sock_addr: *mut bpf_sock_addr) -> SockAddrContext { + SockAddrContext { sock_addr } + } +} + +impl BpfContext for SockAddrContext { + fn as_ptr(&self) -> *mut c_void { + self.sock_addr as *mut _ + } +}