Merge pull request #105 from dave-tucker/socket_filter

bpf: Add macro for socket_filter programs

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

bpf/aya-bpf-macros/src/lib.rs

@@ -1,14 +1,12 @@
 mod expand;
 
 use expand::{
-    Args, BtfTracePoint, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier,
+    Args, BtfTracePoint, CgroupSkb, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint,
-    SkMsg, SkSkb, SkSkbKind, SockOps, TracePoint, Xdp,
+    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()
+}