diff --git a/Cargo.toml b/Cargo.toml index 88752091..0eed0f81 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,2 +1,2 @@ [workspace] -members = ["aya", "xtask"] +members = ["aya", "aya-gen", "xtask"] diff --git a/aya-gen/Cargo.toml b/aya-gen/Cargo.toml new file mode 100644 index 00000000..ed1a3814 --- /dev/null +++ b/aya-gen/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "aya-gen" +version = "0.1.0" +authors = ["Alessandro Decina "] +edition = "2018" + +[dependencies] +bindgen = "0.57" +structopt = {version = "0.3", default-features = false } +anyhow = "1" +thiserror = "1" +syn = "1" +quote = "1" +proc-macro2 = "1" +indexmap = "1.6" \ No newline at end of file diff --git a/aya-gen/src/bin/main.rs b/aya-gen/src/bin/main.rs new file mode 100644 index 00000000..4bba535e --- /dev/null +++ b/aya-gen/src/bin/main.rs @@ -0,0 +1,45 @@ +use aya_gen::btf_types; + +use std::{path::PathBuf, process::exit}; + +use structopt::StructOpt; +#[derive(StructOpt)] +pub struct Options { + #[structopt(subcommand)] + command: Command, +} + +#[derive(StructOpt)] +enum Command { + #[structopt(name = "btf-types")] + BtfTypes { + #[structopt(long, default_value = "/sys/kernel/btf/vmlinux")] + btf: PathBuf, + #[structopt(long)] + probe_read_getters: bool, + names: Vec, + }, +} + +fn main() { + if let Err(e) = try_main() { + eprintln!("{:#}", e); + exit(1); + } +} + +fn try_main() -> Result<(), anyhow::Error> { + let opts = Options::from_args(); + match opts.command { + Command::BtfTypes { + btf, + probe_read_getters, + names, + } => { + let bindings = btf_types::generate(&btf, &names, probe_read_getters)?; + println!("{}", bindings); + } + }; + + Ok(()) +} diff --git a/aya-gen/src/btf_types.rs b/aya-gen/src/btf_types.rs new file mode 100644 index 00000000..1f65d581 --- /dev/null +++ b/aya-gen/src/btf_types.rs @@ -0,0 +1,78 @@ +use std::{io, path::Path, process::Command, str::from_utf8}; + +use bindgen::builder; +use thiserror::Error; + +use crate::{ + getters::{generate_getters_for_items, probe_read_getter}, + rustfmt, +}; + +#[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>( + btf_file: &Path, + types: &[T], + probe_read_getters: bool, +) -> Result { + let mut bindgen = builder() + .use_core() + .ctypes_prefix("::aya_bpf_cty") + .layout_tests(false) + .clang_arg("-Wno-unknown-attributes"); + + let c_header = c_header_from_btf(btf_file)?; + bindgen = bindgen.header_contents("kernel_types.h", &c_header); + + for ty in types { + bindgen = bindgen.whitelist_type(ty); + } + + let bindings = bindgen.generate().or(Err(Error::Bindgen))?.to_string(); + if !probe_read_getters { + return Ok(bindings); + } + + let tree = syn::parse_str::(&bindings).unwrap(); + let bpf_probe_read = syn::parse_str::("::aya_bpf::helpers::bpf_probe_read").unwrap(); + let getters = generate_getters_for_items(&tree.items, |getter| { + probe_read_getter(getter, &bpf_probe_read) + }); + let getters = + rustfmt::format(&getters.to_string()).map_err(|io_error| Error::Rustfmt(io_error))?; + + let bindings = format!("{}\n{}", bindings, getters); + + Ok(bindings) +} + +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(|io_error| Error::BpfTool(io_error))?; + + 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()) +} diff --git a/aya-gen/src/getters.rs b/aya-gen/src/getters.rs new file mode 100644 index 00000000..bf855531 --- /dev/null +++ b/aya-gen/src/getters.rs @@ -0,0 +1,151 @@ +use indexmap::IndexMap; +use proc_macro2::{Span, TokenStream}; +use quote::{quote, TokenStreamExt}; +use syn::{ + self, Fields, FieldsNamed, Generics, Ident, Item, ItemStruct, ItemUnion, Path, Type, TypePath, + Visibility, +}; + +pub struct GetterList<'a> { + slf: Ident, + item_fields: IndexMap, +} + +impl<'a> GetterList<'a> { + pub fn new(items: &'a [Item]) -> GetterList<'a> { + let item_fields = items + .iter() + .filter_map(|item| { + unpack_item(item).map(|(ident, _generics, fields)| (ident.clone(), (item, fields))) + }) + .collect(); + GetterList { + slf: Ident::new("self", Span::call_site()), + item_fields, + } + } + + pub fn iter(&self) -> impl Iterator>)> { + self.item_fields + .values() + .map(move |(item, fields)| (*item, self.getters(&self.slf, fields))) + } + + fn getters(&self, ident: &'a Ident, fields: &'a FieldsNamed) -> Vec> { + let mut getters = Vec::new(); + for field in &fields.named { + if let Visibility::Inherited = field.vis { + continue; + } + + let field_ident = field.ident.as_ref().unwrap(); + let field_s = field_ident.to_string(); + + // FIXME: bindgen generates fields named `_bitfield_N` for bitfields. If a type T has + // two or more unions with bitfields, the getters for the bitfields - generated in impl + // T - will clash. To avoid that we skip getters for bitfields altogether for now. + // See sk_reuseport_md for an example where the clash happens. + if field_s.starts_with("_bitfield") { + continue; + } + + if field_s.starts_with("__bindgen_anon") { + let field_ty_ident = match &field.ty { + Type::Path(TypePath { + path: Path { segments, .. }, + .. + }) => &segments.first().unwrap().ident, + _ => panic!(), + }; + let sub_fields = self + .item_fields + .get(field_ty_ident) + .expect(&field_ty_ident.to_string()) + .1; + getters.extend(self.getters(field_ident, sub_fields).drain(..).map( + |mut getter| { + getter.prefix.insert(0, ident); + getter + }, + )); + } else { + getters.push(Getter { + ident: field_ident, + prefix: vec![ident], + ty: &field.ty, + }); + } + } + + getters + } +} + +pub fn generate_getters_for_items( + items: &[Item], + gen_getter: impl Fn(&Getter<'_>) -> TokenStream, +) -> TokenStream { + let mut tokens = TokenStream::new(); + tokens.append_all(GetterList::new(&items).iter().map(|(item, getters)| { + let getters = getters.iter().map(&gen_getter); + let (ident, generics, _) = unpack_item(item).unwrap(); + quote! { + impl#generics #ident#generics { + #(#getters)* + } + } + })); + + tokens +} + +pub fn probe_read_getter(getter: &Getter<'_>, bpf_probe_read: &Path) -> TokenStream { + let ident = getter.ident; + let ty = getter.ty; + let prefix = &getter.prefix; + match ty { + Type::Ptr(_) => { + quote! { + pub fn #ident(&self) -> Option<#ty> { + let v = unsafe { #bpf_probe_read(&#(#prefix).*.#ident) }.ok()?; + if v.is_null() { + None + } else { + Some(v) + } + } + } + } + _ => { + quote! { + pub fn #ident(&self) -> Option<#ty> { + unsafe { #bpf_probe_read(&#(#prefix).*.#ident) }.ok() + } + } + } + } +} + +pub struct Getter<'a> { + pub ident: &'a Ident, + pub prefix: Vec<&'a Ident>, + pub ty: &'a Type, +} + +fn unpack_item(item: &Item) -> Option<(&Ident, &Generics, &FieldsNamed)> { + match item { + Item::Struct(ItemStruct { + ident, + generics, + fields: Fields::Named(fields), + .. + }) + | Item::Union(ItemUnion { + ident, + generics, + fields, + .. + }) => Some((ident, generics, fields)), + _ => None, + } +} diff --git a/aya-gen/src/lib.rs b/aya-gen/src/lib.rs new file mode 100644 index 00000000..b8b944a2 --- /dev/null +++ b/aya-gen/src/lib.rs @@ -0,0 +1,3 @@ +pub mod btf_types; +pub mod getters; +pub mod rustfmt; diff --git a/aya-gen/src/rustfmt.rs b/aya-gen/src/rustfmt.rs new file mode 100644 index 00000000..cc7ac292 --- /dev/null +++ b/aya-gen/src/rustfmt.rs @@ -0,0 +1,25 @@ +use std::{ + io::{self, Write}, + process::{Command, Stdio}, +}; + +pub fn format(code: &str) -> Result { + let mut child = Command::new("rustfmt") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .spawn()?; + let stdin = child.stdin.as_mut().unwrap(); + stdin.write_all(code.as_bytes())?; + + let output = child.wait_with_output()?; + if !output.status.success() { + return Err(io::Error::new( + io::ErrorKind::Other, + format!( + "rustfmt failed with exit code: {}", + output.status.code().unwrap() + ), + )); + } + Ok(String::from_utf8(output.stdout).unwrap()) +}