From a4bb29a3cc9292c9406b084725202a9fd6cd622e Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Thu, 11 Nov 2021 15:15:51 +0000
Subject: [PATCH] bpf: Add macro for socket_filter programs

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf-macros/src/expand.rs | 32 ++++++++++++++++++++++++++++++
 bpf/aya-bpf-macros/src/lib.rs    | 34 ++++++++++++++++++++++++++++----
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs
index 3d818129..6861bff7 100644
--- a/bpf/aya-bpf-macros/src/expand.rs
+++ b/bpf/aya-bpf-macros/src/expand.rs
@@ -491,3 +491,35 @@ impl SkSkb {
         })
     }
 }
+
+pub struct SocketFilter {
+    item: ItemFn,
+    name: Option<String>,
+}
+
+impl SocketFilter {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SocketFilter> {
+        let name = name_arg(&mut args)?;
+
+        Ok(SocketFilter { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = if let Some(name) = &self.name {
+            format!("socket/{}", name)
+        } else {
+            "socket".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::__sk_buff) -> i64 {
+                return #fn_name(::aya_bpf::programs::SkSkbContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs
index 55f013b6..bfa0b855 100644
--- a/bpf/aya-bpf-macros/src/lib.rs
+++ b/bpf/aya-bpf-macros/src/lib.rs
@@ -1,14 +1,12 @@
 mod expand;
 
 use expand::{
-    Args, BtfTracePoint, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier,
-    SkMsg, SkSkb, SkSkbKind, SockOps, TracePoint, Xdp,
+    Args, BtfTracePoint, CgroupSkb, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint,
+    SchedClassifier, SkMsg, SkSkb, SkSkbKind, SockOps, SocketFilter, TracePoint, Xdp,
 };
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
 
-use crate::expand::CgroupSkb;
-
 #[proc_macro_attribute]
 pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream {
     let args = parse_macro_input!(attrs as Args);
@@ -318,3 +316,31 @@ fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+/// Marks a function as a eBPF Socket Filter program that can be attached to
+/// a socket.
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 3.19
+///
+/// # Examples
+///
+/// ```no_run
+/// use aya_bpf::{macros::socket_filter, programs::SkSkbContext};
+///
+/// #[socket_filter(name = "accept_all")]
+/// pub fn accept_all(_ctx: SkSkbContext) -> i64 {
+///     return 0
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    SocketFilter::from_syn(args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}