bpf: Add support for BPF_PROG_TYPE_CGROUP_SOCK_ADDR

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>

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<SockAddrArgs> {
+        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::<Token![,]>().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<CgroupSkb> {
         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<String>,
+}
+
+impl CgroupSockAddr {
+    pub fn from_syn(mut args: Args, item: ItemFn, attach_type: String) -> Result<CgroupSockAddr> {
+        let name = pop_arg(&mut args, "name");
+        err_on_unknown_args(&args)?;
+
+        Ok(CgroupSockAddr {
+            item,
+            attach_type,
+            name,
+        })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        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<String> {
     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<SkSkb> {
         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<SocketFilter> {
         let name = name_arg(&mut args)?;
+        err_on_unknown_args(&args)?;
 
         Ok(SocketFilter { item, name })
     }

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,
+    Args, BtfTracePoint, CgroupSkb, CgroupSockAddr, CgroupSysctl, FEntry, FExit, Lsm, Map,
-    ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter,
+    PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkMsg, SkSkb, SkSkbKind,
-    TracePoint, Xdp,
+    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);

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;

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 _
+    }
+}