From 263f115fdcd7a257757d4da64ee26380ead31b7b Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@gmail.com>
Date: Sun, 13 Oct 2024 11:36:55 -0400
Subject: [PATCH] Use destructuring

---
 {{project-name}}/src/main.rs | 95 +++++++++++++++++++-----------------
 1 file changed, 49 insertions(+), 46 deletions(-)

diff --git a/{{project-name}}/src/main.rs b/{{project-name}}/src/main.rs
index b97cd5d..fe29ac7 100644
--- a/{{project-name}}/src/main.rs
+++ b/{{project-name}}/src/main.rs
@@ -1,4 +1,4 @@
-{% case program_type -%}
+{%- case program_type -%}
 {%- when "kprobe", "kretprobe" -%}
 use aya::programs::KProbe;
 {%- when "fentry" -%}
@@ -50,16 +50,17 @@ use tokio::signal;
 {% if program_types_with_opts contains program_type -%}
 #[derive(Debug, Parser)]
 struct Opt {
-{%- if program_type == "xdp" or program_type == "classifier" %}
+{%- case program_type -%}
+{%- when "xdp", "classifier" %}
     #[clap(short, long, default_value = "eth0")]
     iface: String,
-{% elsif program_type == "sock_ops" or program_type == "cgroup_skb" or program_type == "cgroup_sysctl" or program_type == "cgroup_sockopt" %}
+{%- when "sock_ops", "cgroup_skb", "cgroup_sysctl", "cgroup_sockopt" %}
     #[clap(short, long, default_value = "/sys/fs/cgroup/unified")]
-    cgroup_path: String,
-{% elsif program_type == "uprobe" or program_type == "uretprobe" %}
+    cgroup_path: std::path::PathBuf,
+{%- when "uprobe", "uretprobe" %}
     #[clap(short, long)]
     pid: Option<i32>,
-{% endif -%}
+{%- endcase %}
 }
 
 {% endif -%}
@@ -93,89 +94,91 @@ async fn main() -> anyhow::Result<()> {
         // This can happen if you remove all log statements from your eBPF program.
         warn!("failed to initialize eBPF logger: {}", e);
     }
-    {% case program_type -%}
-    {%- when "kprobe", "kretprobe" -%}
+    {%- case program_type -%}
+    {%- when "kprobe", "kretprobe" %}
     let program: &mut KProbe = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
     program.attach("{{kprobe}}", 0)?;
-    {%- when "fentry" -%}
+    {%- when "fentry" %}
     let btf = Btf::from_sys_fs()?;
     let program: &mut FEntry = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load("{{fn_name}}", &btf)?;
     program.attach()?;
-    {%- when "fexit" -%}
+    {%- when "fexit" %}
     let btf = Btf::from_sys_fs()?;
     let program: &mut FExit = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load("{{fn_name}}", &btf)?;
     program.attach()?;
-    {%- when "uprobe", "uretprobe" -%}
+    {%- when "uprobe", "uretprobe" %}
+    let Opt { pid } = opt;
     let program: &mut UProbe = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
-    program.attach(Some("{{uprobe_fn_name}}"), 0, "{{uprobe_target}}", opt.pid)?;
-    {%- when "sock_ops" -%}
+    program.attach(Some("{{uprobe_fn_name}}"), 0, "{{uprobe_target}}", pid)?;
+    {%- when "sock_ops", "cgroup_skb", "cgroup_sysctl", "cgroup_sockopt" %}
+    let Opt { cgroup_path } = opt;
+    let cgroup = std::fs::File::open(&cgroup_path);
+    {%- if program_type == "sock_ops" %}
     let program: &mut SockOps = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
-    let cgroup = std::fs::File::open(opt.cgroup_path)?;
     program.load()?;
     program.attach(cgroup, CgroupAttachMode::default())?;
