use anyhow::{bail, Result};
use bindgen::{
self,
callbacks::{EnumVariantCustomBehavior, EnumVariantValue, ParseCallbacks},
Builder,
};
use quote::quote;
use std::env;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::Command;
use syn::visit::Visit;
use bpf_sys::headers::{get_custom_header_path, get_custom_header_version};
use bpf_sys::type_gen::get_custom_vmlinux_path;
use rhdxmr_publish_test_lib::bindgen as bpf_bindgen;
use syn::{
self, parse_str, punctuated::Punctuated, token::Comma, AngleBracketedGenericArguments,
ForeignItemStatic, GenericArgument, Ident, PathArguments::AngleBracketed, Type,
};
use tracing::{debug, warn, Level};
use tracing_subscriber::FmtSubscriber;
fn create_module(path: PathBuf, name: &str, bindings: &str) -> io::Result<()> {
{
let mut file = File::create(&path)?;
writeln!(
&mut file,
r"
mod {name} {{
#![allow(non_camel_case_types)]
#![allow(non_upper_case_globals)]
#![allow(non_snake_case)]
#![allow(unused_unsafe)]
#![allow(clippy::all)]
{bindings}
}}
pub use {name}::*;
",
name = name,
bindings = bindings
)?;
}
let _ = Command::new("rustfmt")
.arg("--edition=2018")
.arg("--emit=files")
.arg(path.to_str().unwrap())
.status();
Ok(())
}
fn rerun_if_changed_dir(dir: &str) {
println!("cargo:rerun-if-changed={}/", dir);
glob::glob(&format!("./{}/**/*.h", dir))
.expect("Failed to glob for source files from build.rs")
.filter_map(|e| e.ok())
.for_each(|path| println!("cargo:rerun-if-changed={}", path.to_string_lossy()));
}
fn generate_bindings_kernel_headers() -> Result<()> {
let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
let types = ["pt_regs", "s32", "bpf_.*"];
let vars = ["BPF_.*"];
let xdp_types = [
"xdp_md",
"ethhdr",
"iphdr",
"ipv6hdr",
"tcphdr",
"udphdr",
"xdp_action",
"__sk_.*",
"sk_.*",
"inet_sock",
"sockaddr",
"sockaddr_in",
"in_addr",
"tcp.*_sock",
"udp.*_sock",
"btf_ptr",
"linux_binprm",
"^sock_type$", "^sock_flags$", ];
let xdp_vars = ["ETH_.*", "IPPROTO_.*", "SOCK_.*", "SK_FL_.*", "AF_.*"];
let mut builder = bpf_bindgen::get_builder_kernel_headers()
.or_else(|e| bail!("error on Builder::get_builder_kernel_headers: {}", e))?
.header("include/redbpf_helpers.h")
.header("include/bpf_helpers.h");
for ty in types.iter().chain(xdp_types.iter()) {
builder = builder.allowlist_type(ty);
}
for var in vars.iter().chain(xdp_vars.iter()) {
builder = builder.allowlist_var(var);
}
builder = builder.opaque_type("xregs_state");
let mut bindings = builder
.generate()
.or_else(|e| bail!("error on Builder::generate: {:?}", e))?
.to_string();
let accessors = bpf_bindgen::generate_read_accessors(
&bindings,
&[
"sock",
"sockaddr",
"sockaddr_in",
"in_addr",
"file",
"inode",
"path",
"dentry",
"qstr",
],
);
bindings.push_str("use crate::helpers::bpf_probe_read;");
bindings.push_str(&accessors);
create_module(out_dir.join("gen_bindings.rs"), "gen_bindings", &bindings)?;
let bindings = bpf_bindgen::get_builder_kernel_headers()
.or_else(|e| bail!("error on Builder::get_builder: {}", e))?
.header("include/redbpf_helpers.h")
.header("include/bpf_helpers.h")
.allowlist_var("bpf_.*")
.generate()
.or_else(|e| {
bail!(
"error on Builder::generate while generating helpers: {:?}",
e
)
})?;
let helpers = gen_helpers(&bindings.to_string());
create_module(out_dir.join("gen_helpers.rs"), "gen_helpers", &helpers)?;
Ok(())
}
#[derive(Debug)]
struct HideEnum;
impl ParseCallbacks for HideEnum {
fn enum_variant_behavior(
&self,
_enum_name: Option<&str>,
_original_variant_name: &str,
_variant_value: EnumVariantValue,
) -> Option<EnumVariantCustomBehavior> {
Some(EnumVariantCustomBehavior::Hide)
}
}
#[derive(Debug)]
struct Callbacks;
impl ParseCallbacks for Callbacks {
fn item_name(&self, name: &str) -> Option<String> {
match name {
"u8" | "u16" | "u32" | "u64" => Some(format!("_cargo_bpf_{}", name)),
_ => None,
}
}
}
fn get_builder_vmlinux(c_dump_file: impl AsRef<Path>) -> Result<Builder, String> {
if let Ok(_) = env::var("DOCS_RS") {
Ok(bindgen::builder()
.use_core()
.ctypes_prefix("::cty")
.parse_callbacks(Box::new(Callbacks))
.header("include/docsrs/vmlinux.h"))
} else {
bpf_bindgen::get_builder_vmlinux(c_dump_file)
}
}
fn generate_bindings_vmlinux() -> Result<()> {
let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
let types = ["^pt_regs$", "^s32$", "^bpf_.*"];
let vars = ["^BPF_.*"];
let xdp_types = [
"^xdp_md$",
"^ethhdr$",
"^iphdr$",
"^ipv6hdr$",
"^tcphdr$",
"^udphdr$",
"^xdp_action$",
"^__sk_.*",
"^sk_.*",
"^inet_sock$",
"^sockaddr$",
"^sockaddr_in$",
"^in_addr$",
"^tcp.*_sock$",
"^udp.*_sock$",
"^btf_ptr$",
"^sock_type$", "^sock_flags$", ];
let xdp_vars = ["^IPPROTO_.*"];
let mut builder = get_builder_vmlinux(out_dir.join("vmlinux.h"))
.or_else(|e| bail!("error on bpf_bindgen::get_builder_vmlinux: {}", e))?
.no_debug("ec_response_motion_sense_fifo_info")
.no_debug("tpm2_.*")
.no_debug("pubkey_hdr")
.no_debug("__pldm.*")
.no_debug("signature_.*hdr");
for ty in types.iter().chain(xdp_types.iter()) {
builder = builder.allowlist_type(ty);
}
for var in vars.iter().chain(xdp_vars.iter()) {
builder = builder.allowlist_var(var);
}
builder = builder.opaque_type("xregs_state");
let mut bindings = builder
.generate()
.or_else(|e| bail!("error on Builder::generate: {:?}", e))?
.to_string();
let accessors = bpf_bindgen::generate_read_accessors(
&bindings,
&[
"sock",
"sockaddr",
"sockaddr_in",
"in_addr",
"file",
"inode",
"path",
"dentry",
"qstr",
],
);
bindings.push_str("use crate::helpers::bpf_probe_read;");
bindings.push_str(&accessors);
let supplement = bindgen::builder()
.use_core()
.ctypes_prefix("::cty")
.header("include/vmlinux_supplement.h")
.allowlist_var("^AF_.*")
.allowlist_var("^ETH_.*")
.allowlist_var("^BPF_.*")
.allowlist_var("^IPPROTO_.*") .allowlist_var("^SOCK_.*")
.allowlist_type("^bpf_map_def$")
.allowlist_type("^bpf_timer$")
.blocklist_type("_bindgen_ty.*") .parse_callbacks(Box::new(HideEnum)) .generate()
.or_else(|e| bail!("error on Builder::generate for supplement: {:?}", e))?
.to_string();
bindings.push_str(&supplement);
create_module(out_dir.join("gen_bindings.rs"), "gen_bindings", &bindings)?;
let bindings = get_builder_vmlinux(out_dir.join("vmlinux_helpers.h"))
.or_else(|e| bail!("error on bpf_bindgen::get_builder_vmlinux: {}", e))?
.header("include/bpf_helpers.h")
.allowlist_var("^bpf_.*")
.generate()
.or_else(|e| bail!("error on Builder::generate for helper: {:?}", e))?;
let helpers = gen_helpers(&bindings.to_string());
create_module(out_dir.join("gen_helpers.rs"), "gen_helpers", &helpers)?;
Ok(())
}
fn generate_bindings() {
if get_custom_vmlinux_path().is_some() {
debug!("Generating bindings with BTF of vmlinux");
generate_bindings_vmlinux().unwrap();
} else if get_custom_header_path().is_some() || get_custom_header_version().is_some() {
debug!("Generating bindings with pre-intalled kernel headers");
generate_bindings_kernel_headers().unwrap();
} else {
debug!("Try generating rust bindings with pre-installed kernel headers");
generate_bindings_kernel_headers()
.or_else(|e| {
warn!("error on generate_bindings_kernel_headers: {:?}", e);
debug!("Try generating rust bindings with vmlinux");
generate_bindings_vmlinux()
})
.unwrap()
}
}
fn main() {
let subscriber = FmtSubscriber::builder()
.with_max_level(Level::TRACE)
.finish();
tracing::subscriber::set_global_default(subscriber).unwrap();
rerun_if_changed_dir("include");
generate_bindings();
}
struct RewriteBpfHelpers {
helpers: Vec<String>,
}
impl Visit<'_> for RewriteBpfHelpers {
fn visit_foreign_item_static(&mut self, item: &ForeignItemStatic) {
if let Type::Path(path) = &*item.ty {
let ident = &item.ident;
let ident_str = ident.to_string();
let last = path.path.segments.last().unwrap();
let ty_ident = last.ident.to_string();
if ident_str.starts_with("bpf_") && ty_ident == "Option" {
let fn_ty = match &last.arguments {
AngleBracketed(AngleBracketedGenericArguments { args, .. }) => {
args.first().unwrap()
}
_ => panic!(),
};
let mut ty_s = quote! {
#[inline(always)]
pub #fn_ty
}
.to_string();
ty_s = ty_s.replace("fn (", &format!("fn {} (", ident_str));
let call_idx = self.helpers.len() + 1;
let args: Punctuated<Ident, Comma> = match fn_ty {
GenericArgument::Type(Type::BareFn(f)) => f
.inputs
.iter()
.map(|arg| arg.name.clone().unwrap().0)
.collect(),
_ => unreachable!(),
};
let body = quote! {
{
let f: #fn_ty = ::core::mem::transmute(#call_idx);
f(#args)
}
}
.to_string();
ty_s.push_str(&body);
let mut helper = ty_s;
if helper.contains("printk") {
helper = format!("/* {} */", helper);
}
self.helpers.push(helper);
}
}
}
}
fn gen_helpers(helpers: &str) -> String {
let tree: syn::File = parse_str(&helpers).unwrap();
let mut tx = RewriteBpfHelpers {
helpers: Vec::new(),
};
tx.visit_file(&tree);
let mut out = String::new();
out.push_str("use crate::bindings::*;\n");
for helper in &tx.helpers {
out.push_str(helper);
}
out
}