aya-bpf: Implement USDT programs

Since these programs require that a few maps are created, we hide this
behind a feature since otherwise the maps will get created regardless of
the program type.

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>

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

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"] }

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<Usdt> {
+        let name = name_arg(&mut args)?.unwrap_or_else(|| item.sig.ident.to_string());
+        Ok(Usdt { item, name })
+    }
+
+    pub fn expand(&self) -> Result<TokenStream> {
+        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;

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

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"
+
+[features]
+usdt = ["aya-common"]
+cookie = ["usdt"]

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<T: FromPtRegs>(&self) -> Option<T> {
+        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<Self>;
+
+    /// Coerces a `T` from the ip value of a pt_regs context.
+    fn from_ip(ctx: &pt_regs) -> Option<Self>;
 }
 
 #[cfg(bpf_target_arch = "x86_64")]
@@ -114,6 +122,10 @@ impl<T> FromPtRegs for *const T {
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *const _).ok() }
     }
+
+    fn from_ip(ctx: &pt_regs) -> Option<Self> {
+        unsafe { bpf_probe_read(&ctx.rip).map(|v| v as *const _).ok() }
+    }
 }
 
 #[cfg(bpf_target_arch = "arm")]
@@ -129,6 +141,10 @@ impl<T> FromPtRegs for *const T {
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *const _).ok() }
     }
+
+    fn from_ip(ctx: &pt_regs) -> Option<Self> {
+        unsafe { bpf_probe_read(&ctx.uregs[12]).map(|v| v as *const _).ok() }
+    }
 }
 
 #[cfg(bpf_target_arch = "aarch64")]
@@ -144,6 +160,10 @@ impl<T> FromPtRegs for *const T {
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *const _).ok() }
     }
+
+    fn from_ip(ctx: &pt_regs) -> Option<Self> {
+        unsafe { bpf_probe_read(&ctx.pc).map(|v| v as *const _).ok() }
+    }
 }
 
 #[cfg(bpf_target_arch = "x86_64")]
@@ -163,6 +183,10 @@ impl<T> FromPtRegs for *mut T {
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.rax).map(|v| v as *mut _).ok() }
     }
+
+    fn from_ip(ctx: &pt_regs) -> Option<Self> {
+        unsafe { bpf_probe_read(&ctx.rip).map(|v| v as *mut _).ok() }
+    }
 }
 
 #[cfg(bpf_target_arch = "arm")]
@@ -178,6 +202,10 @@ impl<T> FromPtRegs for *mut T {
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.uregs[0]).map(|v| v as *mut _).ok() }
     }
+
+    fn from_ip(ctx: &pt_regs) -> Option<Self> {
+        unsafe { bpf_probe_read(&ctx.uregs[12]).map(|v| v as *mut _).ok() }
+    }
 }
 
 #[cfg(bpf_target_arch = "aarch64")]
@@ -193,6 +221,10 @@ impl<T> FromPtRegs for *mut T {
     fn from_retval(ctx: &pt_regs) -> Option<Self> {
         unsafe { bpf_probe_read(&ctx.regs[0]).map(|v| v as *mut _).ok() }
     }
+
+    fn from_ip(ctx: &pt_regs) -> Option<Self> {
+        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<Self> {
                 Some(ctx.rax as *const $type as _)
             }
+
+            fn from_ip(ctx: &pt_regs) -> Option<Self> {
+                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<Self> {
                 Some(ctx.uregs[0] as *const $type as _)
             }
+
+            fn from_ip(ctx: &pt_regs) -> Option<Self> {
+                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<Self> {
                 Some(ctx.regs[0] as *const $type as _)
             }
+
+            fn from_ip(ctx: &pt_regs) -> Option<Self> {
+                Some(ctx.pc as *const $type as _)
+            }
         }
     };
 }

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;

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<UsdtSpec> = Array::with_max_entries(USDT_MAX_SPEC_COUNT, 0);
+
+#[map(name = "__bpf_usdt_ip_to_spec_id")]
+static USDT_IP_TO_SPEC_ID: HashMap<i64, u32> = 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<T: FromPtRegs>(&self) -> Option<T> {
+        T::from_ip(unsafe { &*self.regs })
+    }
+
+    #[cfg(feature = "cookie")]
+    #[inline(always)]
+    fn spec_id(&self) -> Result<u32, ()> {
+        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<u32, ()> {
+        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<u64, ()> {
+        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::<u64>(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
+    }
+}

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};