diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index 32508dc0..9945567a 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -8,7 +8,7 @@ edition = "2018" aya-gen = { path = "../aya-gen" } structopt = {version = "0.3", default-features = false } anyhow = "1" -syn = {version = "1", features = ["visit-mut", "extra-traits"] } +syn = "1" quote = "1" proc-macro2 = "1" indexmap = "1.6" \ No newline at end of file diff --git a/xtask/src/codegen/aya_bpf_bindings.rs b/xtask/src/codegen/aya_bpf_bindings.rs index 317f5455..9783faa4 100644 --- a/xtask/src/codegen/aya_bpf_bindings.rs +++ b/xtask/src/codegen/aya_bpf_bindings.rs @@ -1,23 +1,15 @@ -use std::path::PathBuf; - use anyhow::anyhow; use proc_macro2::TokenStream; -use quote::{quote, ToTokens}; +use quote::ToTokens; +use std::path::PathBuf; use structopt::StructOpt; -use syn::{ - self, parse_str, - punctuated::Punctuated, - token::Comma, - visit_mut::{self, VisitMut}, - AngleBracketedGenericArguments, ForeignItemStatic, GenericArgument, Ident, Item, Path, - PathArguments::AngleBracketed, - Type, -}; use aya_gen::getters::{generate_getters_for_items, probe_read_getter}; +use syn::{parse_str, Item}; use crate::codegen::{ bindings::{self, bindgen}, + helpers::{expand_helpers, extract_helpers}, Architecture, }; @@ -49,13 +41,12 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> { return Err(anyhow!("bindgen failed: {}", output.status)); } - // delete the helpers, then rewrite them in helpers.rs let mut tree = parse_str::(bindings).unwrap(); - - let mut tx = RewriteBpfHelpers { - helpers: Vec::new(), - }; - tx.visit_file_mut(&mut tree); + let (indexes, helpers) = extract_helpers(&tree.items); + let helpers = expand_helpers(&helpers); + for index in indexes { + tree.items[index] = Item::Verbatim(TokenStream::new()) + } bindings::write( &tree.to_token_stream().to_string(), @@ -64,12 +55,12 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> { )?; bindings::write( - &tx.helpers.join(""), + &helpers.to_string(), "use super::bindings::*;", &generated.join("helpers.rs"), )?; - let bpf_probe_read = syn::parse_str::("crate::bpf_probe_read").unwrap(); + let bpf_probe_read = syn::parse_str("crate::bpf_probe_read").unwrap(); bindings::write( &generate_getters_for_items(&tree.items, |getter| { probe_read_getter(getter, &bpf_probe_read) @@ -81,60 +72,3 @@ pub fn codegen(opts: CodegenOptions) -> Result<(), anyhow::Error> { Ok(()) } - -struct RewriteBpfHelpers { - helpers: Vec, -} - -impl VisitMut for RewriteBpfHelpers { - fn visit_item_mut(&mut self, item: &mut Item) { - visit_mut::visit_item_mut(self, item); - if let Item::ForeignMod(_) = item { - *item = Item::Verbatim(TokenStream::new()) - } - } - fn visit_foreign_item_static_mut(&mut self, static_item: &mut ForeignItemStatic) { - if let Type::Path(path) = &*static_item.ty { - let ident = &static_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 = 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); - } - } - } -} diff --git a/xtask/src/codegen/helpers.rs b/xtask/src/codegen/helpers.rs new file mode 100644 index 00000000..efc965ca --- /dev/null +++ b/xtask/src/codegen/helpers.rs @@ -0,0 +1,100 @@ +use proc_macro2::TokenStream; +use quote::{quote, TokenStreamExt}; +use syn::{ + punctuated::Punctuated, AngleBracketedGenericArguments, BareFnArg, ForeignItem, + ForeignItemStatic, GenericArgument, Ident, Item, Path, PathArguments, ReturnType, Token, Type, + TypeBareFn, TypePath, +}; + +pub fn extract_helpers(items: &[Item]) -> (Vec, Vec>) { + let mut helpers = Vec::new(); + let mut indexes = Vec::new(); + for (item_index, item) in items.iter().enumerate() { + if let Item::ForeignMod(module) = item { + for i in &module.items { + if let ForeignItem::Static(s_item) = i { + let ident_s = s_item.ident.to_string(); + if ident_s.starts_with("bpf_") { + helpers.push( + helper_from_item(s_item, helpers.len() + 1) + .expect("unexpected bindgen helper signature"), + ); + indexes.push(item_index); + } + } + } + } + } + + (indexes, helpers) +} + +pub fn helper_from_item(item: &ForeignItemStatic, call_index: usize) -> Option> { + if let Type::Path(TypePath { + path: Path { segments, .. }, + .. + }) = &*item.ty + { + let generics = &segments.last().unwrap().arguments; + if let PathArguments::AngleBracketed(AngleBracketedGenericArguments { args, .. }) = generics + { + if let Some(GenericArgument::Type(ty)) = args.first() { + if let Type::BareFn(TypeBareFn { inputs, output, .. }) = ty { + return Some(Helper { + ident: &item.ident, + ty, + inputs, + output, + call_index, + }); + } + } + }; + } + + None +} + +pub fn expand_helpers(helpers: &[Helper<'_>]) -> TokenStream { + let mut tokens = TokenStream::new(); + tokens.append_all( + helpers + .iter() + .filter(|h| h.ident.to_string() != "bpf_trace_printk") + .map(expand_helper), + ); + + return tokens; +} + +pub fn expand_helper(helper: &Helper<'_>) -> TokenStream { + let Helper { + ident, + ty, + inputs, + output, + call_index, + } = helper; + + let args = inputs + .iter() + .map(|arg| &arg.name.as_ref().unwrap().0) + .collect::>(); + + let helper = quote! { + pub unsafe fn #ident(#inputs) #output { + let fun: #ty = ::core::mem::transmute(#call_index); + fun(#(#args),*) + } + }; + + helper +} + +pub struct Helper<'a> { + ident: &'a Ident, + ty: &'a Type, + inputs: &'a Punctuated, + output: &'a ReturnType, + call_index: usize, +} diff --git a/xtask/src/codegen/mod.rs b/xtask/src/codegen/mod.rs index b929fc16..2e5db0f2 100644 --- a/xtask/src/codegen/mod.rs +++ b/xtask/src/codegen/mod.rs @@ -1,5 +1,6 @@ mod aya_bpf_bindings; mod bindings; +mod helpers; use structopt::StructOpt;