From 6f51b616e127ba6dcaae05d6974024916c05a820 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Thu, 9 Jun 2022 00:20:08 +0100
Subject: [PATCH] bpf: Implement BPF_PROG_TYPE_CGROUP_SOCK

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf-macros/src/expand.rs | 45 ++++++++++++++++++++++++++++++++
 bpf/aya-bpf-macros/src/lib.rs    | 17 +++++++++---
 bpf/aya-bpf/src/programs/mod.rs  |  2 ++
 bpf/aya-bpf/src/programs/sock.rs | 19 ++++++++++++++
 4 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 bpf/aya-bpf/src/programs/sock.rs

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<String>,
+    name: Option<String>,
+}
+
+impl CgroupSock {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<CgroupSock> {
+        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<TokenStream> {
+        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<String> {
     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 _
+    }
+}