-    {%- when "sk_msg" -%}
+    {%- elsif program_type == "cgroup_skb" %}
+    let program: &mut CgroupSkb = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
+    program.load()?;
+    program.attach(
+        cgroup,
+        CgroupSkbAttachType::{{direction}},
+        CgroupAttachMode::default(),
+    )?;
+    {%- elsif program_type == "cgroup_sysctl" %}
+    let program: &mut CgroupSysctl = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
+    program.load()?;
+    program.attach(cgroup, CgroupAttachMode::default())?;
+    {%- elsif program_type == "cgroup_sockopt" %}
+    let program: &mut CgroupSockopt = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
+    program.load()?;
+    program.attach(cgroup, CgroupAttachMode::default())?;
+    {%- endif -%}
+    {%- when "sk_msg" %}
     let sock_map: SockHash<_, SockKey> = ebpf.map("{{sock_map}}").unwrap().try_into()?;
     let map_fd = sock_map.fd().try_clone()?;
-
     let prog: &mut SkMsg = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     prog.load()?;
     prog.attach(&map_fd)?;
     // insert sockets to the map using sock_map.insert here, or from a sock_ops program
-    {%- when "xdp" -%}
+    {%- when "xdp" %}
+    let Opt { iface } = opt;
     let program: &mut Xdp = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
-    program.attach(&opt.iface, XdpFlags::default())
+    program.attach(&iface, XdpFlags::default())
         .context("failed to attach the XDP program with default flags - try changing XdpFlags::default() to XdpFlags::SKB_MODE")?;
-    {%- when "classifier" -%}
+    {%- when "classifier" %}
+    let Opt { iface } = opt;
     // error adding clsact to the interface if it is already added is harmless
     // the full cleanup can be done with 'sudo tc qdisc del dev eth0 clsact'.
-    let _ = tc::qdisc_add_clsact(&opt.iface);
+    let _ = tc::qdisc_add_clsact(&iface);
     let program: &mut SchedClassifier = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
-    program.attach(&opt.iface, TcAttachType::{{direction}})?;
-    {%- when "cgroup_skb" -%}
-    let program: &mut CgroupSkb = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
-    let cgroup = std::fs::File::open(opt.cgroup_path)?;
-    program.load()?;
-    program.attach(
-        cgroup,
-        CgroupSkbAttachType::{{direction}},
-        CgroupAttachMode::default(),
-    )?;
-    {%- when "tracepoint" -%}
+    program.attach(&iface, TcAttachType::{{direction}})?;
+    {%- when "tracepoint" %}
     let program: &mut TracePoint = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
     program.attach("{{tracepoint_category}}", "{{tracepoint_name}}")?;
-    {%- when "lsm" -%}
+    {%- when "lsm" %}
     let btf = Btf::from_sys_fs()?;
     let program: &mut Lsm = ebpf.program_mut("{{lsm_hook}}").unwrap().try_into()?;
     program.load("{{lsm_hook}}", &btf)?;
     program.attach()?;
-    {%- when "tp_btf" -%}
+    {%- when "tp_btf" %}
     let btf = Btf::from_sys_fs()?;
     let program: &mut BtfTracePoint = ebpf.program_mut("{{tracepoint_name}}").unwrap().try_into()?;
     program.load("{{tracepoint_name}}", &btf)?;
     program.attach()?;
-    {%- when "socket_filter" -%}
+    {%- when "socket_filter" %}
     let listener = std::net::TcpListener::bind("localhost:0")?;
     let prog: &mut SocketFilter = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     prog.load()?;
     prog.attach(&listener)?;
-    {%- when "cgroup_sysctl" -%}
-    let program: &mut CgroupSysctl = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
-    let cgroup = std::fs::File::open(opt.cgroup_path)?;
-    program.load()?;
-    program.attach(cgroup, CgroupAttachMode::default())?;
-    {%- when "cgroup_sockopt" -%}
-    let program: &mut CgroupSockopt = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
-    let cgroup = std::fs::File::open(opt.cgroup_path)?;
-    program.load()?;
-    program.attach(cgroup, CgroupAttachMode::default())?;
-    {%- when "perf_event" -%}
+    {%- when "perf_event" %}
     // This will raise scheduled events on each CPU at 1 HZ, triggered by the kernel based
     // on clock ticks.
     let program: &mut PerfEvent = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
@@ -189,7 +192,7 @@ async fn main() -> anyhow::Result<()> {
             true,
         )?;
     }
-    {%- when "raw_tracepoint" -%}
+    {%- when "raw_tracepoint" %}
     let program: &mut RawTracePoint = ebpf.program_mut("{{crate_name}}").unwrap().try_into()?;
     program.load()?;
     program.attach("{{tracepoint_name}}")?;