extern crate proc_macro;
extern crate proc_macro2;
use proc_macro::TokenStream;
use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
use quote::quote;
use std::str;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::token::Comma;
use syn::{
parse_macro_input, parse_quote, parse_str, Expr, ExprLit, File, FnArg, ItemFn, Lit, Pat,
PatIdent, PatType, Result, Stmt,
};
fn inline_string_literal(e: &Expr) -> (TokenStream2, TokenStream2) {
let mut bytes = match e {
Expr::Lit(ExprLit {
lit: Lit::Str(s), ..
}) => s.value().clone().into_bytes(),
_ => panic!("expected string literal"),
};
bytes.push(0u8);
let len = bytes.len();
let bytes = bytes;
let ty = quote!([u8; #len]);
let array_lit = quote!([#(#bytes),*]);
(ty, array_lit)
}
struct Args(Punctuated<Expr, Comma>);
impl Parse for Args {
fn parse(input: ParseStream) -> Result<Args> {
Ok(Args(Punctuated::parse_terminated(input)?))
}
}
#[proc_macro]
pub fn program(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as Args);
let mut args = input.0.iter();
let version = args.next().expect("no version");
let license = args.next().expect("no license");
let (license_ty, license) = inline_string_literal(&license);
let mut tokens = quote! {
#[no_mangle]
#[link_section = "license"]
pub static _license: #license_ty = #license;
#[no_mangle]
#[link_section = "version"]
pub static _version: u32 = #version;
#[panic_handler]
#[no_mangle]
pub extern "C" fn rust_begin_panic(info: &::core::panic::PanicInfo) -> ! {
use ::redbpf_probes::helpers::{bpf_trace_printk, TraceMessage, ufmt};
let mut msg = TraceMessage::new();
ufmt::uwrite!(&mut msg, "panic in {}\n\0", file!());
msg.printk();
unsafe { core::hint::unreachable_unchecked() }
}
};
let mem = str::from_utf8(include_bytes!("mem.rs")).unwrap();
let mem: File = parse_str(&mem).unwrap();
tokens.extend(quote! {
#mem
});
tokens.into()
}
#[proc_macro]
pub fn impl_network_buffer_array(_: TokenStream) -> TokenStream {
let mut tokens = TokenStream2::new();
for i in 1..=512usize {
tokens.extend(quote! {
impl NetworkBufferArray for [u8; #i] {}
});
}
tokens.into()
}
#[proc_macro_attribute]
pub fn map(attrs: TokenStream, item: TokenStream) -> TokenStream {
let attrs = parse_macro_input!(attrs as Expr);
let name = match attrs {
Expr::Lit(ExprLit {
lit: Lit::Str(s), ..
}) => s.value().clone(),
_ => panic!("expected string literal"),
};
let section_name = format!("maps/{}", name);
let item = TokenStream2::from(item);
let tokens = quote! {
#[no_mangle]
#[link_section = #section_name]
#item
};
tokens.into()
}
fn probe_impl(ty: &str, attrs: TokenStream, item: ItemFn) -> TokenStream {
let name = if attrs.is_empty() {
item.sig.ident.to_string()
} else {
match parse_macro_input!(attrs as Expr) {
Expr::Lit(ExprLit {
lit: Lit::Str(s), ..
}) => s.value().clone(),
_ => panic!("expected string literal"),
}
};
let section_name = format!("{}/{}", ty, name);
let tokens = quote! {
#[no_mangle]
#[link_section = #section_name]
#item
};
tokens.into()
}
#[proc_macro_attribute]
pub fn kprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
let item = parse_macro_input!(item as ItemFn);
probe_impl("kprobe", attrs, item).into()
}
#[proc_macro_attribute]
pub fn kretprobe(attrs: TokenStream, item: TokenStream) -> TokenStream {
let item = parse_macro_input!(item as ItemFn);
probe_impl("kretprobe", attrs, item).into()
}
#[proc_macro_attribute]
pub fn xdp(attrs: TokenStream, item: TokenStream) -> TokenStream {
let mut item = parse_macro_input!(item as ItemFn);
let arg = item.sig.inputs.pop().unwrap();
let pat = match arg.value() {
FnArg::Typed(PatType { pat, .. }) => pat,
_ => panic!("unexpected xdp probe signature"),
};
let ident = if let Pat::Ident(PatIdent { ident, .. }) = &**pat {
ident
} else {
panic!("unexpected xdp probe signature")
};
let raw_ctx = Ident::new(&format!("_raw_{}", ident), Span::call_site());
let arg: FnArg = parse_quote! { #raw_ctx: *mut xdp_md };
item.sig.inputs.push(arg);
let ctx: Stmt = parse_quote! { let #ident = XdpContext { ctx: #raw_ctx }; };
item.block.stmts.insert(0, ctx);
probe_impl("xdp", attrs, item).into()
}
#[proc_macro_attribute]
pub fn socket_filter(attrs: TokenStream, item: TokenStream) -> TokenStream {
let item = parse_macro_input!(item as ItemFn);
let ident = item.sig.ident.clone();
let outer_ident = Ident::new(&format!("outer_{}", ident), Span::call_site());
let wrapper = parse_quote! {
fn #outer_ident(skb: *const ::redbpf_probes::bindings::__sk_buff) -> i32 {
let skb = ::redbpf_probes::socket_filter::SkBuff { skb };
return match #ident(skb) {
Ok(::redbpf_probes::socket_filter::SkBuffAction::SendToUserspace) => -1,
_ => 0
};
#item
}
};
probe_impl("socketfilter", attrs, wrapper).into()
}