libbpf-cargo: Support on-the-fly vmlinux.h extraction

libbpf-cargo/Cargo.toml

@@ -28,6 +28,7 @@ path = "src/lib.rs"
 # When turned on, link against system-installed libbpf instead of building
 # and linking against vendored libbpf sources
 novendor = ["libbpf-sys/novendor"]
+bpftool = []
 
 [dependencies]
 anyhow = "1.0.1"

libbpf-cargo/src/build.rs

@@ -279,13 +279,16 @@ pub fn build(
 
 // Only used in libbpf-cargo library
 #[allow(dead_code)]
+#[allow(clippy::too_many_arguments)]
 pub fn build_single(
     debug: bool,
     source: &Path,
     out: &Path,
     clang: Option<&PathBuf>,
     skip_clang_version_checks: bool,
     options: &str,
+    bpftool: Option<&Path>,
+    vmlinux_h: Option<&Path>,
 ) -> Result<()> {
     let clang = extract_clang_or_default(clang);
     check_clang(debug, &clang, skip_clang_version_checks)?;
@@ -297,6 +300,24 @@ pub fn build_single(
         options.to_string()
     };
 
+    if let Some(vmlinux_h) = vmlinux_h {
+        let file = fs::File::create(vmlinux_h)?;
+        let bpftool = bpftool.unwrap_or_else(|| Path::new("bpftool"));
+        Command::new(bpftool)
+            .arg("btf")
+            .arg("dump")
+            .arg("file")
+            .arg("/sys/kernel/btf/vmlinux")
+            .arg("format")
+            .arg("c")
+            .stdout(file)
+            .status()?;
+        compiler_options += &vmlinux_h
+            .parent()
+            .map(|dir| format!(" -I{}", dir.display()))
+            .ok_or_else(|| anyhow!("Failed to find include directory for vmlinux.h"))?;
+    }
+
     // Explicitly disable stack protector logic, which doesn't work with
     // BPF. See https://lkml.org/lkml/2020/2/21/1000.
     compiler_options += " -fno-stack-protector";

libbpf-cargo/src/lib.rs

@@ -120,6 +120,9 @@ pub struct SkeletonBuilder {
     clang: Option<PathBuf>,
     clang_args: String,
     skip_clang_version_check: bool,
+    bpftool: Option<PathBuf>,
+    extract_vmlinux_header: bool,
+    vmlinux_h: Option<PathBuf>,
     rustfmt: PathBuf,
     dir: Option<TempDir>,
 }
@@ -139,6 +142,9 @@ impl SkeletonBuilder {
             clang: None,
             clang_args: String::new(),
             skip_clang_version_check: false,
+            bpftool: None,
+            extract_vmlinux_header: false,
+            vmlinux_h: None,
             rustfmt: "rustfmt".into(),
             dir: None,
         }
@@ -202,6 +208,22 @@ impl SkeletonBuilder {
         self
     }
 
+    /// Specify which `bpftool` binary to use
+    ///
+    /// Default searchs `$PATH` for `bpftool`
+    pub fn bpftool<P: AsRef<Path>>(&mut self, bpftool: P) -> &mut SkeletonBuilder {
+        self.bpftool = Some(bpftool.as_ref().to_path_buf());
+        self
+    }
+
+    /// Tell whether to extract vmlinux.h from the running kernel
+    ///
+    /// Default is off
+    pub fn extract_vmlinux_header(&mut self, extract: bool) -> &mut SkeletonBuilder {
+        self.extract_vmlinux_header = extract;
+        self
+    }
+
     /// Specify which `rustfmt` binary to use
     ///
     /// Default searches `$PATH` for `rustfmt`
@@ -240,15 +262,25 @@ impl SkeletonBuilder {
             )));
         }
 
-        if self.obj.is_none() {
-            let name = filename.split('.').next().unwrap();
+        if self.obj.is_none() || self.extract_vmlinux_header {
             let dir = tempdir().map_err(|e| Error::Build(e.into()))?;
-            let objfile = dir.path().join(format!("{name}.o"));
-            self.obj = Some(objfile);
             // Hold onto tempdir so that it doesn't get deleted early
             self.dir = Some(dir);
         }
 
+        if self.obj.is_none() {
+            let name = filename.split('.').next().unwrap();
+            // SANITY: `dir` is guaranteed to be `Some` at this point.
+            let objfile = self.dir.as_ref().unwrap().path().join(format!("{name}.o"));
+            self.obj = Some(objfile);
+        }
+
+        if self.extract_vmlinux_header {
+            // SANITY: `dir` is guaranteed to be `Some` at this point.
+            let vmlinux_h = self.dir.as_ref().unwrap().path().join("vmlinux.h");
+            self.vmlinux_h = Some(vmlinux_h);
+        }
+
         build::build_single(
             self.debug,
             source,
@@ -257,6 +289,8 @@ impl SkeletonBuilder {
             self.clang.as_ref(),
             self.skip_clang_version_check,
             &self.clang_args,
+            self.bpftool.as_deref(),
+            self.vmlinux_h.as_deref(),
         )
         .map_err(Error::Build)?;
 

libbpf-cargo/src/test.rs

@@ -850,6 +850,47 @@ fn test_skeleton_builder_clang_opts() {
         .unwrap();
 }
 
+#[cfg(feature = "bpftool")]
+#[test]
+fn test_skeleton_builder_bpftool() {
+    let (_dir, proj_dir, _cargo_toml) = setup_temp_project();
+
+    // Add prog dir
+    create_dir(proj_dir.join("src/bpf")).expect("failed to create prog dir");
+
+    // Add a prog
+    let mut prog = OpenOptions::new()
+        .write(true)
+        .create(true)
+        .open(proj_dir.join("src/bpf/prog.bpf.c"))
+        .expect("failed to open prog.bpf.c");
+
+    write!(
+        prog,
+        r#"
+        #include "vmlinux.h"
+        "#,
+    )
+    .expect("failed to write prog.bpf.c");
+
+    let skel = NamedTempFile::new().unwrap();
+
+    // Should fail b/c vmlinux.h is not present
+    SkeletonBuilder::new()
+        .source(proj_dir.join("src/bpf/prog.bpf.c"))
+        .debug(true)
+        .build_and_generate(skel.path())
+        .unwrap_err();
+
+    // Should succeed b/c vmlinux.h is extracted
+    SkeletonBuilder::new()
+        .source(proj_dir.join("src/bpf/prog.bpf.c"))
+        .debug(true)
+        .extract_vmlinux_header(true)
+        .build_and_generate(skel.path())
+        .unwrap();
+}
+
 #[test]
 fn test_skeleton_builder_arrays_ptrs() {
     let (_dir, proj_dir, cargo_toml) = setup_temp_project();