From 417ab352cca9179ddee6a56e5d1c516f001d6bb7 Mon Sep 17 00:00:00 2001
From: Daiki Ueno <dueno@redhat.com>
Date: Tue, 21 Nov 2023 11:18:09 +0900
Subject: [PATCH] agent: Generate vmlinux.h on the fly in build.rs

This adds a logic to generate "vmlinux.h" in build.rs, needed for
compiling BPF programs.  If vmlinux.h is supplied by the makefile, it
will be copied to the build directory; otherwise bpftool will be used
to extract it from the running kernel.

Signed-off-by: Daiki Ueno <dueno@redhat.com>
---
 .github/workflows/rust.yml              |  4 ++-
 GNUmakefile                             |  4 +--
 agent/build.rs                          | 40 ++++++++++++++++++++-----
 agent/tests/agenttest/build.rs          | 40 ++++++++++++++++++++-----
 agent/tests/agenttest/src/bpf/vmlinux.h |  1 -
 5 files changed, 69 insertions(+), 20 deletions(-)
 delete mode 120000 agent/tests/agenttest/src/bpf/vmlinux.h

diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 1d33734..743d4dc 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -29,7 +29,9 @@ jobs:
     - name: Install dependencies
       run: dnf install -y cargo clang clippy kernel-devel libbpf-devel llvm-devel rustfmt
     - name: Copy vmlinux.h
-      run: cp $(rpm -ql kernel-devel | grep '/vmlinux.h$' | tail -1) agent/src/bpf
+      run: |
+        cp $(rpm -ql kernel-devel | grep '/vmlinux.h$' | tail -1) agent/src/bpf
+        cp $(rpm -ql kernel-devel | grep '/vmlinux.h$' | tail -1) agent/tests/agenttest/src/bpf
     - name: Build
       run: cargo build --verbose
     - name: Run tests
diff --git a/GNUmakefile b/GNUmakefile
index 4ab5f77..c6f76a6 100644
--- a/GNUmakefile
+++ b/GNUmakefile
@@ -28,10 +28,10 @@ conffiles = \
 .PHONY: all
 all: $(programs)
 
-agent/src/bpf/vmlinux.h:
+agent/src/bpf/vmlinux.h agent/tests/agenttest/src/bpf/vmlinux.h:
 	bpftool btf dump file /sys/kernel/btf/vmlinux format c > $@-t && mv $@-t $@
 
-$(programs): agent/src/bpf/vmlinux.h
+$(programs): agent/src/bpf/vmlinux.h agent/tests/agenttest/src/bpf/vmlinux.h
 	cargo build --target-dir="${TARGETDIR}" ${CARGO_ARGS}
 
 .PHONY: install-programs
diff --git a/agent/build.rs b/agent/build.rs
index 15a12ba..166bd28 100644
--- a/agent/build.rs
+++ b/agent/build.rs
@@ -1,17 +1,41 @@
 // SPDX-License-Identifier: GPL-2.0
 
 use libbpf_cargo::SkeletonBuilder;
-use std::{env, path::PathBuf};
-
-const SRC: &str = "src/bpf/audit.bpf.c";
+use std::{
+    env,
+    fs::{self, File},
+    path::PathBuf,
+    process::Command,
+};
 
 fn main() {
-    let mut out =
+    let builddir =
         PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR must be set in build script"));
-    out.push("audit.skel.rs");
+    let srcdir = PathBuf::from(
+        env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR must be set in build script"),
+    );
+
+    let vmlinux_h = srcdir.join("src").join("bpf").join("vmlinux.h");
+    if vmlinux_h.exists() {
+        fs::copy(&vmlinux_h, &builddir.join("vmlinux.h")).expect("unable to copy vmlinux.h");
+    } else {
+        let file = File::create(&builddir.join("vmlinux.h")).expect("unable to create vmlinux.h");
+        Command::new("bpftool")
+            .arg("btf")
+            .arg("dump")
+            .arg("file")
+            .arg("/sys/kernel/btf/vmlinux")
+            .arg("format")
+            .arg("c")
+            .stdout(file)
+            .status()
+            .expect("unable to run bpftool");
+    }
+    let src = srcdir.join("src").join("bpf").join("audit.bpf.c");
     SkeletonBuilder::new()
-        .source(SRC)
-        .build_and_generate(&out)
+        .source(&src)
+        .clang_args(&format!("-I{}", builddir.display()))
+        .build_and_generate(&builddir.join("audit.skel.rs"))
         .unwrap();
-    println!("cargo:rerun-if-changed={}", SRC);
+    println!("cargo:rerun-if-changed={}", src.display());
 }
diff --git a/agent/tests/agenttest/build.rs b/agent/tests/agenttest/build.rs
index ffbaf23..b1caf7a 100644
--- a/agent/tests/agenttest/build.rs
+++ b/agent/tests/agenttest/build.rs
@@ -1,17 +1,41 @@
 // SPDX-License-Identifier: GPL-2.0
 
 use libbpf_cargo::SkeletonBuilder;
-use std::{env, path::PathBuf};
-
-const SRC: &str = "src/bpf/agent.bpf.c";
+use std::{
+    env,
+    fs::{self, File},
+    path::PathBuf,
+    process::Command,
+};
 
 fn main() {
-    let mut out =
+    let builddir =
         PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR must be set in build script"));
-    out.push("agent.skel.rs");
+    let srcdir = PathBuf::from(
+        env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR must be set in build script"),
+    );
+
+    let vmlinux_h = srcdir.join("src").join("bpf").join("vmlinux.h");
+    if vmlinux_h.exists() {
+        fs::copy(&vmlinux_h, &builddir.join("vmlinux.h")).expect("unable to copy vmlinux.h");
+    } else {
+        let file = File::create(&builddir.join("vmlinux.h")).expect("unable to create vmlinux.h");
+        Command::new("bpftool")
+            .arg("btf")
+            .arg("dump")
+            .arg("file")
+            .arg("/sys/kernel/btf/vmlinux")
+            .arg("format")
+            .arg("c")
+            .stdout(file)
+            .status()
+            .expect("unable to run bpftool");
+    }
+    let src = srcdir.join("src").join("bpf").join("agent.bpf.c");
     SkeletonBuilder::new()
-        .source(SRC)
-        .build_and_generate(&out)
+        .source(&src)
+        .clang_args(&format!("-I{}", builddir.display()))
+        .build_and_generate(&builddir.join("agent.skel.rs"))
         .unwrap();
-    println!("cargo:rerun-if-changed={}", SRC);
+    println!("cargo:rerun-if-changed={}", src.display());
 }
diff --git a/agent/tests/agenttest/src/bpf/vmlinux.h b/agent/tests/agenttest/src/bpf/vmlinux.h
deleted file mode 120000
index 71c5aea..0000000
--- a/agent/tests/agenttest/src/bpf/vmlinux.h
+++ /dev/null
@@ -1 +0,0 @@
-../../../../src/bpf/vmlinux.h
\ No newline at end of file