From 0c7e9b94ea5977db062538080c57db9cbbdf1845 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Wed, 10 Nov 2021 17:12:00 +0000
Subject: [PATCH] bpf: Add macros for sk_skb programs

This commit adds the stream_parser and stream_verdict macros for use in
bpf programs

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf-macros/src/expand.rs | 50 ++++++++++++++++++++++
 bpf/aya-bpf-macros/src/lib.rs    | 72 +++++++++++++++++++++++++++++++-
 2 files changed, 121 insertions(+), 1 deletion(-)

diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs
index 3ac63028..3d818129 100644
--- a/bpf/aya-bpf-macros/src/expand.rs
+++ b/bpf/aya-bpf-macros/src/expand.rs
@@ -441,3 +441,53 @@ impl BtfTracePoint {
         })
     }
 }
+
+#[allow(clippy::enum_variant_names)]
+#[derive(Debug, Copy, Clone)]
+pub enum SkSkbKind {
+    StreamVerdict,
+    StreamParser,
+}
+
+impl std::fmt::Display for SkSkbKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        use SkSkbKind::*;
+        match self {
+            StreamVerdict => write!(f, "stream_verdict"),
+            StreamParser => write!(f, "stream_parser"),
+        }
+    }
+}
+
+pub struct SkSkb {
+    kind: SkSkbKind,
+    item: ItemFn,
+    name: Option<String>,
+}
+
+impl SkSkb {
+    pub fn from_syn(kind: SkSkbKind, mut args: Args, item: ItemFn) -> Result<SkSkb> {
+        let name = pop_arg(&mut args, "name");
+        Ok(SkSkb { item, kind, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let kind = &self.kind;
+        let section_name = if let Some(name) = &self.name {
+            format!("sk_skb/{}/{}", kind, name)
+        } else {
+            format!("sk_skb/{}", kind)
+        };
+        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) -> u32 {
+                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 59eeaceb..55f013b6 100644
--- a/bpf/aya-bpf-macros/src/lib.rs
+++ b/bpf/aya-bpf-macros/src/lib.rs
@@ -2,7 +2,7 @@ mod expand;
 
 use expand::{
     Args, BtfTracePoint, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier,
-    SkMsg, SockOps, TracePoint, Xdp,
+    SkMsg, SkSkb, SkSkbKind, SockOps, TracePoint, Xdp,
 };
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
@@ -248,3 +248,73 @@ pub fn btf_tracepoint(attrs: TokenStream, item: TokenStream) -> TokenStream {
         .unwrap_or_else(|err| err.to_compile_error())
         .into()
 }
+
+/// Marks a function as a SK_SKB Stream Parser eBPF program that can be attached
+/// to a SockMap
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 4.14
+///
+/// # Examples
+///
+/// ```no_run
+/// use aya_bpf::{macros::stream_parser, programs::SkSkbContext};
+///
+///
+///#[stream_parser]
+///fn stream_parser(ctx: SkSkbContext) -> u32 {
+///    match { try_stream_parser(ctx) } {
+///        Ok(ret) => ret,
+///        Err(ret) => ret,
+///    }
+///}
+///
+///fn try_stream_parser(ctx: SkSkbContext) -> Result<u32, u32> {
+///    Ok(ctx.len())
+///}
+/// ```
+#[proc_macro_attribute]
+pub fn stream_parser(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    sk_skb(SkSkbKind::StreamParser, attrs, item)
+}
+
+/// Marks a function as a SK_SKB Stream Verdict eBPF program that can be attached
+/// to a SockMap
+///
+/// # Minimum kernel version
+///
+/// The minimum kernel version required to use this feature is 4.14
+///
+/// # Examples
+///
+/// ```no_run
+/// use aya_bpf::{macros::stream_verdict, programs::SkSkbContext, bindings::sk_action};
+///
+///
+///#[stream_verdict]
+///fn stream_verdict(ctx: SkSkbContext) -> u32 {
+///    match { try_stream_verdict(ctx) } {
+///        Ok(ret) => ret,
+///        Err(ret) => ret,
+///    }
+///}
+///
+///fn try_stream_verdict(_ctx: SkSkbContext) -> Result<u32, u32> {
+///    Ok(sk_action::SK_PASS)
+///}
+/// ```
+#[proc_macro_attribute]
+pub fn stream_verdict(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    sk_skb(SkSkbKind::StreamVerdict, attrs, item)
+}
+
+fn sk_skb(kind: SkSkbKind, attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    SkSkb::from_syn(kind, args, item)
+        .and_then(|u| u.expand())
+        .unwrap_or_else(|err| err.to_compile_error())
+        .into()
+}