bpf: add support for BPF_PROG_TYPE_CGROUP_SKB programs

Example:

fn cgroup_skb_egress(skb: SkSkbContext) -> i32 {
    // allow data to go through
    1
}

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

@@ -43,8 +43,8 @@ pub struct Map {
 }
 
 impl Map {
-    pub fn from_syn(args: Args, item: ItemStatic) -> Result<Map> {
-        let name = name_arg(&args)?.unwrap_or_else(|| item.ident.to_string());
+    pub fn from_syn(mut args: Args, item: ItemStatic) -> Result<Map> {
+        let name = name_arg(&mut args)?.unwrap_or_else(|| item.ident.to_string());
         Ok(Map { item, name })
     }
 
@@ -66,8 +66,8 @@ pub struct Probe {
 }
 
 impl Probe {
-    pub fn from_syn(kind: ProbeKind, args: Args, item: ItemFn) -> Result<Probe> {
-        let name = name_arg(&args)?.unwrap_or_else(|| item.sig.ident.to_string());
+    pub fn from_syn(kind: ProbeKind, mut args: Args, item: ItemFn) -> Result<Probe> {
+        let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string());
 
         Ok(Probe { kind, item, name })
     }
@@ -95,8 +95,8 @@ pub struct SockOps {
 }
 
 impl SockOps {
-    pub fn from_syn(args: Args, item: ItemFn) -> Result<SockOps> {
-        let name = name_arg(&args)?;
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SockOps> {
+        let name = name_arg(&mut args)?;
 
         Ok(SockOps { item, name })
     }
@@ -127,8 +127,8 @@ pub struct SkMsg {
 }
 
 impl SkMsg {
-    pub fn from_syn(args: Args, item: ItemFn) -> Result<SkMsg> {
-        let name = name_arg(&args)?.unwrap_or_else(|| item.sig.ident.to_string());
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SkMsg> {
+        let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string());
 
         Ok(SkMsg { item, name })
     }
@@ -155,8 +155,8 @@ pub struct Xdp {
 }
 
 impl Xdp {
-    pub fn from_syn(args: Args, item: ItemFn) -> Result<Xdp> {
-        let name = name_arg(&args)?;
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<Xdp> {
+        let name = name_arg(&mut args)?;
 
         Ok(Xdp { item, name })
     }
@@ -187,8 +187,8 @@ pub struct SchedClassifier {
 }
 
 impl SchedClassifier {
-    pub fn from_syn(args: Args, item: ItemFn) -> Result<SchedClassifier> {
-        let name = name_arg(&args)?;
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<SchedClassifier> {
+        let name = name_arg(&mut args)?;
 
         Ok(SchedClassifier { item, name })
     }
@@ -213,16 +213,65 @@ impl SchedClassifier {
     }
 }
 
-fn name_arg(args: &Args) -> Result<Option<String>> {
-    for arg in &args.args {
-        if arg.name == "name" {
-            return Ok(Some(arg.value.value()));
+pub struct CgroupSkb {
+    item: ItemFn,
+    expected_attach_type: String,
+    name: Option<String>,
+}
+
+impl CgroupSkb {
+    pub fn from_syn(mut args: Args, item: ItemFn) -> Result<CgroupSkb> {
+        let name = pop_arg(&mut args, "name");
+        let expected_attach_type = pop_arg(&mut args, "attach").unwrap_or_else(|| "skb".to_owned());
+
+        Ok(CgroupSkb {
+            item,
+            name,
+            expected_attach_type,
+        })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        let attach = &self.expected_attach_type;
+        let section_name = if let Some(name) = &self.name {
+            format!("cgroup_skb/{}/{}", attach, name)
         } else {
-            return Err(Error::new_spanned(&arg.name, "invalid argument"));
-        }
+            format!("cgroup_skb/{}", attach)
+        };
+        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 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()),
+        None => None,
+    }
+}
+
+fn err_on_unknown_args(args: &Args) -> Result<()> {
+    for arg in &args.args {
+        return Err(Error::new_spanned(&arg.name, "invalid argument"));
+    }
+
+    Ok(())
+}
+
+fn name_arg(args: &mut Args) -> Result<Option<String>> {
+    let name = pop_arg(args, "name");
+    err_on_unknown_args(args)?;
 
-    Ok(None)
+    Ok(name)
 }
 
 #[derive(Debug, Copy, Clone)]

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

@@ -4,6 +4,8 @@ use expand::{Args, Map, Probe, ProbeKind, SchedClassifier, SkMsg, SockOps, 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);
@@ -79,6 +81,17 @@ pub fn classifier(attrs: TokenStream, item: TokenStream) -> TokenStream {
         .into()
 }
 
+#[proc_macro_attribute]
+pub fn cgroup_skb(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let args = parse_macro_input!(attrs as Args);
+    let item = parse_macro_input!(item as ItemFn);
+
+    CgroupSkb::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);

bpf/aya-bpf/src/programs/sk_skb.rs

@@ -3,7 +3,9 @@ use core::{
     mem::{self, MaybeUninit},
 };
 
-use aya_bpf_bindings::helpers::{bpf_l3_csum_replace, bpf_l4_csum_replace, bpf_skb_load_bytes, bpf_skb_store_bytes};
+use aya_bpf_bindings::helpers::{
+    bpf_l3_csum_replace, bpf_l4_csum_replace, bpf_skb_load_bytes, bpf_skb_store_bytes,
+};
 use aya_bpf_cty::c_long;
 
 use crate::{bindings::__sk_buff, BpfContext};
@@ -17,6 +19,11 @@ impl SkSkbContext {
         SkSkbContext { skb }
     }
 
+    #[inline]
+    pub fn len(&self) -> u32 {
+        unsafe { *self.skb }.len
+    }
+
     #[inline]
     pub fn set_mark(&mut self, mark: u32) {
         unsafe { *self.skb }.mark = mark;