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>( input_file: InputFile, types: &[T], additional_flags: &[T], ) -> Result { let additional_flags = additional_flags .iter() .map(|s| s.as_ref().into()) .collect::>(); let mut bindgen = bindgen::bpf_builder(); let (additional_flags, ctypes_prefix) = extract_ctypes_prefix(&additional_flags); if let Some(prefix) = ctypes_prefix { bindgen = bindgen.ctypes_prefix(prefix) } for ty in types { bindgen = bindgen.allowlist_type(ty); } 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(), ), }; 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); 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 { 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 extract_ctypes_prefix(s: &[String]) -> (Vec, Option) { if let Some(index) = s.iter().position(|el| el == "--ctypes-prefix") { if index < s.len() - 1 { let mut flags = Vec::new(); flags.extend_from_slice(&s[0..index]); // skip ["--ctypes-prefix", "value"] flags.extend_from_slice(&s[index + 2..]); return (flags, s.get(index + 1).cloned()); } } (s.to_vec(), None) } fn combine_flags(s1: &[String], s2: &[String]) -> Vec { let mut flags = Vec::new(); let mut extra = Vec::new(); for s in [s1, s2] { let mut s = s.splitn(2, |el| el == "--"); // append args flags.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() { flags.push("--".to_string()); flags.extend(extra); } flags } #[cfg(test)] mod test { use super::{combine_flags, extract_ctypes_prefix}; fn to_vec(s: &str) -> Vec { s.split(' ').map(|x| x.into()).collect() } #[test] fn test_extract_ctypes_prefix() { let (flags, prefix) = extract_ctypes_prefix(&to_vec("foo --ctypes-prefix bar baz")); assert_eq!(flags, to_vec("foo baz")); assert_eq!(prefix, Some("bar".to_string())); let (flags, prefix) = extract_ctypes_prefix(&to_vec("foo --ctypes-prefi bar baz")); assert_eq!(flags, to_vec("foo --ctypes-prefi bar baz")); assert_eq!(prefix, None); let (flags, prefix) = extract_ctypes_prefix(&to_vec("--foo bar --ctypes-prefix")); assert_eq!(flags, to_vec("--foo bar --ctypes-prefix")); assert_eq!(prefix, None); let (flags, prefix) = extract_ctypes_prefix(&to_vec("--ctypes-prefix foo")); let empty: Vec = Vec::new(); assert_eq!(flags, empty); assert_eq!(prefix, Some("foo".to_string())); let (flags, prefix) = extract_ctypes_prefix(&to_vec("--ctypes-prefix")); assert_eq!(flags, to_vec("--ctypes-prefix")); assert_eq!(prefix, None); } #[test] fn test_combine_flags() { 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(), ); } }