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

This adds the `extract_vmlinux_header` method to `SkelBuilder`, which
lets libbpf-cargo auto-generate the "vmlinux.h" header from the
running kernel using bpftool.  This would be particularly useful when
packaging a crate with "cargo package", as it doesn't need the header
file to be added to the repository.

Signed-off-by: Daiki Ueno <[email protected]>

Author: ueno

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

@@ -242,11 +242,11 @@ fn compile(debug: bool, objs: &[UnprocessedObj], clang: &Path, target_dir: &Path
     Ok(())
 }
 
-fn extract_clang_or_default(clang: Option<&PathBuf>) -> PathBuf {
-    match clang {
+fn extract_command_or_default(command: Option<&PathBuf>, default: &str) -> PathBuf {
+    match command {
         Some(c) => c.into(),
         // Searches $PATH
-        None => "clang".into(),
+        None => default.into(),
     }
 }
 
@@ -269,7 +269,7 @@ pub fn build(
 
     check_progs(&to_compile)?;
 
-    let clang = extract_clang_or_default(clang);
+    let clang = extract_command_or_default(clang, "clang");
     check_clang(debug, &clang, skip_clang_version_checks)
         .with_context(|| anyhow!("{} is invalid", clang.display()))?;
     compile(debug, &to_compile, &clang, &target_dir).context("Failed to compile progs")?;
@@ -286,8 +286,10 @@ pub fn build_single(
     clang: Option<&PathBuf>,
     skip_clang_version_checks: bool,
     options: &str,
+    bpftool: Option<&PathBuf>,
+    vmlinux_h: Option<&PathBuf>,
 ) -> Result<()> {
-    let clang = extract_clang_or_default(clang);
+    let clang = extract_command_or_default(clang, "clang");
     check_clang(debug, &clang, skip_clang_version_checks)?;
     let header_parent_dir = tempdir()?;
     let header_dir = extract_libbpf_headers_to_disk(header_parent_dir.path())?;
@@ -297,6 +299,21 @@ pub fn build_single(
         options.to_string()
     };
 
+    if let Some(vmlinux_h) = vmlinux_h {
+        let file = fs::File::create(vmlinux_h)?;
+        let bpftool = extract_command_or_default(bpftool, "bpftool");
+        Command::new(bpftool)
+            .arg("btf")
+            .arg("dump")
+            .arg("file")
+            .arg("/sys/kernel/btf/vmlinux")
+            .arg("format")
+            .arg("c")
+            .stdout(file)
+            .status()?;
+        compiler_options += &format!(" -I{}", vmlinux_h.parent().unwrap().display());
+    }
+
     // 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();
+            // Unwrap is safe here since we guarantee that dir.is_some() above
+            let objfile = self.dir.as_ref().unwrap().path().join(format!("{name}.o"));
+            self.obj = Some(objfile);
+        }
+
+        if self.extract_vmlinux_header {
+            // Unwrap is safe here since we guarantee that dir.is_some() above
+            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_ref(),
+            self.vmlinux_h.as_ref(),
         )
         .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();