aya-gen cli rework (#297)

* update aya-gen from structopt to clap 3

* aya-gen: add --header option

* aya-gen: run bindgen as a child process

* aya-gen: add support for passing additional bindgen args

aya-gen/Cargo.toml

@@ -6,10 +6,11 @@ edition = "2018"
 
 [dependencies]
 bindgen = "0.60"
-structopt = {version = "0.3", default-features = false }
+clap = { version = "3", features = ["derive"] }
 anyhow = "1"
 thiserror = "1"
 syn = "1"
 quote = "1"
 proc-macro2 = "1"
 indexmap = "1.6"
+tempfile = "3"

aya-gen/src/bin/aya-gen.rs

@@ -1,21 +1,26 @@
-use aya_gen::btf_types;
+use aya_gen::generate::{generate, InputFile};
 
 use std::{path::PathBuf, process::exit};
 
-use structopt::StructOpt;
-#[derive(StructOpt)]
+use clap::Parser;
+
+#[derive(Parser)]
 pub struct Options {
-    #[structopt(subcommand)]
+    #[clap(subcommand)]
     command: Command,
 }
 
-#[derive(StructOpt)]
+#[derive(Parser)]
 enum Command {
-    #[structopt(name = "btf-types")]
-    BtfTypes {
-        #[structopt(long, default_value = "/sys/kernel/btf/vmlinux")]
+    #[clap(name = "generate")]
+    Generate {
+        #[clap(long, default_value = "/sys/kernel/btf/vmlinux")]
         btf: PathBuf,
+        #[clap(long, conflicts_with = "btf")]
+        header: Option<PathBuf>,
         names: Vec<String>,
+        #[clap(last = true)]
+        bindgen_args: Vec<String>,
     },
 }
 
@@ -27,10 +32,20 @@ fn main() {
 }
 
 fn try_main() -> Result<(), anyhow::Error> {
-    let opts = Options::from_args();
+    let opts = Options::parse();
     match opts.command {
-        Command::BtfTypes { btf, names } => {
-            let bindings = btf_types::generate(&btf, &names)?;
+        Command::Generate {
+            btf,
+            header,
+            names,
+            bindgen_args,
+        } => {
+            let bindings: String;
+            if let Some(header) = header {
+                bindings = generate(InputFile::Header(header), &names, &bindgen_args)?;
+            } else {
+                bindings = generate(InputFile::Btf(btf), &names, &bindgen_args)?;
+            }
             println!("{}", bindings);
         }
     };

aya-gen/src/btf_types.rs

@@ -1,53 +0,0 @@
-use std::{io, path::Path, process::Command, str::from_utf8};
-
-use thiserror::Error;
-
-use crate::bindgen;
-
-#[derive(Error, Debug)]
-pub enum Error {
-    #[error("error executing bpftool")]
-    BpfTool(#[source] io::Error),
-
-    #[error("{stderr}\nbpftool failed with exit code {code}")]
-    BpfToolExit { code: i32, stderr: String },
-
-    #[error("bindgen failed")]
-    Bindgen,
-
-    #[error("rustfmt failed")]
-    Rustfmt(#[source] io::Error),
-}
-
-pub fn generate<T: AsRef<str>>(btf_file: &Path, types: &[T]) -> Result<String, Error> {
-    let mut bindgen = bindgen::bpf_builder();
-
-    let c_header = c_header_from_btf(btf_file)?;
-    bindgen = bindgen.header_contents("kernel_types.h", &c_header);
-
-    for ty in types {
-        bindgen = bindgen.allowlist_type(ty);
-    }
-
-    let bindings = bindgen.generate().or(Err(Error::Bindgen))?.to_string();
-
-    Ok(bindings)
-}
-
-fn c_header_from_btf(path: &Path) -> Result<String, Error> {
-    let output = Command::new("bpftool")
-        .args(&["btf", "dump", "file"])
-        .arg(path)
-        .args(&["format", "c"])
-        .output()
-        .map_err(Error::BpfTool)?;
-
-    if !output.status.success() {
-        return Err(Error::BpfToolExit {
-            code: output.status.code().unwrap(),
-            stderr: from_utf8(&output.stderr).unwrap().to_owned(),
-        });
-    }
-
-    Ok(from_utf8(&output.stdout).unwrap().to_owned())
-}

aya-gen/src/generate.rs

@@ -0,0 +1,170 @@
+use std::{
+    fs::{self, File},
+    io::{self, Write},
+    path::{Path, PathBuf},
+    process::Command,
+    str,
+};
+
+use tempfile::tempdir;
+
+use thiserror::Error;
+
+use crate::bindgen;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("error executing bpftool")]
+    BpfTool(#[source] io::Error),
+
+    #[error("{stderr}\nbpftool failed with exit code {code}")]
+    BpfToolExit { code: i32, stderr: String },
+
+    #[error("bindgen failed")]
+    Bindgen(#[source] io::Error),
+
+    #[error("{stderr}\nbindgen failed with exit code {code}")]
+    BindgenExit { code: i32, stderr: String },
+
+    #[error("rustfmt failed")]
+    Rustfmt(#[source] io::Error),
+
+    #[error("error reading header file")]
+    ReadHeaderFile(#[source] io::Error),
+}
+
+pub enum InputFile {
+    Btf(PathBuf),
+    Header(PathBuf),
+}
+
+pub fn generate<T: AsRef<str>>(
+    input_file: InputFile,
+    types: &[T],
+    additional_flags: &[T],
+) -> Result<String, Error> {
+    let mut bindgen = bindgen::bpf_builder();
+
+    let (c_header, name) = match &input_file {
+        InputFile::Btf(path) => (c_header_from_btf(path)?, "kernel_types.h"),
+        InputFile::Header(header) => (
+            fs::read_to_string(&header).map_err(Error::ReadHeaderFile)?,
+            header.file_name().unwrap().to_str().unwrap(),
+        ),
+    };
+
+    for ty in types {
+        bindgen = bindgen.allowlist_type(ty);
+    }
+
+    let dir = tempdir().unwrap();
+    let file_path = dir.path().join(name);
+    let mut file = File::create(&file_path).unwrap();
+    let _ = file.write(c_header.as_bytes()).unwrap();
+
+    let flags = combine_flags(
+        &bindgen.command_line_flags(),
+        &additional_flags
+            .iter()
+            .map(|s| s.as_ref().into())
+            .collect::<Vec<_>>(),
+    );
+
+    let output = Command::new("bindgen")
+        .arg(file_path)
+        .args(&flags)
+        .output()
+        .map_err(Error::Bindgen)?;
+
+    if !output.status.success() {
+        return Err(Error::BindgenExit {
+            code: output.status.code().unwrap(),
+            stderr: str::from_utf8(&output.stderr).unwrap().to_owned(),
+        });
+    }
+
+    Ok(str::from_utf8(&output.stdout).unwrap().to_owned())
+}
+
+fn c_header_from_btf(path: &Path) -> Result<String, Error> {
+    let output = Command::new("bpftool")
+        .args(&["btf", "dump", "file"])
+        .arg(path)
+        .args(&["format", "c"])
+        .output()
+        .map_err(Error::BpfTool)?;
+
+    if !output.status.success() {
+        return Err(Error::BpfToolExit {
+            code: output.status.code().unwrap(),
+            stderr: str::from_utf8(&output.stderr).unwrap().to_owned(),
+        });
+    }
+
+    Ok(str::from_utf8(&output.stdout).unwrap().to_owned())
+}
+
+fn combine_flags(s1: &[String], s2: &[String]) -> Vec<String> {
+    let mut args = Vec::new();
+    let mut extra = Vec::new();
+
+    for s in [s1, s2] {
+        let mut s = s.splitn(2, |el| el == "--");
+        // append args
+        args.extend(s.next().unwrap().iter().cloned());
+        if let Some(e) = s.next() {
+            // append extra args
+            extra.extend(e.iter().cloned());
+        }
+    }
+
+    // append extra args
+    if !extra.is_empty() {
+        args.push("--".to_string());
+        args.extend(extra);
+    }
+
+    args
+}
+
+#[cfg(test)]
+mod test {
+    use super::combine_flags;
+
+    fn to_vec(s: &str) -> Vec<String> {
+        s.split(" ").map(|x| x.into()).collect()
+    }
+
+    #[test]
+    fn combine_arguments_test() {
+        assert_eq!(
+            combine_flags(&to_vec("a b"), &to_vec("c d"),).join(" "),
+            "a b c d".to_string(),
+        );
+
+        assert_eq!(
+            combine_flags(&to_vec("a -- b"), &to_vec("a b"),).join(" "),
+            "a a b -- b".to_string(),
+        );
+
+        assert_eq!(
+            combine_flags(&to_vec("a -- b"), &to_vec("c d"),).join(" "),
+            "a c d -- b".to_string(),
+        );
+
+        assert_eq!(
+            combine_flags(&to_vec("a b"), &to_vec("c -- d"),).join(" "),
+            "a b c -- d".to_string(),
+        );
+
+        assert_eq!(
+            combine_flags(&to_vec("a -- b"), &to_vec("c -- d"),).join(" "),
+            "a c -- b d".to_string(),
+        );
+
+        assert_eq!(
+            combine_flags(&to_vec("a -- b"), &to_vec("-- c d"),).join(" "),
+            "a -- b c d".to_string(),
+        );
+    }
+}

aya-gen/src/lib.rs

@@ -5,7 +5,7 @@ use std::{
 };
 
 pub mod bindgen;
-pub mod btf_types;
+pub mod generate;
 pub mod rustfmt;
 
 pub fn write_to_file<T: AsRef<Path>>(path: T, code: &str) -> Result<(), io::Error> {