diff --git a/Cargo.toml b/Cargo.toml index 81a5d761..a345c22e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [workspace] members = [ - "aya", "aya-tool", "aya-log", "aya-log-common", "aya-log-parser", "test/integration-test", "test/integration-test-macros", "xtask", + "aya", "aya-common", "aya-tool", "aya-log", "aya-log-common", "aya-log-parser", "test/integration-test", "test/integration-test-macros", "xtask", # macros "aya-bpf-macros", "aya-log-ebpf-macros", # ebpf crates @@ -19,4 +19,4 @@ opt-level = 2 overflow-checks = false [profile.release.package.integration-ebpf] -debug = 2 \ No newline at end of file +debug = 2 diff --git a/aya-bpf-macros/Cargo.toml b/aya-bpf-macros/Cargo.toml index a601444d..4e9d396a 100644 --- a/aya-bpf-macros/Cargo.toml +++ b/aya-bpf-macros/Cargo.toml @@ -13,4 +13,4 @@ quote = "1.0" syn = {version = "1.0", features = ["full"]} [dev-dependencies] -aya-bpf = { path = "../bpf/aya-bpf" } +aya-bpf = { path = "../bpf/aya-bpf", features = ["usdt"] } diff --git a/aya-bpf-macros/src/expand.rs b/aya-bpf-macros/src/expand.rs index 42088969..fc9f1775 100644 --- a/aya-bpf-macros/src/expand.rs +++ b/aya-bpf-macros/src/expand.rs @@ -820,6 +820,34 @@ impl SkLookup { } } +pub struct Usdt { + item: ItemFn, + name: String, +} + +impl Usdt { + pub fn from_syn(mut args: Args, item: ItemFn) -> Result { + let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string()); + Ok(Usdt { item, name }) + } + + pub fn expand(&self) -> Result { + let section_name = format!("usdt/{}", self.name); + let fn_name = &self.item.sig.ident; + let item = &self.item; + Ok(quote! { + #[no_mangle] + #[link_section = #section_name] + fn #fn_name(ctx: *mut ::core::ffi::c_void) -> u32 { + let _ = #fn_name(::aya_bpf::programs::UsdtContext::new(ctx)); + return 0; + + #item + } + }) + } +} + #[cfg(test)] mod tests { use syn::parse_quote; diff --git a/aya-bpf-macros/src/lib.rs b/aya-bpf-macros/src/lib.rs index be255952..2f1a2f0d 100644 --- a/aya-bpf-macros/src/lib.rs +++ b/aya-bpf-macros/src/lib.rs @@ -3,7 +3,8 @@ mod expand; use expand::{ Args, BtfTracePoint, CgroupSkb, CgroupSock, CgroupSockAddr, CgroupSockopt, CgroupSysctl, FEntry, FExit, Lsm, Map, PerfEvent, Probe, ProbeKind, RawTracePoint, SchedClassifier, SkLookup, - SkMsg, SkSkb, SkSkbKind, SockAddrArgs, SockOps, SocketFilter, SockoptArgs, TracePoint, Xdp, + SkMsg, SkSkb, SkSkbKind, SockAddrArgs, SockOps, SocketFilter, SockoptArgs, TracePoint, Usdt, + Xdp, }; use proc_macro::TokenStream; use syn::{parse_macro_input, ItemFn, ItemStatic}; @@ -507,3 +508,30 @@ pub fn sk_lookup(attrs: TokenStream, item: TokenStream) -> TokenStream { .unwrap_or_else(|err| err.to_compile_error()) .into() } + +/// Marks a function as an User Statically-Defined Tracepoint program. +/// +/// # Minimum kernel version +/// +/// The minimum kernel version required to use this feature is 4.20 +/// +/// # Examples +/// +/// ```no_run +/// use aya_bpf::{macros::usdt, programs::UsdtContext}; +/// +/// #[usdt] +/// pub fn tick(_ctx: UsdtContext) -> u32 { +/// return 0 +/// } +/// ``` +#[proc_macro_attribute] +pub fn usdt(attrs: TokenStream, item: TokenStream) -> TokenStream { + let args = parse_macro_input!(attrs as Args); + let item = parse_macro_input!(item as ItemFn); + + Usdt::from_syn(args, item) + .and_then(|u| u.expand()) + .unwrap_or_else(|err| err.to_compile_error()) + .into() +} diff --git a/bpf/aya-bpf/Cargo.toml b/bpf/aya-bpf/Cargo.toml index a1da1e87..79c108c2 100644 --- a/bpf/aya-bpf/Cargo.toml +++ b/bpf/aya-bpf/Cargo.toml @@ -8,6 +8,11 @@ edition = "2021" aya-bpf-cty = { path = "../aya-bpf-cty" } aya-bpf-macros = { path = "../../aya-bpf-macros" } aya-bpf-bindings = { path = "../aya-bpf-bindings" } +aya-common = { path = "../../aya-common", optional = true } [build-dependencies] -rustversion = "1.0" \ No newline at end of file +rustversion = "1.0" + +[features] +usdt = ["aya-common"] +cookie = ["usdt"] diff --git a/bpf/aya-bpf/src/args.rs b/bpf/aya-bpf/src/args.rs index fea024eb..e613356f 100644 --- a/bpf/aya-bpf/src/args.rs +++ b/bpf/aya-bpf/src/args.rs @@ -76,6 +76,11 @@ impl PtRegs { T::from_retval(unsafe { &*self.regs }) } + /// Returns the value of the register used to pass the IP + pub fn ip(&self) -> Option { + T::from_ip(unsafe { &*self.regs }) + } + /// Returns a pointer to the wrapped value. pub fn as_ptr(&self) -> *mut pt_regs { self.regs @@ -95,6 +100,9 @@ pub trait FromPtRegs: Sized { /// Coerces a `T` from the return value of a pt_regs context. fn from_retval(ctx: &pt_regs) -> Option; + + /// Coerces a `T` from the ip value of a pt_regs context. + fn from_ip(ctx: &pt_regs) -> Option; } #[cfg(bpf_target_arch = "x86_64")] @@ -114,6 +122,10 @@ impl FromPtRegs for *const T { fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *const _).ok() } } + + fn from_ip(ctx: &pt_regs) -> Option { + unsafe { bpf_probe_read(&ctx.rip).map(|v| v as *const _).ok() } + } } #[cfg(bpf_target_arch = "arm")] @@ -129,6 +141,10 @@ impl FromPtRegs for *const T { fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *const _).ok() } } + + fn from_ip(ctx: &pt_regs) -> Option { + unsafe { bpf_probe_read(&ctx.uregs[12]).map(|v| v as *const _).ok() } + } } #[cfg(bpf_target_arch = "aarch64")] @@ -144,6 +160,10 @@ impl FromPtRegs for *const T { fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *const _).ok() } } + + fn from_ip(ctx: &pt_regs) -> Option { + unsafe { bpf_probe_read(&ctx.pc).map(|v| v as *const _).ok() } + } } #[cfg(bpf_target_arch = "x86_64")] @@ -163,6 +183,10 @@ impl FromPtRegs for *mut T { fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *mut _).ok() } } + + fn from_ip(ctx: &pt_regs) -> Option { + unsafe { bpf_probe_read(&ctx.rip).map(|v| v as *mut _).ok() } + } } #[cfg(bpf_target_arch = "arm")] @@ -178,6 +202,10 @@ impl FromPtRegs for *mut T { fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *mut _).ok() } } + + fn from_ip(ctx: &pt_regs) -> Option { + unsafe { bpf_probe_read(&ctx.uregs[12]).map(|v| v as *mut _).ok() } + } } #[cfg(bpf_target_arch = "aarch64")] @@ -193,6 +221,10 @@ impl FromPtRegs for *mut T { fn from_retval(ctx: &pt_regs) -> Option { unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *mut _).ok() } } + + fn from_ip(ctx: &pt_regs) -> Option { + unsafe { bpf_probe_read(&ctx.pc).map(|v| v as *mut _).ok() } + } } /// Helper macro to implement [`FromPtRegs`] for a primitive type. @@ -215,6 +247,10 @@ macro_rules! impl_from_pt_regs { fn from_retval(ctx: &pt_regs) -> Option { Some(ctx.rax as *const $type as _) } + + fn from_ip(ctx: &pt_regs) -> Option { + Some(ctx.rip as *const $type as _) + } } #[cfg(bpf_target_arch = "arm")] @@ -230,6 +266,10 @@ macro_rules! impl_from_pt_regs { fn from_retval(ctx: &pt_regs) -> Option { Some(ctx.uregs[0] as *const $type as _) } + + fn from_ip(ctx: &pt_regs) -> Option { + Some(ctx.uregs[12] as *const $type as _) + } } #[cfg(bpf_target_arch = "aarch64")] @@ -245,6 +285,10 @@ macro_rules! impl_from_pt_regs { fn from_retval(ctx: &pt_regs) -> Option { Some(ctx.regs[0] as *const $type as _) } + + fn from_ip(ctx: &pt_regs) -> Option { + Some(ctx.pc as *const $type as _) + } } }; } diff --git a/bpf/aya-bpf/src/programs/mod.rs b/bpf/aya-bpf/src/programs/mod.rs index edd69397..2b21413e 100644 --- a/bpf/aya-bpf/src/programs/mod.rs +++ b/bpf/aya-bpf/src/programs/mod.rs @@ -15,6 +15,8 @@ pub mod sysctl; pub mod tc; pub mod tp_btf; pub mod tracepoint; +#[cfg(feature = "usdt")] +pub mod usdt; pub mod xdp; pub use fentry::FEntryContext; @@ -34,4 +36,6 @@ pub use sysctl::SysctlContext; pub use tc::TcContext; pub use tp_btf::BtfTracePointContext; pub use tracepoint::TracePointContext; +#[cfg(feature = "usdt")] +pub use usdt::UsdtContext; pub use xdp::XdpContext; diff --git a/bpf/aya-bpf/src/programs/usdt.rs b/bpf/aya-bpf/src/programs/usdt.rs new file mode 100644 index 00000000..8cf136b1 --- /dev/null +++ b/bpf/aya-bpf/src/programs/usdt.rs @@ -0,0 +1,104 @@ +use core::ffi::c_void; + +use aya_common::{ + UsdtArgType, UsdtSpec, USDT_MAX_ARG_COUNT, USDT_MAX_IP_COUNT, USDT_MAX_SPEC_COUNT, +}; + +use crate::{ + args::FromPtRegs, + helpers::{bpf_probe_read_kernel, bpf_probe_read_user}, + macros::map, + maps::{Array, HashMap}, + BpfContext, +}; + +// aarch64 uses user_pt_regs instead of pt_regs +#[cfg(not(bpf_target_arch = "aarch64"))] +use crate::bindings::pt_regs; +#[cfg(bpf_target_arch = "aarch64")] +use crate::bindings::user_pt_regs as pt_regs; + +#[map(name = "__bpf_usdt_specs")] +static USDT_SPECS: Array = Array::with_max_entries(USDT_MAX_SPEC_COUNT, 0); + +#[map(name = "__bpf_usdt_ip_to_spec_id")] +static USDT_IP_TO_SPEC_ID: HashMap = HashMap::with_max_entries(USDT_MAX_IP_COUNT, 0); + +pub struct UsdtContext { + pub regs: *mut pt_regs, +} + +impl UsdtContext { + pub fn new(ctx: *mut c_void) -> UsdtContext { + UsdtContext { + regs: ctx as *mut pt_regs, + } + } + + #[inline(always)] + fn ip(&self) -> Option { + T::from_ip(unsafe { &*self.regs }) + } + + #[cfg(feature = "cookie")] + #[inline(always)] + fn spec_id(&self) -> Result { + unsafe { Ok(aya_bpf_bindings::helpers::bpf_get_attach_cookie(self.as_ptr()) as u32) } + } + + #[cfg(not(feature = "cookie"))] + #[inline(always)] + fn spec_id(&self) -> Result { + let ip: i64 = self.ip().ok_or(())?; + let spec = unsafe { USDT_IP_TO_SPEC_ID.get(&ip).ok_or(())? }; + Ok(*spec) + } + + #[inline(always)] + pub fn arg(&self, n: usize) -> Result { + if n > USDT_MAX_ARG_COUNT { + return Err(()); + } + let spec_id = self.spec_id()?; + let spec = USDT_SPECS.get(spec_id).ok_or(())?; + + if n > (spec.arg_count as usize) { + return Err(()); + } + + let arg_spec = &spec.args[n]; + let mut val = match arg_spec.arg_type { + UsdtArgType::Const => arg_spec.val_off, + UsdtArgType::Reg => unsafe { + bpf_probe_read_kernel(self.as_ptr().offset(arg_spec.reg_off as isize) as *const _) + .map_err(|_| ())? + }, + UsdtArgType::RegDeref => unsafe { + let ptr: u64 = bpf_probe_read_kernel( + self.as_ptr().offset(arg_spec.reg_off as isize) as *const _, + ) + .map_err(|_| ())?; + let ptr = ptr as *const u64; + bpf_probe_read_user::(ptr.offset(arg_spec.val_off as isize)).map_err(|_| ())? + // TODO: libbpf applies a bitshift here if the arch is big endian + }, + }; + + // cast arg from 1, 2, or 4 bytes to final 8 byte size clearing + // necessary upper arg_bitshift bits, with sign extension if argument + // is signed + val <<= arg_spec.arg_bitshift; + if arg_spec.arg_signed { + val = ((val as i64) >> arg_spec.arg_bitshift) as u64 + } else { + val >>= arg_spec.arg_bitshift; + } + Ok(val) + } +} + +impl BpfContext for UsdtContext { + fn as_ptr(&self) -> *mut c_void { + self.regs as *mut c_void + } +} diff --git a/xtask/src/codegen/aya_common.rs b/xtask/src/codegen/aya_common.rs index 0676944d..db4b6c55 100644 --- a/xtask/src/codegen/aya_common.rs +++ b/xtask/src/codegen/aya_common.rs @@ -1,7 +1,7 @@ use anyhow::anyhow; use std::path::PathBuf; -use aya_gen::{bindgen, write_to_file}; +use aya_tool::{bindgen, write_to_file}; use crate::codegen::{Architecture, Options};