From fa037a88e2f0820d2a64bbaae12464bf5dce083d Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Sat, 20 Nov 2021 19:25:15 +0000
Subject: [PATCH 1/2] bpf: Fix cgroup_skb macro

This commit ensures that if no attach type is provided, that we use the
cgroup/skb section. If an attach type is provided we use the
cgroup_skb/$attach_type section.

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf-macros/src/expand.rs | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs
index d079323f..8e35eb74 100644
--- a/bpf/aya-bpf-macros/src/expand.rs
+++ b/bpf/aya-bpf-macros/src/expand.rs
@@ -215,14 +215,14 @@ impl SchedClassifier {
 
 pub struct CgroupSkb {
     item: ItemFn,
-    expected_attach_type: String,
+    expected_attach_type: Option<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());
+        let expected_attach_type = pop_arg(&mut args, "attach");
 
         Ok(CgroupSkb {
             item,
@@ -232,11 +232,16 @@ impl CgroupSkb {
     }
 
     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)
+        let section_name = if let Some(attach) = &self.expected_attach_type {
+            if let Some(name) = &self.name {
+                format!("cgroup_skb/{}/{}", attach, name)
+            } else {
+                format!("cgroup_skb/{}", attach)
+            }
+        } else if let Some(name) = &self.name {
+            format!("cgroup/skb/{}", name)
         } else {
-            format!("cgroup_skb/{}", attach)
+            ("cgroup/skb").to_owned()
         };
         let fn_name = &self.item.sig.ident;
         let item = &self.item;

From eadba79debe7fec69e2b28e73844a58e9f577731 Mon Sep 17 00:00:00 2001
From: Dave Tucker <dave@dtucker.co.uk>
Date: Sat, 20 Nov 2021 23:30:35 +0000
Subject: [PATCH 2/2] bpf: Add tests for cgroup_skb macro

This checks that the expand function produces the expected link_section
given certain input

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
---
 bpf/aya-bpf-macros/src/expand.rs | 75 ++++++++++++++++++++++++++++++++
 1 file changed, 75 insertions(+)

diff --git a/bpf/aya-bpf-macros/src/expand.rs b/bpf/aya-bpf-macros/src/expand.rs
index 8e35eb74..39791e7a 100644
--- a/bpf/aya-bpf-macros/src/expand.rs
+++ b/bpf/aya-bpf-macros/src/expand.rs
@@ -528,3 +528,78 @@ impl SocketFilter {
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use syn::parse_quote;
+
+    use super::*;
+
+    #[test]
+    fn cgroup_skb_with_attach_and_name() {
+        let prog = CgroupSkb::from_syn(
+            parse_quote!(name = "foo", attach = "ingress"),
+            parse_quote!(
+                fn foo(ctx: SkBuffContext) -> i32 {
+                    0
+                }
+            ),
+        )
+        .unwrap();
+        let stream = prog.expand().unwrap();
+        assert!(stream
+            .to_string()
+            .contains("[link_section = \"cgroup_skb/ingress/foo\"]"));
+    }
+
+    #[test]
+    fn cgroup_skb_with_name() {
+        let prog = CgroupSkb::from_syn(
+            parse_quote!(name = "foo"),
+            parse_quote!(
+                fn foo(ctx: SkBuffContext) -> i32 {
+                    0
+                }
+            ),
+        )
+        .unwrap();
+        let stream = prog.expand().unwrap();
+        assert!(stream
+            .to_string()
+            .contains("[link_section = \"cgroup/skb/foo\"]"));
+    }
+
+    #[test]
+    fn cgroup_skb_no_name() {
+        let prog = CgroupSkb::from_syn(
+            parse_quote!(),
+            parse_quote!(
+                fn foo(ctx: SkBuffContext) -> i32 {
+                    0
+                }
+            ),
+        )
+        .unwrap();
+        let stream = prog.expand().unwrap();
+        assert!(stream
+            .to_string()
+            .contains("[link_section = \"cgroup/skb\"]"));
+    }
+
+    #[test]
+    fn cgroup_skb_with_attach_no_name() {
+        let prog = CgroupSkb::from_syn(
+            parse_quote!(attach = "egress"),
+            parse_quote!(
+                fn foo(ctx: SkBuffContext) -> i32 {
+                    0
+                }
+            ),
+        )
+        .unwrap();
+        let stream = prog.expand().unwrap();
+        assert!(stream
+            .to_string()
+            .contains("[link_section = \"cgroup_skb/egress\"]"));
+    }
+}