Add aya-gen

aya-gen can be used to generate bindings for kernel types, eg:

	aya-gen btf-types ethhdr --probe-read-getters

Will generate:

	// ...

	#[repr(C)]
	#[derive(Debug, Copy, Clone)]
	pub struct ethhdr {
	    pub h_dest: [::aya_bpf_cty::c_uchar; 6usize],
	    pub h_source: [::aya_bpf_cty::c_uchar; 6usize],
	    pub h_proto: __be16,
	}

	impl ethhdr {
	    pub fn h_dest(&self) -> Option<[::aya_bpf_cty::c_uchar; 6usize]> {
		unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.h_dest) }.ok()
	    }
	    pub fn h_source(&self) -> Option<[::aya_bpf_cty::c_uchar; 6usize]> {
		unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.h_source) }.ok()
	    }
	    pub fn h_proto(&self) -> Option<__be16> {
		unsafe { ::aya_bpf::helpers::bpf_probe_read(&self.h_proto) }.ok()
	    }
	}

Cargo.toml

@@ -1,2 +1,2 @@
 [workspace]
-members = ["aya", "xtask"]
+members = ["aya", "aya-gen", "xtask"]

aya-gen/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "aya-gen"
+version = "0.1.0"
+authors = ["Alessandro Decina <alessandro.d@gmail.com>"]
+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"

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<String>,
+    },
+}
+
+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(())
+}

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<T: AsRef<str>>(
+    btf_file: &Path,
+    types: &[T],
+    probe_read_getters: bool,
+) -> Result<String, Error> {
+    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::<syn::File>(&bindings).unwrap();
+    let bpf_probe_read = syn::parse_str::<syn::Path>("::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<String, Error> {
+    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())
+}

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<Ident, (&'a Item, &'a FieldsNamed)>,
+}
+
+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<Item = (&'a Item, Vec<Getter<'_>>)> {
+        self.item_fields
+            .values()
+            .map(move |(item, fields)| (*item, self.getters(&self.slf, fields)))
+    }
+
+    fn getters(&self, ident: &'a Ident, fields: &'a FieldsNamed) -> Vec<Getter<'a>> {
+        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,
+    }
+}

aya-gen/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod btf_types;
+pub mod getters;
+pub mod rustfmt;

aya-gen/src/rustfmt.rs

@@ -0,0 +1,25 @@
+use std::{
+    io::{self, Write},
+    process::{Command, Stdio},
+};
+
+pub fn format(code: &str) -> Result<String, io::Error> {
+    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())
+}