From 78a9f91cf35aaae9d6d7e3f5fde94a660fa474b7 Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@gmail.com>
Date: Thu, 10 Oct 2024 16:14:20 -0400
Subject: [PATCH] Add rustfmt to CI

---
 .github/workflows/ci.yml          |  2 +-
 rustfmt.toml                      |  3 +++
 test.sh                           |  1 +
 {{project-name}}-ebpf/src/main.rs |  1 -
 {{project-name}}/src/main.rs      | 20 +++++++++++---------
 5 files changed, 16 insertions(+), 11 deletions(-)
 create mode 100644 rustfmt.toml

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f4d40e9..0044b06 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -51,7 +51,7 @@ jobs:
 
       - uses: dtolnay/rust-toolchain@stable
         with:
-          components: clippy
+          components: clippy,rustfmt
 
       - uses: Swatinem/rust-cache@v2
 
diff --git a/rustfmt.toml b/rustfmt.toml
new file mode 100644
index 0000000..f7a9e5b
--- /dev/null
+++ b/rustfmt.toml
@@ -0,0 +1,3 @@
+group_imports = "StdExternalCrate"
+imports_granularity = "Crate"
+reorder_imports = true
diff --git a/test.sh b/test.sh
index 9e0354c..1d298b1 100755
--- a/test.sh
+++ b/test.sh
@@ -51,6 +51,7 @@ esac
 
 cargo generate --path "${TEMPLATE_DIR}" -n test -d program_type="${PROG_TYPE}" ${ADDITIONAL_ARGS}
 pushd test
+cargo fmt --all -- --check
 cargo build --package test
 cargo build --package test --release
 # We cannot run clippy over the whole workspace at once due to feature unification. Since both test
diff --git a/{{project-name}}-ebpf/src/main.rs b/{{project-name}}-ebpf/src/main.rs
index 9a0a61f..13822b6 100644
--- a/{{project-name}}-ebpf/src/main.rs
+++ b/{{project-name}}-ebpf/src/main.rs
@@ -135,7 +135,6 @@ use aya_ebpf::{
     programs::SkMsgContext,
 };
 use aya_log_ebpf::info;
-
 use {{crate_name}}_common::SockKey;
 
 #[map]
diff --git a/{{project-name}}/src/main.rs b/{{project-name}}/src/main.rs
index 6a2fe38..8b7e037 100644
--- a/{{project-name}}/src/main.rs
+++ b/{{project-name}}/src/main.rs
@@ -10,11 +10,10 @@ use aya::programs::UProbe;
 {%- when "sock_ops" -%}
 use aya::programs::{SockOps, links::CgroupAttachMode};
 {%- when "sk_msg" -%}
-use aya::maps::SockHash;
-use aya::programs::SkMsg;
+use aya::{maps::SockHash, programs::SkMsg};
 use {{crate_name}}_common::SockKey;
 {%- when "xdp" -%}
-use anyhow::Context;
+use anyhow::Context as _;
 use aya::programs::{Xdp, XdpFlags};
 {%- when "classifier" -%}
 use aya::programs::{tc, SchedClassifier, TcAttachType};
@@ -39,12 +38,12 @@ use aya::programs::SocketFilter;
 {%- when "raw_tracepoint" -%}
 use aya::programs::RawTracePoint;
 {%- endcase %}
-use aya::{include_bytes_aligned, Ebpf};
-use aya_log::EbpfLogger;
 {% if program_types_with_opts contains program_type -%}
 use clap::Parser;
 {% endif -%}
-use log::{info, warn, debug};
+
+#[rustfmt::skip]
+use log::{debug, info, warn};
 use tokio::signal;
 
 {% if program_types_with_opts contains program_type -%}
@@ -85,8 +84,11 @@ async fn main() -> Result<(), anyhow::Error> {
     // runtime. This approach is recommended for most real-world use cases. If you would
     // like to specify the eBPF program at runtime rather than at compile-time, you can
     // reach for `Bpf::load_file` instead.
-    let mut ebpf = Ebpf::load(include_bytes_aligned!(concat!(env!("OUT_DIR"), "/{{project-name}}")))?;
-    if let Err(e) = EbpfLogger::init(&mut ebpf) {
+    let mut ebpf = aya::Ebpf::load(aya::include_bytes_aligned!(concat!(
+        env!("OUT_DIR"),
+        "/test"
+    )))?;
+    if let Err(e) = aya_log::EbpfLogger::init(&mut ebpf) {
         // This can happen if you remove all log statements from your eBPF program.
         warn!("failed to initialize eBPF logger: {}", e);
     }
@@ -115,7 +117,7 @@ async fn main() -> Result<(), anyhow::Error> {
     program.load()?;
     program.attach(cgroup, CgroupAttachMode::default())?;
     {%- when "sk_msg" -%}
-    let sock_map: SockHash::<_, SockKey> = ebpf.map("{{sock_map}}").unwrap().try_into()?;
+    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()?;