From afcc5dc662e7efd9a9c99992067eb88abb47f426 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alessandro=C2=A0Decina?= <alessandro.d@gmail.com>
Date: Mon, 3 May 2021 06:48:44 +0000
Subject: [PATCH] bpf: add support for BPF_PROG_TYPE_SCHED_CLS programs

---
 bpf/aya-bpf-bindings/include/bindings.h |  2 ++
 bpf/aya-bpf-bindings/src/lib.rs         | 23 ++++++++++++++----
 bpf/aya-bpf-macros/src/expand.rs        | 32 +++++++++++++++++++++++++
 bpf/aya-bpf-macros/src/lib.rs           | 13 +++++++++-
 4 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/bpf/aya-bpf-bindings/include/bindings.h b/bpf/aya-bpf-bindings/include/bindings.h
index b07d77dc..6f91bdb6 100644
--- a/bpf/aya-bpf-bindings/include/bindings.h
+++ b/bpf/aya-bpf-bindings/include/bindings.h
@@ -1,4 +1,6 @@
 #include <linux/types.h>
 #include <linux/bpf.h>
 #include <linux/ptrace.h>
+// needed for TC_ACT_*
+#include <linux/pkt_cls.h>
 #include "bpf_helpers.h"
\ No newline at end of file
diff --git a/bpf/aya-bpf-bindings/src/lib.rs b/bpf/aya-bpf-bindings/src/lib.rs
index 6ae2a904..9d30ed30 100644
--- a/bpf/aya-bpf-bindings/src/lib.rs
+++ b/bpf/aya-bpf-bindings/src/lib.rs
@@ -7,11 +7,14 @@ mod x86_64;
 #[cfg(bpf_target_arch = "aarch64")]
 mod aarch64;
 
-#[cfg(bpf_target_arch = "x86_64")]
-pub use x86_64::*;
+mod gen {
+    #[cfg(bpf_target_arch = "x86_64")]
+    pub use super::x86_64::*;
 
-#[cfg(bpf_target_arch = "aarch64")]
-pub use aarch64::*;
+    #[cfg(bpf_target_arch = "aarch64")]
+    pub use super::aarch64::*;
+}
+pub use gen::*;
 
 use aya_bpf_cty::{c_long, c_void};
 use core::mem::{self, MaybeUninit};
@@ -30,3 +33,15 @@ unsafe fn bpf_probe_read<T>(src: *const T) -> Result<T, c_long> {
 
     Ok(v.assume_init())
 }
+
+pub const TC_ACT_OK: i32 = gen::bindings::TC_ACT_OK as i32;
+pub const TC_ACT_RECLASSIFY: i32 = gen::bindings::TC_ACT_RECLASSIFY as i32;
+pub const TC_ACT_SHOT: i32 = gen::bindings::TC_ACT_SHOT as i32;
+pub const TC_ACT_PIPE: i32 = gen::bindings::TC_ACT_PIPE as i32;
+pub const TC_ACT_STOLEN: i32 = gen::bindings::TC_ACT_STOLEN as i32;
+pub const TC_ACT_QUEUED: i32 = gen::bindings::TC_ACT_QUEUED as i32;
+pub const TC_ACT_REPEAT: i32 = gen::bindings::TC_ACT_REPEAT as i32;
+pub const TC_ACT_REDIRECT: i32 = gen::bindings::TC_ACT_REDIRECT as i32;
+pub const TC_ACT_TRAP: i32 = gen::bindings::TC_ACT_TRAP as i32;
+pub const TC_ACT_VALUE_MAX: i32 = gen::bindings::TC_ACT_VALUE_MAX as i32;
+pub const TC_ACT_EXT_VAL_MASK: i32 = 268435455;
diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs
index ba627605..fb58d4a8 100644
--- a/bpf/aya-bpf-macros/src/expand.rs
+++ b/bpf/aya-bpf-macros/src/expand.rs
@@ -181,6 +181,38 @@ impl Xdp {
     }
 }
 
+pub struct SchedClassifier {
+    item: ItemFn,
+    name: Option<String>,
+}
+
+impl SchedClassifier {
+    pub fn from_syn(args: Args, item: ItemFn) -> Result<SchedClassifier> {
+        let name = name_arg(&args)?;
+
+        Ok(SchedClassifier { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let section_name = if let Some(name) = &self.name {
+            format!("classifier/{}", name)
+        } else {
+            "classifier".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) -> i32 {
+                return #fn_name(::aya_bpf::programs::SkSkbContext::new(ctx));
+
+                #item
+            }
+        })
+    }
+}
+
 fn name_arg(args: &Args) -> Result<Option<String>> {
     for arg in &args.args {
         if arg.name == "name" {
diff --git a/bpf/aya-bpf-macros/src/lib.rs b/bpf/aya-bpf-macros/src/lib.rs
index 906ec358..1a5a8b5f 100644
--- a/bpf/aya-bpf-macros/src/lib.rs
+++ b/bpf/aya-bpf-macros/src/lib.rs
@@ -1,6 +1,6 @@
 mod expand;
 
-use expand::{Args, Map, Probe, ProbeKind, SkMsg, SockOps, Xdp};
+use expand::{Args, Map, Probe, ProbeKind, SchedClassifier, SkMsg, SockOps, Xdp};
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, ItemFn, ItemStatic};
 
@@ -68,6 +68,17 @@ pub fn xdp(attrs: TokenStream, item: TokenStream) -> TokenStream {
         .into()
 }
 
+#[proc_macro_attribute]
+pub fn classifier(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    SchedClassifier::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);