Merge dfd26f9071 into 49404367d8

.github/workflows/build-aya.yml

@@ -42,7 +42,7 @@ jobs:
         env:
           RUST_BACKTRACE: full
         run: |
-          cross test --verbose --target ${{matrix.arch}}
+          cross test --verbose --target ${{matrix.arch}} --all-features
 
   test:
     runs-on: ubuntu-20.04

.vim/coc-settings.json

@@ -1,4 +1,3 @@
 {
-  "rust-analyzer.checkOnSave.allTargets": false,
   "rust-analyzer.checkOnSave.command": "clippy"
 }

.vscode/settings.json

@@ -1,4 +1,3 @@
 {
-  "rust-analyzer.checkOnSave.allTargets": false,
   "rust-analyzer.checkOnSave.command": "clippy"
 }

Cargo.toml

@@ -1,12 +1,12 @@
 [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
     "bpf/aya-bpf", "bpf/aya-bpf-bindings", "bpf/aya-log-ebpf", "test/integration-ebpf"
 ]
-default-members = ["aya", "aya-tool", "aya-log", "aya-bpf-macros", "aya-log-ebpf-macros"]
+default-members = ["aya", "aya-common", "aya-tool", "aya-log", "aya-bpf-macros", "aya-log-ebpf-macros"]
 
 [profile.dev]
 panic = "abort"
@@ -19,4 +19,4 @@ opt-level = 2
 overflow-checks = false
 
 [profile.release.package.integration-ebpf]
-debug = 2
+debug = 2

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,32 @@ 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 {
+///     let arg = ctx.arg(0);
+///     // Use aya-log to print the value to userspace
+///     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()
+}

aya-common/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "aya-common"
+version = "0.1.0"
+description = "Code shared between aya and aya-bpf crates."
+keywords = ["ebpf", "bpf", "linux", "kernel"]
+license = "MIT OR Apache-2.0"
+authors = ["The Aya Contributors"]
+repository = "https://github.com/aya-rs/aya"
+readme = "README.md"
+documentation = "https://docs.rs/aya"
+edition = "2021"
+
+[dependencies]
+regex = {version = "1", optional = true}
+thiserror = {version = "1", optional = true}
+memoffset = {version = "0.6", optional = true}
+lazy_static = {version = "1", optional = true}
+
+[features]
+default = []
+user = [ "regex", "thiserror", "memoffset", "lazy_static" ]
+
+[lib]
+path = "src/lib.rs"

aya-common/include/linux_wrapper.h

@@ -0,0 +1 @@
+#include <linux/ptrace.h>

aya-common/src/generated/linux_bindings_aarch64.rs

@@ -0,0 +1,11 @@
+/* automatically generated by rust-bindgen 0.60.1 */
+
+pub type __u64 = ::std::os::raw::c_ulonglong;
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct user_pt_regs {
+    pub regs: [__u64; 31usize],
+    pub sp: __u64,
+    pub pc: __u64,
+    pub pstate: __u64,
+}

aya-common/src/generated/linux_bindings_armv7.rs

@@ -0,0 +1,7 @@
+/* automatically generated by rust-bindgen 0.60.1 */
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct pt_regs {
+    pub uregs: [::std::os::raw::c_long; 18usize],
+}

aya-common/src/generated/linux_bindings_riscv64.rs

@@ -0,0 +1 @@
+/* automatically generated by rust-bindgen 0.60.1 */

aya-common/src/generated/linux_bindings_x86_64.rs

@@ -0,0 +1,27 @@
+/* automatically generated by rust-bindgen 0.60.1 */
+
+#[repr(C)]
+#[derive(Debug, Copy, Clone)]
+pub struct pt_regs {
+    pub r15: ::std::os::raw::c_ulong,
+    pub r14: ::std::os::raw::c_ulong,
+    pub r13: ::std::os::raw::c_ulong,
+    pub r12: ::std::os::raw::c_ulong,
+    pub rbp: ::std::os::raw::c_ulong,
+    pub rbx: ::std::os::raw::c_ulong,
+    pub r11: ::std::os::raw::c_ulong,
+    pub r10: ::std::os::raw::c_ulong,
+    pub r9: ::std::os::raw::c_ulong,
+    pub r8: ::std::os::raw::c_ulong,
+    pub rax: ::std::os::raw::c_ulong,
+    pub rcx: ::std::os::raw::c_ulong,
+    pub rdx: ::std::os::raw::c_ulong,
+    pub rsi: ::std::os::raw::c_ulong,
+    pub rdi: ::std::os::raw::c_ulong,
+    pub orig_rax: ::std::os::raw::c_ulong,
+    pub rip: ::std::os::raw::c_ulong,
+    pub cs: ::std::os::raw::c_ulong,
+    pub eflags: ::std::os::raw::c_ulong,
+    pub rsp: ::std::os::raw::c_ulong,
+    pub ss: ::std::os::raw::c_ulong,
+}

aya-common/src/generated/mod.rs

@@ -0,0 +1,28 @@
+#![allow(
+    dead_code,
+    non_camel_case_types,
+    non_snake_case,
+    clippy::all,
+    missing_docs
+)]
+
+#[cfg(target_arch = "aarch64")]
+mod linux_bindings_aarch64;
+#[cfg(target_arch = "arm")]
+mod linux_bindings_armv7;
+#[cfg(target_arch = "riscv64")]
+mod linux_bindings_riscv64;
+#[cfg(target_arch = "x86_64")]
+mod linux_bindings_x86_64;
+
+#[cfg(target_arch = "x86_64")]
+pub use linux_bindings_x86_64::*;
+
+#[cfg(target_arch = "arm")]
+pub use linux_bindings_armv7::*;
+
+#[cfg(target_arch = "aarch64")]
+pub use linux_bindings_aarch64::*;
+
+#[cfg(target_arch = "riscv64")]
+pub use linux_bindings_riscv64::*;

aya-common/src/lib.rs

@@ -0,0 +1,306 @@
+#![no_std]
+
+pub const USDT_MAX_SPEC_COUNT: u32 = 256;
+pub const USDT_MAX_IP_COUNT: u32 = 4 * USDT_MAX_SPEC_COUNT;
+pub const USDT_MAX_ARG_COUNT: usize = 12;
+
+/// The type of argument in a USDT program.
+#[repr(u32)]
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "user", derive(Debug))]
+pub enum UsdtArgType {
+    /// Value is Constant.
+    Const,
+    /// Value is stored in a Register.
+    Reg,
+    /// Value is stored in a Register and requires dereferencing.
+    RegDeref,
+}
+
+impl Default for UsdtArgType {
+    fn default() -> Self {
+        UsdtArgType::Const
+    }
+}
+
+/// The specifcation of an argument in a USDT program.
+#[repr(C)]
+#[derive(Copy, Clone, Default, PartialEq, Eq)]
+#[cfg_attr(feature = "user", derive(Debug))]
+pub struct UsdtArgSpec {
+    /// Meaning of val_off differs based on `arg_type`.
+    /// If Constant, this holds the scalar value of unknow type, up to u64 in size.
+    /// If RegDeref, this contains an offset which is an i64.
+    pub val_off: u64,
+    /// Type of Argument.
+    pub arg_type: UsdtArgType,
+    /// Offset of the register within the BPF context
+    pub reg_off: i16,
+    /// Whether the value should be interpreted as signed
+    pub arg_signed: bool,
+    /// Number of bits that need to be cleared and, optionally,
+    /// sign-extended to cast arguments that are 1, 2, or 4 bytes
+    /// long into final 8-byte u64/s64 value returned to user.
+    pub arg_bitshift: i8,
+}
+
+/// The specification of a USDT
+#[repr(C)]
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "user", derive(Debug))]
+pub struct UsdtSpec {
+    /// Specification used to access arguments.
+    pub args: [UsdtArgSpec; USDT_MAX_ARG_COUNT],
+    /// User supplied cookie since the BPF Attach Cookie is used internally.
+    pub cookie: u64,
+    /// Number of args in this tracepoint
+    pub arg_count: i16,
+}
+
+#[cfg(feature = "user")]
+mod generated;
+
+#[cfg(feature = "user")]
+extern crate std;
+#[cfg(feature = "user")]
+pub mod with_std {
+    use crate::{UsdtArgSpec, UsdtArgType, UsdtSpec, USDT_MAX_ARG_COUNT};
+    use lazy_static::lazy_static;
+    use regex::Regex;
+    use std::{format, string::String};
+    use thiserror::Error;
+
+    #[derive(Error, Debug)]
+    pub enum ParseError {
+        #[error("error parsing usdt arg spec: {0}")]
+        UsdtArgSpecError(String),
+        #[error("error parsing usdt spec: {0}")]
+        UsdtSpecError(String),
+    }
+
+    lazy_static! {
+        static ref USDT_REGEX: Regex =
+            Regex::new(r"^(-?[0-9]+)@((-?[0-9]+)\(%(.*)\)|%(.*)|\$([0-9]+))$").unwrap();
+    }
+
+    impl std::str::FromStr for UsdtArgSpec {
+        type Err = ParseError;
+
+        fn from_str(s: &str) -> Result<Self, Self::Err> {
+            let mut spec = UsdtArgSpec::default();
+            let caps = USDT_REGEX.captures(s).unwrap();
+
+            if caps.len() != 7 {
+                return Err(ParseError::UsdtArgSpecError(format!(
+                    "could not parse {}",
+                    s
+                )));
+            }
+            let mut arg_size: isize = caps.get(1).unwrap().as_str().parse().unwrap();
+            if caps.get(3).is_some() && caps.get(4).is_some() {
+                spec.arg_type = UsdtArgType::RegDeref;
+                spec.val_off = caps.get(3).unwrap().as_str().parse::<i64>().map_err(|e| {
+                    ParseError::UsdtArgSpecError(format!("could not parse {}: {}", s, e))
+                })? as u64;
+                spec.reg_off = calc_pt_regs_offset(caps.get(4).unwrap().as_str())?;
+            } else if caps.get(5).is_some() {
+                spec.arg_type = UsdtArgType::Reg;
+                spec.reg_off = calc_pt_regs_offset(caps.get(5).unwrap().as_str())?;
+            } else if caps.get(6).is_some() {
+                spec.arg_type = UsdtArgType::Const;
+                spec.val_off = caps.get(6).unwrap().as_str().parse::<i64>().map_err(|e| {
+                    ParseError::UsdtArgSpecError(format!("could not parse {}: {}", s, e))
+                })? as u64;
+            }
+            if arg_size < 0 {
+                spec.arg_signed = true;
+                arg_size = -arg_size;
+            }
+            match arg_size {
+                1 | 2 | 4 | 8 => spec.arg_bitshift = (arg_size * 8) as i8,
+                _ => {
+                    return Err(ParseError::UsdtArgSpecError(format!(
+                        "arg size was not 1,2,4,8: {}",
+                        s
+                    )))
+                }
+            }
+            Ok(spec)
+        }
+    }
+
+    #[cfg(target_arch = "x86_64")]
+    fn calc_pt_regs_offset(reg: &str) -> Result<i16, ParseError> {
+        use crate::generated::pt_regs;
+        use memoffset::offset_of;
+        match reg {
+            "rip" | "eip" => Ok(offset_of!(pt_regs, rip) as i16),
+            "rax" | "eax" | "ax" | "al" => Ok(offset_of!(pt_regs, rax) as i16),
+            "rbx" | "ebx" | "bx" | "bl" => Ok(offset_of!(pt_regs, rbx) as i16),
+            "rcx" | "ecx" | "cx" | "cl" => Ok(offset_of!(pt_regs, rcx) as i16),
+            "rdx" | "edx" | "dx" | "dl" => Ok(offset_of!(pt_regs, rdx) as i16),
+            "rsi" | "esi" | "si" | "sil" => Ok(offset_of!(pt_regs, rsi) as i16),
+            "rdi" | "edi" | "di" | "dil" => Ok(offset_of!(pt_regs, rdi) as i16),
+            "rbp" | "ebp" | "bp" | "bpl" => Ok(offset_of!(pt_regs, rbp) as i16),
+            "rsp" | "esp" | "sp" | "bsl" => Ok(offset_of!(pt_regs, rsp) as i16),
+            "r8" | "r8d" | "r8w" | "r8b" => Ok(offset_of!(pt_regs, r8) as i16),
+            "r9" | "r9d" | "r9w" | "r9b" => Ok(offset_of!(pt_regs, r9) as i16),
+            "r10" | "r10d" | "r10w" | "r10b" => Ok(offset_of!(pt_regs, r10) as i16),
+            "r11" | "r11d" | "r11w" | "r11b" => Ok(offset_of!(pt_regs, r11) as i16),
+            "r12" | "r12d" | "r12w" | "r12b" => Ok(offset_of!(pt_regs, r12) as i16),
+            "r13" | "r13d" | "r13w" | "r13b" => Ok(offset_of!(pt_regs, r13) as i16),
+            "r14" | "r14d" | "r14w" | "r14b" => Ok(offset_of!(pt_regs, r14) as i16),
+            "r15" | "r15d" | "r15w" | "r15b" => Ok(offset_of!(pt_regs, r15) as i16),
+            _ => Err(ParseError::UsdtArgSpecError(format!(
+                "unknown register: {}",
+                reg
+            ))),
+        }
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    fn calc_pt_regs_offset(reg: &str) -> Result<i16, ParseError> {
+        use crate::generated::user_pt_regs;
+        use memoffset::offset_of;
+        use std::mem;
+        match reg {
+            r if r.starts_with('x') => {
+                let n: usize = r.strip_prefix('x').unwrap().parse().map_err(|_| {
+                    ParseError::UsdtArgSpecError(format!(
+                        "invalid register format. expected: xN. got: {}",
+                        reg
+                    ))
+                })?;
+                Ok((offset_of!(user_pt_regs, regs) + (n * mem::size_of::<u64>())) as i16)
+            }
+            "sp" => Ok(offset_of!(user_pt_regs, sp) as i16),
+            _ => Err(ParseError::UsdtArgSpecError(format!(
+                "unknown register: {}",
+                reg
+            ))),
+        }
+    }
+
+    #[cfg(target_arch = "arm")]
+    fn calc_pt_regs_offset(reg: &str) -> Result<i16, ParseError> {
+        use crate::generated::pt_regs;
+        use memoffset::offset_of;
+        use std::mem;
+        match reg {
+            r if r.starts_with('r') => {
+                let n: usize = r.strip_prefix('r').unwrap().parse().map_err(|_| {
+                    ParseError::UsdtArgSpecError(format!(
+                        "invalid register format. expected: rN. got: {}",
+                        reg
+                    ))
+                })?;
+                Ok((offset_of!(pt_regs, uregs) + (n * mem::size_of::<u32>())) as i16)
+            }
+            _ => Err(ParseError::UsdtArgSpecError(format!(
+                "unknown register: {}",
+                reg
+            ))),
+        }
+    }
+
+    #[cfg(target_arch = "riscv64")]
+    fn calc_pt_regs_offset(reg: &str) -> Result<i16, ParseError> {
+        unimplemented!("riscv support for usdt probes not implemented")
+    }
+
+    impl std::str::FromStr for UsdtSpec {
+        type Err = ParseError;
+
+        fn from_str(s: &str) -> Result<Self, Self::Err> {
+            use std::vec::Vec;
+            let parts: Vec<&str> = s.split_whitespace().collect();
+            if parts.len() > USDT_MAX_ARG_COUNT {
+                return Err(ParseError::UsdtSpecError(format!("too many args: {}", s)));
+            }
+            let mut args = parts
+                .iter()
+                .map(|s| s.parse::<UsdtArgSpec>().unwrap())
+                .collect::<Vec<UsdtArgSpec>>();
+            let arg_count = args.len() as i16;
+            args.resize(USDT_MAX_ARG_COUNT, UsdtArgSpec::default());
+            Ok(UsdtSpec {
+                args: args.try_into().unwrap(),
+                cookie: 0,
+                arg_count,
+            })
+        }
+    }
+
+    #[cfg(test)]
+    mod test {
+        use super::*;
+        use memoffset::offset_of;
+
+        #[cfg_attr(target_arch = "x86_64", test)]
+        fn test_parse_specs() {
+            use crate::generated::pt_regs;
+
+            let s = "-8@%rax -8@%rcx";
+            let res: UsdtSpec = s.parse().unwrap();
+            assert_eq!(res.arg_count, 2);
+            assert_eq!(
+                res.args[0],
+                UsdtArgSpec {
+                    val_off: 0,
+                    arg_type: UsdtArgType::Reg,
+                    reg_off: offset_of!(pt_regs, rax) as i16,
+                    arg_signed: true,
+                    arg_bitshift: 64
+                }
+            );
+            assert_eq!(
+                res.args[1],
+                UsdtArgSpec {
+                    val_off: 0,
+                    arg_type: UsdtArgType::Reg,
+                    reg_off: offset_of!(pt_regs, rcx) as i16,
+                    arg_signed: true,
+                    arg_bitshift: 64
+                }
+            );
+        }
+
+        #[cfg_attr(target_arch = "x86_64", test)]
+        fn test_parse_args() {
+            use crate::generated::pt_regs;
+
+            assert_eq!(
+                "-4@-1204(%rbp)".parse::<UsdtArgSpec>().unwrap(),
+                UsdtArgSpec {
+                    val_off: -1204i64 as u64,
+                    arg_type: UsdtArgType::RegDeref,
+                    reg_off: offset_of!(pt_regs, rbp) as i16,
+                    arg_signed: true,
+                    arg_bitshift: 32
+                }
+            );
+
+            assert_eq!(
+                "-4@%edi".parse::<UsdtArgSpec>().unwrap(),
+                UsdtArgSpec {
+                    val_off: 0,
+                    arg_type: UsdtArgType::Reg,
+                    reg_off: offset_of!(pt_regs, rdi) as i16,
+                    arg_signed: true,
+                    arg_bitshift: 32
+                }
+            );
+
+            assert_eq!(
+                "-4@$5".parse::<UsdtArgSpec>().unwrap(),
+                UsdtArgSpec {
+                    val_off: 5,
+                    arg_type: UsdtArgType::Const,
+                    reg_off: 0,
+                    arg_signed: true,
+                    arg_bitshift: 32
+                }
+            );
+        }
+    }
+}

aya-log/src/lib.rs

@@ -480,7 +480,6 @@ mod test {
     use super::*;
     use aya_log_common::{write_record_header, WriteToBuf};
     use log::logger;
-    use testing_logger;
 
     fn new_log(args: usize) -> Result<(usize, Vec<u8>), ()> {
         let mut buf = vec![0; 8192];

aya/Cargo.toml

@@ -21,6 +21,7 @@ parking_lot = { version = "0.12.0", features = ["send_guard"] }
 tokio = { version = "1.2.0", features = ["macros", "rt", "rt-multi-thread", "net"], optional = true }
 async-io = { version = "1.3", optional = true }
 log = "0.4"
+aya-common = { version = "0.1.0", path = "../aya-common", features = ["user"] }
 
 [dev-dependencies]
 matches = "0.1.8"

aya/src/bpf.rs

@@ -25,12 +25,13 @@ use crate::{
         BtfTracePoint, CgroupSkb, CgroupSkbAttachType, CgroupSock, CgroupSockAddr, CgroupSockopt,
         CgroupSysctl, Extension, FEntry, FExit, KProbe, LircMode2, Lsm, PerfEvent, ProbeKind,
         Program, ProgramData, ProgramError, RawTracePoint, SchedClassifier, SkLookup, SkMsg, SkSkb,
-        SkSkbKind, SockOps, SocketFilter, TracePoint, UProbe, Xdp,
+        SkSkbKind, SockOps, SocketFilter, TracePoint, UProbe, Usdt, Xdp,
     },
     sys::{
-        bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_btf_datasec_supported,
-        is_btf_decl_tag_supported, is_btf_float_supported, is_btf_func_global_supported,
-        is_btf_func_supported, is_btf_supported, is_btf_type_tag_supported, is_prog_name_supported,
+        bpf_load_btf, bpf_map_freeze, bpf_map_update_elem_ptr, is_bpf_cookie_supported,
+        is_btf_datasec_supported, is_btf_decl_tag_supported, is_btf_float_supported,
+        is_btf_func_global_supported, is_btf_func_supported, is_btf_supported,
+        is_btf_type_tag_supported, is_perf_link_supported, is_prog_name_supported,
         retry_with_verifier_logs,
     },
     util::{bytes_of, possible_cpus, VerifierLog, POSSIBLE_CPUS},
@@ -116,10 +117,16 @@ impl Default for PinningType {
     }
 }
 
+lazy_static! {
+    pub(crate) static ref FEATURES: Features = Features::new();
+}
+
 // Features implements BPF and BTF feature detection
 #[derive(Default, Debug)]
 pub(crate) struct Features {
     pub bpf_name: bool,
+    pub bpf_cookie: bool,
+    pub bpf_perf_link: bool,
     pub btf: bool,
     pub btf_func: bool,
     pub btf_func_global: bool,
@@ -130,38 +137,45 @@ pub(crate) struct Features {
 }
 
 impl Features {
-    fn probe_features(&mut self) {
-        self.bpf_name = is_prog_name_supported();
-        debug!("[FEAT PROBE] BPF program name support: {}", self.bpf_name);
-
-        self.btf = is_btf_supported();
-        debug!("[FEAT PROBE] BTF support: {}", self.btf);
+    fn new() -> Self {
+        let mut f = Features {
+            bpf_name: is_prog_name_supported(),
+            bpf_cookie: is_bpf_cookie_supported(),
+            bpf_perf_link: is_perf_link_supported(),
+            btf: is_btf_supported(),
+            ..Default::default()
+        };
+        debug!("[FEAT PROBE] BPF program name support: {}", f.bpf_name);
+        debug!("[FEAT PROBE] BPF cookie support: {}", f.bpf_cookie);
+        debug!("[FEAT PROBE] BPF probe link support: {}", f.bpf_perf_link);
+        debug!("[FEAT PROBE] BTF support: {}", f.btf);
 
-        if self.btf {
-            self.btf_func = is_btf_func_supported();
-            debug!("[FEAT PROBE] BTF func support: {}", self.btf_func);
+        if f.btf {
+            f.btf_func = is_btf_func_supported();
+            debug!("[FEAT PROBE] BTF func support: {}", f.btf_func);
 
-            self.btf_func_global = is_btf_func_global_supported();
+            f.btf_func_global = is_btf_func_global_supported();
             debug!(
                 "[FEAT PROBE] BTF global func support: {}",
-                self.btf_func_global
+                f.btf_func_global
             );
 
-            self.btf_datasec = is_btf_datasec_supported();
+            f.btf_datasec = is_btf_datasec_supported();
             debug!(
                 "[FEAT PROBE] BTF var and datasec support: {}",
-                self.btf_datasec
+                f.btf_datasec
             );
 
-            self.btf_float = is_btf_float_supported();
-            debug!("[FEAT PROBE] BTF float support: {}", self.btf_float);
+            f.btf_float = is_btf_float_supported();
+            debug!("[FEAT PROBE] BTF float support: {}", f.btf_float);
 
-            self.btf_decl_tag = is_btf_decl_tag_supported();
-            debug!("[FEAT PROBE] BTF decl_tag support: {}", self.btf_decl_tag);
+            f.btf_decl_tag = is_btf_decl_tag_supported();
+            debug!("[FEAT PROBE] BTF decl_tag support: {}", f.btf_decl_tag);
 
-            self.btf_type_tag = is_btf_type_tag_supported();
-            debug!("[FEAT PROBE] BTF type_tag support: {}", self.btf_type_tag);
+            f.btf_type_tag = is_btf_type_tag_supported();
+            debug!("[FEAT PROBE] BTF type_tag support: {}", f.btf_type_tag);
         }
+        f
     }
 }
 
@@ -192,7 +206,6 @@ pub struct BpfLoader<'a> {
     map_pin_path: Option<PathBuf>,
     globals: HashMap<&'a str, &'a [u8]>,
     max_entries: HashMap<&'a str, u32>,
-    features: Features,
     extensions: HashSet<&'a str>,
     verifier_log_level: VerifierLogLevel,
 }
@@ -222,14 +235,11 @@ impl Default for VerifierLogLevel {
 impl<'a> BpfLoader<'a> {
     /// Creates a new loader instance.
     pub fn new() -> BpfLoader<'a> {
-        let mut features = Features::default();
-        features.probe_features();
         BpfLoader {
             btf: Btf::from_sys_fs().ok().map(Cow::Owned),
             map_pin_path: None,
             globals: HashMap::new(),
             max_entries: HashMap::new(),
-            features,
             extensions: HashSet::new(),
             verifier_log_level: VerifierLogLevel::default(),
         }
@@ -413,12 +423,12 @@ impl<'a> BpfLoader<'a> {
         let mut obj = Object::parse(data)?;
         obj.patch_map_data(self.globals.clone())?;
 
-        let btf_fd = if self.features.btf {
+        let btf_fd = if FEATURES.btf {
             if let Some(ref mut obj_btf) = obj.btf {
                 // fixup btf
                 let section_data = obj.section_sizes.clone();
                 let symbol_offsets = obj.symbol_offset_by_name.clone();
-                obj_btf.fixup_and_sanitize(&section_data, &symbol_offsets, &self.features)?;
+                obj_btf.fixup_and_sanitize(&section_data, &symbol_offsets, &FEATURES)?;
                 // load btf to the kernel
                 let raw_btf = obj_btf.to_bytes();
                 Some(load_btf(raw_btf)?)
@@ -504,7 +514,7 @@ impl<'a> BpfLoader<'a> {
             .programs
             .drain()
             .map(|(name, obj)| {
-                let prog_name = if self.features.bpf_name {
+                let prog_name = if FEATURES.bpf_name {
                     Some(name.clone())
                 } else {
                     None
@@ -633,6 +643,9 @@ impl<'a> BpfLoader<'a> {
                                 attach_type: *attach_type,
                             })
                         }
+                        ProgramSection::Usdt { .. } => Program::Usdt(Usdt {
+                            data: ProgramData::new(prog_name, obj, btf_fd, verifier_log_level),
+                        }),
                     }
                 };
                 (name, program)
@@ -933,6 +946,11 @@ pub enum BpfError {
     #[error("program error")]
     /// A program error
     ProgramError(#[from] ProgramError),
+
+    /// Required map not found
+    #[error("required map {0} not found. did you enable the usdt feature (aya) or include usdt.bpf.h (libbpf)?")]
+    /// A program error
+    MissingMap(String),
 }
 
 fn load_btf(raw_btf: Vec<u8>) -> Result<RawFd, BtfError> {

aya/src/obj/mod.rs

@@ -293,6 +293,9 @@ pub enum ProgramSection {
         name: String,
         attach_type: CgroupSockAttachType,
     },
+    Usdt {
+        name: String,
+    },
 }
 
 impl ProgramSection {
@@ -326,6 +329,7 @@ impl ProgramSection {
             ProgramSection::Extension { name } => name,
             ProgramSection::SkLookup { name } => name,
             ProgramSection::CgroupSock { name, .. } => name,
+            ProgramSection::Usdt { name } => name,
         }
     }
 }
@@ -349,6 +353,7 @@ impl FromStr for ProgramSection {
             "kprobe" => KProbe { name },
             "kretprobe" => KRetProbe { name },
             "uprobe" => UProbe { name },
+            "usdt" => Usdt { name },
             "uretprobe" => URetProbe { name },
             "xdp" => Xdp { name },
             "tp_btf" => BtfTracePoint { name },
@@ -1483,7 +1488,6 @@ mod tests {
             map_flags: 5,
             id: 0,
             pinning: PinningType::None,
-            ..Default::default()
         };
 
         assert_eq!(
@@ -1502,7 +1506,6 @@ mod tests {
             map_flags: 5,
             id: 6,
             pinning: PinningType::ByName,
-            ..Default::default()
         };
 
         assert_eq!(parse_map_def("foo", bytes_of(&def)).unwrap(), def);
@@ -1518,7 +1521,6 @@ mod tests {
             map_flags: 5,
             id: 6,
             pinning: PinningType::ByName,
-            ..Default::default()
         };
         let mut buf = [0u8; 128];
         unsafe { ptr::write_unaligned(buf.as_mut_ptr() as *mut _, def) };
@@ -1549,7 +1551,6 @@ mod tests {
                         map_flags: 5,
                         id: 0,
                         pinning: PinningType::None,
-                        ..Default::default()
                     })
                 ),
                 "foo"
@@ -1684,7 +1685,7 @@ mod tests {
         buf.extend(&map_data);
         buf.extend(&map_data);
         // throw in some padding
-        buf.extend(&[0, 0, 0, 0]);
+        buf.extend([0, 0, 0, 0]);
         buf.extend(&map_data);
         assert_matches!(
             obj.parse_section(fake_section(BpfSectionKind::Maps, "maps", buf.as_slice(),)),
@@ -2218,7 +2219,6 @@ mod tests {
                     map_flags: BPF_F_RDONLY_PROG,
                     id: 1,
                     pinning: PinningType::None,
-                    ..Default::default()
                 },
                 section_index: 1,
                 symbol_index: 1,

aya/src/programs/cgroup_skb.rs

@@ -97,7 +97,7 @@ impl CgroupSkb {
         };
         let k_ver = kernel_version().unwrap();
         if k_ver >= (5, 7, 0) {
-            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, 0).map_err(
+            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, None, 0).map_err(
                 |(_, io_error)| ProgramError::SyscallError {
                     call: "bpf_link_create".to_owned(),
                     io_error,

aya/src/programs/cgroup_sock.rs

@@ -73,7 +73,7 @@ impl CgroupSock {
         let attach_type = self.data.expected_attach_type.unwrap();
         let k_ver = kernel_version().unwrap();
         if k_ver >= (5, 7, 0) {
-            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, 0).map_err(
+            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, None, 0).map_err(
                 |(_, io_error)| ProgramError::SyscallError {
                     call: "bpf_link_create".to_owned(),
                     io_error,

aya/src/programs/cgroup_sock_addr.rs

@@ -74,7 +74,7 @@ impl CgroupSockAddr {
         let attach_type = self.data.expected_attach_type.unwrap();
         let k_ver = kernel_version().unwrap();
         if k_ver >= (5, 7, 0) {
-            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, 0).map_err(
+            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, None, 0).map_err(
                 |(_, io_error)| ProgramError::SyscallError {
                     call: "bpf_link_create".to_owned(),
                     io_error,

aya/src/programs/cgroup_sockopt.rs

@@ -71,7 +71,7 @@ impl CgroupSockopt {
         let attach_type = self.data.expected_attach_type.unwrap();
         let k_ver = kernel_version().unwrap();
         if k_ver >= (5, 7, 0) {
-            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, 0).map_err(
+            let link_fd = bpf_link_create(prog_fd, cgroup_fd, attach_type, None, None, 0).map_err(
                 |(_, io_error)| ProgramError::SyscallError {
                     call: "bpf_link_create".to_owned(),
                     io_error,

aya/src/programs/cgroup_sysctl.rs

@@ -66,12 +66,11 @@ impl CgroupSysctl {
 
         let k_ver = kernel_version().unwrap();
         if k_ver >= (5, 7, 0) {
-            let link_fd = bpf_link_create(prog_fd, cgroup_fd, BPF_CGROUP_SYSCTL, None, 0).map_err(
-                |(_, io_error)| ProgramError::SyscallError {
+            let link_fd = bpf_link_create(prog_fd, cgroup_fd, BPF_CGROUP_SYSCTL, None, None, 0)
+                .map_err(|(_, io_error)| ProgramError::SyscallError {
                     call: "bpf_link_create".to_owned(),
                     io_error,
-                },
-            )? as RawFd;
+                })? as RawFd;
             self.data
                 .links
                 .insert(CgroupSysctlLink(CgroupSysctlLinkInner::Fd(FdLink::new(

aya/src/programs/extension.rs

@@ -90,11 +90,18 @@ impl Extension {
         let target_fd = self.data.attach_prog_fd.ok_or(ProgramError::NotLoaded)?;
         let btf_id = self.data.attach_btf_id.ok_or(ProgramError::NotLoaded)?;
         // the attach type must be set as 0, which is bpf_attach_type::BPF_CGROUP_INET_INGRESS
-        let link_fd = bpf_link_create(prog_fd, target_fd, BPF_CGROUP_INET_INGRESS, Some(btf_id), 0)
-            .map_err(|(_, io_error)| ProgramError::SyscallError {
-                call: "bpf_link_create".to_owned(),
-                io_error,
-            })? as RawFd;
+        let link_fd = bpf_link_create(
+            prog_fd,
+            target_fd,
+            BPF_CGROUP_INET_INGRESS,
+            Some(btf_id),
+            None,
+            0,
+        )
+        .map_err(|(_, io_error)| ProgramError::SyscallError {
+            call: "bpf_link_create".to_owned(),
+            io_error,
+        })? as RawFd;
         self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
     }
 
@@ -118,11 +125,18 @@ impl Extension {
         let (_, btf_id) = get_btf_info(target_fd, func_name)?;
         let prog_fd = self.data.fd_or_err()?;
         // the attach type must be set as 0, which is bpf_attach_type::BPF_CGROUP_INET_INGRESS
-        let link_fd = bpf_link_create(prog_fd, target_fd, BPF_CGROUP_INET_INGRESS, Some(btf_id), 0)
-            .map_err(|(_, io_error)| ProgramError::SyscallError {
-                call: "bpf_link_create".to_owned(),
-                io_error,
-            })? as RawFd;
+        let link_fd = bpf_link_create(
+            prog_fd,
+            target_fd,
+            BPF_CGROUP_INET_INGRESS,
+            Some(btf_id),
+            None,
+            0,
+        )
+        .map_err(|(_, io_error)| ProgramError::SyscallError {
+            call: "bpf_link_create".to_owned(),
+            io_error,
+        })? as RawFd;
         self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
     }
 

aya/src/programs/kprobe.rs

@@ -6,7 +6,7 @@ use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
     programs::{
         define_link_wrapper, load_program,
-        perf_attach::{PerfLink, PerfLinkId},
+        perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{attach, ProbeKind},
         ProgramData, ProgramError,
     },
@@ -90,8 +90,8 @@ define_link_wrapper!(
     KProbeLink,
     /// The type returned by [KProbe::attach]. Can be passed to [KProbe::detach].
     KProbeLinkId,
-    PerfLink,
-    PerfLinkId
+    PerfLinkInner,
+    PerfLinkIdInner
 );
 
 /// The type returned when attaching a [`KProbe`] fails.

aya/src/programs/mod.rs

@@ -60,6 +60,7 @@ pub mod tc;
 pub mod tp_btf;
 pub mod trace_point;
 pub mod uprobe;
+pub mod usdt;
 mod utils;
 pub mod xdp;
 
@@ -98,6 +99,7 @@ pub use tc::{SchedClassifier, TcAttachType, TcError};
 pub use tp_btf::BtfTracePoint;
 pub use trace_point::{TracePoint, TracePointError};
 pub use uprobe::{UProbe, UProbeError};
+pub use usdt::{Usdt, UsdtError};
 pub use xdp::{Xdp, XdpError, XdpFlags};
 
 use crate::{
@@ -198,6 +200,10 @@ pub enum ProgramError {
     #[error(transparent)]
     Btf(#[from] BtfError),
 
+    /// An error occurred while working with a Usdt program.
+    #[error(transparent)]
+    Usdt(#[from] UsdtError),
+
     /// The program is not attached.
     #[error("the program name `{name}` is invalid")]
     InvalidName {
@@ -265,6 +271,8 @@ pub enum Program {
     SkLookup(SkLookup),
     /// A [`CgroupSock`] program
     CgroupSock(CgroupSock),
+    /// A [`Usdt`] program
+    Usdt(Usdt),
 }
 
 impl Program {
@@ -295,6 +303,7 @@ impl Program {
             Program::CgroupSockAddr(_) => BPF_PROG_TYPE_CGROUP_SOCK_ADDR,
             Program::SkLookup(_) => BPF_PROG_TYPE_SK_LOOKUP,
             Program::CgroupSock(_) => BPF_PROG_TYPE_CGROUP_SOCK,
+            Program::Usdt(_) => BPF_PROG_TYPE_KPROBE,
         }
     }
 
@@ -324,6 +333,7 @@ impl Program {
             Program::CgroupSockAddr(p) => p.pin(path),
             Program::SkLookup(p) => p.pin(path),
             Program::CgroupSock(p) => p.pin(path),
+            Program::Usdt(p) => p.pin(path),
         }
     }
 
@@ -353,6 +363,7 @@ impl Program {
             Program::CgroupSockAddr(p) => p.unload(),
             Program::SkLookup(p) => p.unload(),
             Program::CgroupSock(p) => p.unload(),
+            Program::Usdt(p) => p.unload(),
         }
     }
 
@@ -385,6 +396,7 @@ impl Program {
             Program::CgroupSockAddr(p) => p.fd(),
             Program::SkLookup(p) => p.fd(),
             Program::CgroupSock(p) => p.fd(),
+            Program::Usdt(p) => p.fd(),
         }
     }
 }
@@ -637,6 +649,7 @@ impl_program_unload!(
     SkLookup,
     SockOps,
     CgroupSock,
+    Usdt,
 );
 
 macro_rules! impl_fd {
@@ -676,6 +689,7 @@ impl_fd!(
     SkLookup,
     SockOps,
     CgroupSock,
+    Usdt,
 );
 
 macro_rules! impl_program_pin{
@@ -720,6 +734,7 @@ impl_program_pin!(
     SkLookup,
     SockOps,
     CgroupSock,
+    Usdt,
 );
 
 macro_rules! impl_try_from_program {
@@ -774,6 +789,7 @@ impl_try_from_program!(
     CgroupSockAddr,
     SkLookup,
     CgroupSock,
+    Usdt,
 );
 
 /// Provides information about a loaded program, like name, id and statistics

aya/src/programs/perf_attach.rs

@@ -3,11 +3,44 @@ use libc::close;
 use std::os::unix::io::RawFd;
 
 use crate::{
-    programs::{probe::detach_debug_fs, Link, ProbeKind, ProgramData, ProgramError},
-    sys::perf_event_ioctl,
-    PERF_EVENT_IOC_DISABLE, PERF_EVENT_IOC_ENABLE, PERF_EVENT_IOC_SET_BPF,
+    generated::bpf_attach_type::BPF_PERF_EVENT,
+    programs::{probe::detach_debug_fs, Link, ProbeKind, ProgramError},
+    sys::{bpf_link_create, perf_event_ioctl},
+    FEATURES, PERF_EVENT_IOC_DISABLE, PERF_EVENT_IOC_ENABLE, PERF_EVENT_IOC_SET_BPF,
 };
 
+use crate::programs::links::FdLink;
+
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub(crate) enum PerfLinkIdInner {
+    FdLinkId(<FdLink as Link>::Id),
+    PerfLinkId(<PerfLink as Link>::Id),
+}
+
+#[derive(Debug)]
+pub(crate) enum PerfLinkInner {
+    FdLink(FdLink),
+    PerfLink(PerfLink),
+}
+
+impl Link for PerfLinkInner {
+    type Id = PerfLinkIdInner;
+
+    fn id(&self) -> Self::Id {
+        match self {
+            PerfLinkInner::FdLink(link) => PerfLinkIdInner::FdLinkId(link.id()),
+            PerfLinkInner::PerfLink(link) => PerfLinkIdInner::PerfLinkId(link.id()),
+        }
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        match self {
+            PerfLinkInner::FdLink(link) => link.detach(),
+            PerfLinkInner::PerfLink(link) => link.detach(),
+        }
+    }
+}
+
 /// The identifer of a PerfLink.
 #[derive(Debug, Hash, Eq, PartialEq)]
 pub struct PerfLinkId(RawFd);
@@ -41,29 +74,39 @@ impl Link for PerfLink {
     }
 }
 
-pub(crate) fn perf_attach<T: Link + From<PerfLink>>(
-    data: &mut ProgramData<T>,
+pub(crate) fn perf_attach(
+    prog_fd: RawFd,
     fd: RawFd,
-) -> Result<T::Id, ProgramError> {
-    perf_attach_either(data, fd, None, None)
+    cookie: Option<u64>,
+) -> Result<PerfLinkInner, ProgramError> {
+    if FEATURES.bpf_perf_link {
+        let link_fd = bpf_link_create(prog_fd, fd, BPF_PERF_EVENT, None, cookie, 0).map_err(
+            |(_, io_error)| ProgramError::SyscallError {
+                call: "bpf_link_create".to_owned(),
+                io_error,
+            },
+        )? as RawFd;
+        Ok(PerfLinkInner::FdLink(FdLink::new(link_fd)))
+    } else {
+        perf_attach_either(prog_fd, fd, None, None)
+    }
 }
 
-pub(crate) fn perf_attach_debugfs<T: Link + From<PerfLink>>(
-    data: &mut ProgramData<T>,
+pub(crate) fn perf_attach_debugfs(
+    prog_fd: RawFd,
     fd: RawFd,
     probe_kind: ProbeKind,
     event_alias: String,
-) -> Result<T::Id, ProgramError> {
-    perf_attach_either(data, fd, Some(probe_kind), Some(event_alias))
+) -> Result<PerfLinkInner, ProgramError> {
+    perf_attach_either(prog_fd, fd, Some(probe_kind), Some(event_alias))
 }
 
-fn perf_attach_either<T: Link + From<PerfLink>>(
-    data: &mut ProgramData<T>,
+fn perf_attach_either(
+    prog_fd: RawFd,
     fd: RawFd,
     probe_kind: Option<ProbeKind>,
     event_alias: Option<String>,
-) -> Result<T::Id, ProgramError> {
-    let prog_fd = data.fd_or_err()?;
+) -> Result<PerfLinkInner, ProgramError> {
     perf_event_ioctl(fd, PERF_EVENT_IOC_SET_BPF, prog_fd).map_err(|(_, io_error)| {
         ProgramError::SyscallError {
             call: "PERF_EVENT_IOC_SET_BPF".to_owned(),
@@ -77,12 +120,9 @@ fn perf_attach_either<T: Link + From<PerfLink>>(
         }
     })?;
 
-    data.links.insert(
-        PerfLink {
-            perf_fd: fd,
-            probe_kind,
-            event_alias,
-        }
-        .into(),
-    )
+    Ok(PerfLinkInner::PerfLink(PerfLink {
+        perf_fd: fd,
+        probe_kind,
+        event_alias,
+    }))
 }

aya/src/programs/perf_event.rs

@@ -12,8 +12,9 @@ use crate::{
         },
     },
     programs::{
+        links::define_link_wrapper,
         load_program, perf_attach,
-        perf_attach::{PerfLink, PerfLinkId},
+        perf_attach::{PerfLinkIdInner, PerfLinkInner},
         ProgramData, ProgramError,
     },
     sys::perf_event_open,
@@ -118,7 +119,7 @@ pub enum PerfEventScope {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_PERF_EVENT")]
 pub struct PerfEvent {
-    pub(crate) data: ProgramData<PerfLink>,
+    pub(crate) data: ProgramData<PerfEventLink>,
 }
 
 impl PerfEvent {
@@ -140,7 +141,7 @@ impl PerfEvent {
         config: u64,
         scope: PerfEventScope,
         sample_policy: SamplePolicy,
-    ) -> Result<PerfLinkId, ProgramError> {
+    ) -> Result<PerfEventLinkId, ProgramError> {
         let (sample_period, sample_frequency) = match sample_policy {
             SamplePolicy::Period(period) => (period, None),
             SamplePolicy::Frequency(frequency) => (0, Some(frequency)),
@@ -167,13 +168,14 @@ impl PerfEvent {
             io_error,
         })? as i32;
 
-        perf_attach(&mut self.data, fd)
+        let link = perf_attach(self.data.fd_or_err()?, fd, None)?;
+        self.data.links.insert(PerfEventLink(link))
     }
 
     /// Detaches the program.
     ///
     /// See [PerfEvent::attach].
-    pub fn detach(&mut self, link_id: PerfLinkId) -> Result<(), ProgramError> {
+    pub fn detach(&mut self, link_id: PerfEventLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
 
@@ -181,7 +183,16 @@ impl PerfEvent {
     ///
     /// The link will be detached on `Drop` and the caller is now responsible
     /// for managing its lifetime.
-    pub fn take_link(&mut self, link_id: PerfLinkId) -> Result<PerfLink, ProgramError> {
+    pub fn take_link(&mut self, link_id: PerfEventLinkId) -> Result<PerfEventLink, ProgramError> {
         self.data.take_link(link_id)
     }
 }
+
+define_link_wrapper!(
+    /// The link used by [PerfEvent] programs.
+    PerfEventLink,
+    /// The type returned by [PerfEvent::attach]. Can be passed to [PerfEvent::detach].
+    PerfEventLinkId,
+    PerfLinkInner,
+    PerfLinkIdInner
+);

aya/src/programs/probe.rs

@@ -7,7 +7,7 @@ use std::{
 
 use crate::{
     programs::{
-        kprobe::KProbeError, perf_attach, perf_attach::PerfLink, perf_attach_debugfs,
+        kprobe::KProbeError, perf_attach, perf_attach::PerfLinkInner, perf_attach_debugfs,
         trace_point::read_sys_fs_trace_point_id, uprobe::UProbeError, Link, ProgramData,
         ProgramError,
     },
@@ -36,7 +36,7 @@ impl ProbeKind {
     }
 }
 
-pub(crate) fn attach<T: Link + From<PerfLink>>(
+pub(crate) fn attach<T: Link + From<PerfLinkInner>>(
     program_data: &mut ProgramData<T>,
     kind: ProbeKind,
     fn_name: &str,
@@ -49,12 +49,19 @@ pub(crate) fn attach<T: Link + From<PerfLink>>(
     if k_ver < (4, 17, 0) {
         let (fd, event_alias) = create_as_trace_point(kind, fn_name, offset, pid)?;
 
-        return perf_attach_debugfs(program_data, fd, kind, event_alias);
+        let link = T::from(perf_attach_debugfs(
+            program_data.fd_or_err()?,
+            fd,
+            kind,
+            event_alias,
+        )?);
+        return program_data.links.insert(link);
     };
 
-    let fd = create_as_probe(kind, fn_name, offset, pid)?;
+    let fd = create_as_probe(kind, fn_name, offset, pid, None)?;
 
-    perf_attach(program_data, fd)
+    let link = T::from(perf_attach(program_data.fd_or_err()?, fd, None)?);
+    program_data.links.insert(link)
 }
 
 pub(crate) fn detach_debug_fs(kind: ProbeKind, event_alias: &str) -> Result<(), ProgramError> {
@@ -70,11 +77,12 @@ pub(crate) fn detach_debug_fs(kind: ProbeKind, event_alias: &str) -> Result<(),
     Ok(())
 }
 
-fn create_as_probe(
+pub(crate) fn create_as_probe(
     kind: ProbeKind,
     fn_name: &str,
     offset: u64,
     pid: Option<pid_t>,
+    ref_cnt_offset: Option<u64>,
 ) -> Result<i32, ProgramError> {
     use ProbeKind::*;
 
@@ -97,7 +105,7 @@ fn create_as_probe(
         _ => None,
     };
 
-    let fd = perf_event_open_probe(perf_ty, ret_bit, fn_name, offset, pid).map_err(
+    let fd = perf_event_open_probe(perf_ty, ret_bit, fn_name, offset, pid, ref_cnt_offset).map_err(
         |(_code, io_error)| ProgramError::SyscallError {
             call: "perf_event_open".to_owned(),
             io_error,

aya/src/programs/sk_lookup.rs

@@ -64,7 +64,7 @@ impl SkLookup {
         let prog_fd = self.data.fd_or_err()?;
         let netns_fd = netns.as_raw_fd();
 
-        let link_fd = bpf_link_create(prog_fd, netns_fd, BPF_SK_LOOKUP, None, 0).map_err(
+        let link_fd = bpf_link_create(prog_fd, netns_fd, BPF_SK_LOOKUP, None, None, 0).map_err(
             |(_, io_error)| ProgramError::SyscallError {
                 call: "bpf_link_create".to_owned(),
                 io_error,

aya/src/programs/trace_point.rs

@@ -6,7 +6,7 @@ use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT,
     programs::{
         define_link_wrapper, load_program,
-        perf_attach::{perf_attach, PerfLink, PerfLinkId},
+        perf_attach::{perf_attach, PerfLinkIdInner, PerfLinkInner},
         ProgramData, ProgramError,
     },
     sys::perf_event_open_trace_point,
@@ -85,7 +85,8 @@ impl TracePoint {
             }
         })? as i32;
 
-        perf_attach(&mut self.data, fd)
+        let link = TracePointLink(perf_attach(self.data.fd_or_err()?, fd, None)?);
+        self.data.links.insert(link)
     }
 
     /// Detaches from a trace point.
@@ -109,8 +110,8 @@ define_link_wrapper!(
     TracePointLink,
     /// The type returned by [TracePoint::attach]. Can be passed to [TracePoint::detach].
     TracePointLinkId,
-    PerfLink,
-    PerfLinkId
+    PerfLinkInner,
+    PerfLinkIdInner
 );
 
 pub(crate) fn read_sys_fs_trace_point_id(

aya/src/programs/uprobe.rs

@@ -1,13 +1,8 @@
 //! User space probes.
 use libc::pid_t;
-use object::{Object, ObjectSymbol};
 use std::{
     error::Error,
-    ffi::CStr,
-    fs,
-    io::{self, BufRead, Cursor, Read},
-    mem,
-    os::raw::c_char,
+    io,
     path::{Path, PathBuf},
     sync::Arc,
 };
@@ -17,20 +12,13 @@ use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
     programs::{
         define_link_wrapper, load_program,
-        perf_attach::{PerfLink, PerfLinkId},
+        perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{attach, ProbeKind},
+        utils::{resolve_symbol, ProcMap, ProcMapError, LD_SO_CACHE, LD_SO_CACHE_FILE},
         ProgramData, ProgramError,
     },
 };
 
-const LD_SO_CACHE_FILE: &str = "/etc/ld.so.cache";
-
-lazy_static! {
-    static ref LD_SO_CACHE: Result<LdSoCache, Arc<io::Error>> =
-        LdSoCache::load(LD_SO_CACHE_FILE).map_err(Arc::new);
-}
-const LD_SO_CACHE_HEADER: &str = "glibc-ld.so.cache1.1";
-
 /// An user space probe.
 ///
 /// User probes are eBPF programs that can be attached to any userspace
@@ -83,10 +71,14 @@ impl UProbe {
         let target_str = &*target.as_os_str().to_string_lossy();
 
         let mut path = if let Some(pid) = pid {
-            find_lib_in_proc_maps(pid, target_str).map_err(|io_error| UProbeError::FileError {
-                filename: format!("/proc/{}/maps", pid),
-                io_error,
-            })?
+            let proc_map_libs =
+                ProcMap::new(pid).map_err(|e| UProbeError::ProcMapError { pid, source: e })?;
+            proc_map_libs
+                .find_by_name(target_str)
+                .map_err(|io_error| UProbeError::FileError {
+                    filename: format!("/proc/{}/maps", pid),
+                    io_error,
+                })?
         } else {
             None
         };
@@ -143,8 +135,8 @@ define_link_wrapper!(
     UProbeLink,
     /// The type returned by [UProbe::attach]. Can be passed to [UProbe::detach].
     UProbeLinkId,
-    PerfLink,
-    PerfLinkId
+    PerfLinkInner,
+    PerfLinkIdInner
 );
 
 /// The type returned when attaching an [`UProbe`] fails.
@@ -184,150 +176,14 @@ pub enum UProbeError {
         #[source]
         io_error: io::Error,
     },
-}
-
-fn proc_maps_libs(pid: pid_t) -> Result<Vec<(String, String)>, io::Error> {
-    let maps_file = format!("/proc/{}/maps", pid);
-    let data = fs::read_to_string(maps_file)?;
-
-    Ok(data
-        .lines()
-        .filter_map(|line| {
-            let line = line.split_whitespace().last()?;
-            if line.starts_with('/') {
-                let path = PathBuf::from(line);
-                let key = path.file_name().unwrap().to_string_lossy().into_owned();
-                Some((key, path.to_string_lossy().to_string()))
-            } else {
-                None
-            }
-        })
-        .collect())
-}
-
-fn find_lib_in_proc_maps(pid: pid_t, lib: &str) -> Result<Option<String>, io::Error> {
-    let libs = proc_maps_libs(pid)?;
-
-    let ret = if lib.contains(".so") {
-        libs.iter().find(|(k, _)| k.as_str().starts_with(lib))
-    } else {
-        let lib = lib.to_string();
-        let lib1 = lib.clone() + ".so";
-        let lib2 = lib + "-";
-        libs.iter()
-            .find(|(k, _)| k.starts_with(&lib1) || k.starts_with(&lib2))
-    };
-
-    Ok(ret.map(|(_, v)| v.clone()))
-}
-
-#[derive(Debug)]
-pub(crate) struct CacheEntry {
-    key: String,
-    value: String,
-    _flags: i32,
-}
-
-#[derive(Debug)]
-pub(crate) struct LdSoCache {
-    entries: Vec<CacheEntry>,
-}
-
-impl LdSoCache {
-    pub fn load<T: AsRef<Path>>(path: T) -> Result<Self, io::Error> {
-        let data = fs::read(path)?;
-        Self::parse(&data)
-    }
-
-    fn parse(data: &[u8]) -> Result<Self, io::Error> {
-        let mut cursor = Cursor::new(data);
-
-        let read_u32 = |cursor: &mut Cursor<_>| -> Result<u32, io::Error> {
-            let mut buf = [0u8; mem::size_of::<u32>()];
-            cursor.read_exact(&mut buf)?;
 
-            Ok(u32::from_ne_bytes(buf))
-        };
-
-        let read_i32 = |cursor: &mut Cursor<_>| -> Result<i32, io::Error> {
-            let mut buf = [0u8; mem::size_of::<i32>()];
-            cursor.read_exact(&mut buf)?;
-
-            Ok(i32::from_ne_bytes(buf))
-        };
-
-        let mut buf = [0u8; LD_SO_CACHE_HEADER.len()];
-        cursor.read_exact(&mut buf)?;
-        let header = std::str::from_utf8(&buf).map_err(|_| {
-            io::Error::new(io::ErrorKind::InvalidData, "invalid ld.so.cache header")
-        })?;
-        if header != LD_SO_CACHE_HEADER {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidData,
-                "invalid ld.so.cache header",
-            ));
-        }
-
-        let num_entries = read_u32(&mut cursor)?;
-        let _str_tab_len = read_u32(&mut cursor)?;
-        cursor.consume(5 * mem::size_of::<u32>());
-
-        let mut entries = Vec::new();
-        for _ in 0..num_entries {
-            let flags = read_i32(&mut cursor)?;
-            let k_pos = read_u32(&mut cursor)? as usize;
-            let v_pos = read_u32(&mut cursor)? as usize;
-            cursor.consume(12);
-            let key =
-                unsafe { CStr::from_ptr(cursor.get_ref()[k_pos..].as_ptr() as *const c_char) }
-                    .to_string_lossy()
-                    .into_owned();
-            let value =
-                unsafe { CStr::from_ptr(cursor.get_ref()[v_pos..].as_ptr() as *const c_char) }
-                    .to_string_lossy()
-                    .into_owned();
-            entries.push(CacheEntry {
-                key,
-                value,
-                _flags: flags,
-            });
-        }
-
-        Ok(LdSoCache { entries })
-    }
-
-    pub fn resolve(&self, lib: &str) -> Option<&str> {
-        let lib = if !lib.contains(".so") {
-            lib.to_string() + ".so"
-        } else {
-            lib.to_string()
-        };
-        self.entries
-            .iter()
-            .find(|entry| entry.key.starts_with(&lib))
-            .map(|entry| entry.value.as_str())
-    }
-}
-
-#[derive(Error, Debug)]
-enum ResolveSymbolError {
-    #[error(transparent)]
-    Io(#[from] io::Error),
-
-    #[error("error parsing ELF")]
-    Object(#[from] object::Error),
-
-    #[error("unknown symbol `{0}`")]
-    Unknown(String),
-}
-
-fn resolve_symbol(path: &str, symbol: &str) -> Result<u64, ResolveSymbolError> {
-    let data = fs::read(path)?;
-    let obj = object::read::File::parse(&*data)?;
-
-    obj.dynamic_symbols()
-        .chain(obj.symbols())
-        .find(|sym| sym.name().map(|name| name == symbol).unwrap_or(false))
-        .map(|s| s.address())
-        .ok_or_else(|| ResolveSymbolError::Unknown(symbol.to_string()))
+    /// There was en error resolving a path
+    #[error("error fetching libs for {pid}")]
+    ProcMapError {
+        /// The pid
+        pid: i32,
+        /// The [`ProcMapError`] that caused the error
+        #[source]
+        source: ProcMapError,
+    },
 }

aya/src/programs/usdt.rs

@@ -0,0 +1,525 @@
+//! User statically-defined tracepoints.
+use aya_common::{UsdtSpec, USDT_MAX_SPEC_COUNT};
+use libc::pid_t;
+use object::{elf::*, read::elf::*, Endianness};
+use std::{
+    collections::{HashMap, VecDeque},
+    convert::TryInto,
+    ffi::CStr,
+    fs,
+    io::{self, BufRead, Cursor, Read},
+    mem,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+use thiserror::Error;
+
+use crate::{
+    generated::{bpf_prog_type::BPF_PROG_TYPE_KPROBE, BPF_NOEXIST},
+    maps::{MapError, MapRefMut},
+    programs::{
+        define_link_wrapper, load_program,
+        perf_attach::{perf_attach, PerfLinkIdInner, PerfLinkInner},
+        probe::create_as_probe,
+        utils::{LD_SO_CACHE, LD_SO_CACHE_FILE},
+        Link, ProbeKind, ProgramData, ProgramError,
+    },
+    Pod, FEATURES,
+};
+
+use crate::programs::utils::{ProcMap, ProcMapError};
+
+unsafe impl Pod for UsdtSpec {}
+
+/// Name of the map used for USDT specs.
+pub const USDT_SPEC_MAP: &str = "__bpf_usdt_specs";
+/// Name of the map used for USDT to IP mappings.
+pub const USDT_IP_TO_SPEC_MAP: &str = "__bpf_usdt_ip_to_spec_id";
+
+/// A user statically-defined tracepoint program.
+///
+/// USDT programs are the fastest of the userspace tracing programs and can be
+/// used to trace events instrumented libraries or binaries. Unlike uprobes and
+/// uretprobes that have access to all CPU registers, USDTs provide a structured
+/// specification for accessing the arguments for each tracepoint. When compliled
+/// a single tracepoint may have mutliple different entrypoints in the same program.
+/// In order to simply access to arguments from eBPF, Aya keeps state in 2 maps:
+///
+///   - [`USDT_SPEC_MAP`] which keeps track of USDT specifications
+///   - [`USDT_IP_TO_SPEC_MAP`] which keeps track of Instructio Pointers to USDT specs.
+///
+/// The last map is not used on kernels which support the BPF Attach Cookie feature.
+///
+/// # Minimum kernel version
+///
+/// While support was added to the kenel in 4.19, Aya depends on a feature that
+/// allows the kernel to manage semaphore reference counting which was added in
+/// 4.20.
+///
+/// The minimum supported kernel version is 4.20.
+///
+/// # Examples
+///
+/// ```no_run
+/// # let mut bpf = Bpf::load_file("ebpf_programs.o")?;
+/// use aya::{Bpf, programs::{Usdt, usdt::{USDT_SPEC_MAP, USDT_IP_TO_SPEC_MAP}}};
+/// use aya::maps::{Array, HashMap};
+/// use std::convert::TryInto;
+///
+/// let spec_map = Array::try_from(bpf.map_mut(USDT_SPEC_MAP).unwrap())?;
+/// let ip_to_spec_map = HashMap::try_from(bpf.map_mut(USDT_IP_TO_SPEC_MAP).unwrap())?;
+/// let program: &mut Usdt = bpf.program_mut("usdt").unwrap().try_into()?;
+/// program.load()?;
+/// program.attach(
+///     spec_map,
+///     ip_to_spec_map,
+///     "clock",
+///     "loop",
+///     "/path/to/target/debug/clock",
+///     Some(12345),
+/// )?;
+/// # Ok::<(), aya::BpfError>(())
+/// ```
+#[derive(Debug)]
+#[doc(alias = "BPF_PROG_TYPE_KPROBE")]
+pub struct Usdt {
+    pub(crate) data: ProgramData<UsdtLink>,
+}
+
+impl Usdt {
+    /// Loads the program inside the kernel.
+    pub fn load(&mut self) -> Result<(), ProgramError> {
+        load_program(BPF_PROG_TYPE_KPROBE, &mut self.data)
+    }
+
+    /// Attaches the program.
+    ///
+    /// Attaches the USDT to the tracepoint with a matching `tp_provider` and `tp_name`
+    /// in the `target`.
+    ///
+    /// If `pid` is not `None`, the program executes only when the target
+    /// function is executed by the given `pid`. This is only supported in kernels which
+    /// provide the BPF Attach Cookie feature.
+    ///
+    /// The `target` argument can be an absolute path to a binary or library, or
+    /// a library name (eg: `"libc"`).
+    ///
+    /// Since there a single tracepoint can have multiple entrypoints, a single `UsdtLinkId`
+    /// may be comprised of multiple links.
+    ///
+    /// The returned value can be used to detach, see [Usdt::detach].
+    pub fn attach<T: AsRef<Path>>(
+        &mut self,
+        mut spec_map: crate::maps::Array<MapRefMut, UsdtSpec>,
+        mut ip_to_spec_map: crate::maps::HashMap<MapRefMut, i64, u32>,
+        tp_provider: &str,
+        tp_name: &str,
+        target: T,
+        pid: Option<pid_t>,
+    ) -> Result<UsdtLinkId, ProgramError> {
+        let target = target.as_ref();
+        let target_str = &*target.as_os_str().to_string_lossy();
+
+        let mut path = if let Some(pid) = pid {
+            let proc_map_libs =
+                ProcMap::new(pid).map_err(|e| UsdtError::ProcMapError { pid, source: e })?;
+            proc_map_libs
+                .find_by_name(target_str)
+                .map_err(|io_error| UsdtError::FileError {
+                    filename: format!("/proc/{}/maps", pid),
+                    io_error,
+                })?
+        } else {
+            None
+        };
+
+        if path.is_none() {
+            path = if target.is_absolute() {
+                Some(target_str)
+            } else {
+                let cache = LD_SO_CACHE
+                    .as_ref()
+                    .map_err(|error| UsdtError::InvalidLdSoCache {
+                        io_error: error.clone(),
+                    })?;
+                cache.resolve(target_str)
+            }
+            .map(String::from)
+        };
+
+        let path = path.ok_or(UsdtError::InvalidTarget {
+            path: target.to_owned(),
+        })?;
+
+        let tracepoints = collect_usdts(&path, tp_provider, tp_name, pid)?;
+        let mut perf_links = vec![];
+        let mut spec_ids = VecDeque::with_capacity(USDT_MAX_SPEC_COUNT as usize);
+        for i in 0..USDT_MAX_SPEC_COUNT {
+            spec_ids.push_back(i)
+        }
+        let mut spec_id_map = HashMap::new();
+        for t in tracepoints {
+            let id = if spec_id_map.contains_key(&t.args) {
+                *(spec_id_map.get(&t.args).unwrap())
+            } else {
+                let id = spec_ids.pop_front().unwrap();
+                spec_id_map.insert(t.args.clone(), id);
+                spec_map.set(id, t.spec, 0)?;
+                id
+            };
+            let mut cookie = Some(id as u64);
+            if !FEATURES.bpf_cookie {
+                cookie.take();
+                if let Err(MapError::SyscallError { call, io_error }) =
+                    ip_to_spec_map.insert(t.abs_ip.try_into().unwrap(), id, BPF_NOEXIST.into())
+                {
+                    if io_error.raw_os_error().unwrap() != (-libc::EEXIST) {
+                        return Err(ProgramError::MapError(MapError::SyscallError {
+                            call,
+                            io_error,
+                        }));
+                    }
+                }
+            }
+            let fd = create_as_probe(ProbeKind::UProbe, &path, t.rel_ip, pid, Some(t.sem_off))?;
+            let link = perf_attach(self.data.fd_or_err()?, fd, cookie)?;
+            perf_links.push(link);
+        }
+        let link = UsdtLink(MultiPerfLink { perf_links });
+        self.data.links.insert(link)
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [UProbe::attach].
+    pub fn detach(&mut self, link_id: UsdtLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn take_link(&mut self, link_id: UsdtLinkId) -> Result<UsdtLink, ProgramError> {
+        self.data.take_link(link_id)
+    }
+}
+
+/// The identifer of a MultiPerfLink.
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub(crate) struct MultiPerfLinkId(Vec<PerfLinkIdInner>);
+
+// A wrapper around multiple PerfLinkInner
+#[derive(Debug)]
+pub(crate) struct MultiPerfLink {
+    perf_links: Vec<PerfLinkInner>,
+}
+
+impl Link for MultiPerfLink {
+    type Id = MultiPerfLinkId;
+
+    fn id(&self) -> Self::Id {
+        let ids = self.perf_links.iter().map(|p| p.id()).collect();
+        MultiPerfLinkId(ids)
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        for l in self.perf_links {
+            l.detach()?;
+        }
+        Ok(())
+    }
+}
+
+define_link_wrapper!(
+    /// The link used by [Usdt] programs.
+    UsdtLink,
+    /// The type returned by [Usdt::attach]. Can be passed to [Usdt::detach].
+    UsdtLinkId,
+    MultiPerfLink,
+    MultiPerfLinkId
+);
+
+/// The type returned when attaching a [`Usdt`] fails.
+#[derive(Debug, Error)]
+pub enum UsdtError {
+    /// There was an error parsing `/etc/ld.so.cache`.
+    #[error("error reading `{}` file", LD_SO_CACHE_FILE)]
+    InvalidLdSoCache {
+        /// the original [`io::Error`]
+        #[source]
+        io_error: Arc<io::Error>,
+    },
+
+    /// The target program could not be found.
+    #[error("could not resolve uprobe target `{path}`")]
+    InvalidTarget {
+        /// path to target
+        path: PathBuf,
+    },
+
+    /// There was an error resolving the target symbol.
+    #[error("error resolving symbol")]
+    SymbolError {
+        /// symbol name
+        symbol: String,
+        /// the original error
+        #[source]
+        error: Box<dyn std::error::Error + Send + Sync>,
+    },
+
+    /// There was an error accessing `filename`.
+    #[error("`{filename}`")]
+    FileError {
+        /// The file name
+        filename: String,
+        /// The [`io::Error`] returned from the file operation
+        #[source]
+        io_error: io::Error,
+    },
+
+    /// There was en error resolving a path
+    #[error("error fetching libs for {pid}")]
+    ProcMapError {
+        /// The pid
+        pid: i32,
+        /// The [`ProcMapError`] that caused the error
+        #[source]
+        source: ProcMapError,
+    },
+
+    /// Unsupported file type
+    #[error("unsupported file type")]
+    Unsupported,
+
+    /// An [`io::Error`]
+    #[error("io error")]
+    Io(#[from] io::Error),
+
+    /// An [`object::Error`]
+    #[error("error parsing ELF")]
+    Object(#[from] object::Error),
+
+    /// Can't find matching offset in shard libs
+    #[error("can't find matching offset in shared libs")]
+    OffsetError,
+
+    /// Section is not executable
+    #[error("section is not executable")]
+    NoExec,
+
+    /// Segment is not found
+    #[error("segment not found")]
+    SegmentNotFound,
+
+    /// BPF Cookies are not supported
+    #[error("bpf cookies are required to support attachment without a pid")]
+    NoCookie,
+}
+
+fn collect_usdts(
+    path: &str,
+    provider: &str,
+    name: &str,
+    pid: Option<pid_t>,
+) -> Result<Vec<UsdtTarget>, UsdtError> {
+    let file = fs::read(path)?;
+    let data = &*file;
+    if let Ok(elf) = object::elf::FileHeader32::parse(data) {
+        if mem::size_of::<usize>() != 4 {
+            return Err(UsdtError::Unsupported);
+        }
+        return collect_usdts_from_elf(elf, data, provider, name, pid);
+    } else if let Ok(elf) = object::elf::FileHeader64::parse(data) {
+        if mem::size_of::<usize>() != 8 {
+            return Err(UsdtError::Unsupported);
+        }
+        return collect_usdts_from_elf(elf, data, provider, name, pid);
+    }
+    Err(UsdtError::Unsupported)
+}
+
+fn collect_usdts_from_elf<Elf: FileHeader<Endian = Endianness>>(
+    elf: &Elf,
+    data: &[u8],
+    provider: &str,
+    name: &str,
+    pid: Option<pid_t>,
+) -> Result<Vec<UsdtTarget>, UsdtError> {
+    let endian = elf.endian()?;
+    let sections = elf.sections(endian, data)?;
+    let program_headers = elf.program_headers(endian, data)?;
+    let mut results = vec![];
+    let mut base_addr: Option<u64> = None;
+    if let Some((_, base_section)) = sections.section_by_name(endian, b".stapsdt.base") {
+        base_addr = Some(base_section.sh_addr(endian).into())
+    };
+    if let Some((_, notes_section)) = sections.section_by_name(endian, b".note.stapsdt") {
+        if let Some(mut notes) = notes_section.notes(endian, data)? {
+            while let Ok(Some(note)) = notes.next() {
+                if note.name() != b"stapsdt" {
+                    continue;
+                }
+                if note.n_type(endian) != 3 {
+                    continue;
+                }
+                let note_data = note.desc();
+                let n = UsdtNote::parse(endian, note_data)?;
+                if n.provider != provider || n.name != name {
+                    continue;
+                }
+
+                let mut abs_ip = n.loc_addr;
+                if let Some(addr) = base_addr {
+                    abs_ip += addr - n.base_addr;
+                }
+
+                let seg = find_segment_by_address::<Elf>(program_headers, endian, abs_ip)
+                    .ok_or(UsdtError::SegmentNotFound)?;
+                if seg.p_flags(endian) & PF_X == 0 {
+                    return Err(UsdtError::NoExec);
+                }
+                let rel_ip = abs_ip - seg.p_vaddr(endian).into() + seg.p_offset(endian).into();
+
+                // If attaching to a sharef library and bpf cookies are not supported.
+                // Abs address of attach points are required
+                if elf.e_type(endian) == ET_DYN && !FEATURES.bpf_cookie {
+                    if pid.is_none() {
+                        return Err(UsdtError::NoCookie);
+                    }
+                    let proc_map_libs =
+                        ProcMap::new(pid.unwrap()).map_err(|e| UsdtError::ProcMapError {
+                            pid: pid.unwrap(),
+                            source: e,
+                        })?;
+                    let res = proc_map_libs
+                        .find_by_offset(rel_ip)
+                        .ok_or(UsdtError::OffsetError)?;
+                    abs_ip = res.address - res.offset + rel_ip;
+                }
+
+                let mut sem_off = 0;
+                if n.sem_addr != 0x0 {
+                    // semaphore refcnt support was in 4.20, which is min supported version so we assume its supported
+                    let seg = find_segment_by_address::<Elf>(program_headers, endian, n.sem_addr)
+                        .ok_or(UsdtError::SegmentNotFound)?;
+                    if seg.p_flags(endian) & PF_X == 0 {
+                        return Err(UsdtError::NoExec);
+                    }
+                    sem_off = n.sem_addr - seg.p_vaddr(endian).into() + seg.p_offset(endian).into();
+                }
+                let spec = n.args.parse().unwrap();
+                results.push(UsdtTarget {
+                    abs_ip,
+                    rel_ip,
+                    sem_off,
+                    args: n.args,
+                    spec,
+                })
+            }
+        }
+    }
+    Ok(results)
+}
+
+fn find_segment_by_address<Elf: FileHeader<Endian = Endianness>>(
+    program_headers: &[Elf::ProgramHeader],
+    endian: Endianness,
+    addr: u64,
+) -> Option<&Elf::ProgramHeader> {
+    program_headers.iter().find(|&header| {
+        header.p_vaddr(endian).into() < addr
+            && addr < (header.p_vaddr(endian).into() + header.p_memsz(endian).into())
+    })
+}
+
+// A resolved Usdt target.
+#[derive(Debug)]
+pub(crate) struct UsdtTarget {
+    abs_ip: u64,
+    rel_ip: u64,
+    sem_off: u64,
+    args: String,
+    spec: UsdtSpec,
+}
+
+// A parsed note from an ELF stapsdt note.
+#[derive(Debug)]
+pub(crate) struct UsdtNote {
+    loc_addr: u64,
+    base_addr: u64,
+    sem_addr: u64,
+    provider: String,
+    name: String,
+    args: String,
+}
+
+impl UsdtNote {
+    pub(crate) fn parse(endianness: Endianness, data: &[u8]) -> Result<UsdtNote, UsdtError> {
+        let mut cursor = Cursor::new(data);
+        let read_u64 = |cursor: &mut Cursor<_>| -> Result<u64, io::Error> {
+            let mut buf = [0u8; mem::size_of::<u64>()];
+            cursor.read_exact(&mut buf)?;
+            match endianness {
+                Endianness::Big => Ok(u64::from_be_bytes(buf)),
+                Endianness::Little => Ok(u64::from_le_bytes(buf)),
+            }
+        };
+        let read_string = |cursor: &mut Cursor<_>| -> Result<String, io::Error> {
+            let mut buf = vec![];
+            cursor.read_until(b'\0', &mut buf)?;
+            Ok(CStr::from_bytes_with_nul(&buf)
+                .unwrap()
+                .to_string_lossy()
+                .to_string())
+        };
+        let loc_addr = read_u64(&mut cursor)?;
+        let base_addr = read_u64(&mut cursor)?;
+        let sem_addr = read_u64(&mut cursor)?;
+        let provider = read_string(&mut cursor)?;
+        let name = read_string(&mut cursor)?;
+        let args = read_string(&mut cursor)?;
+
+        let res = UsdtNote {
+            loc_addr,
+            base_addr,
+            sem_addr,
+            provider,
+            name,
+            args,
+        };
+        Ok(res)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_parse_stapsdt() {
+        /*
+        /usr/bin/mariadb:     file format elf64-x86-64
+
+        Contents of section .note.stapsdt:
+         0000 08000000 34000000 03000000 73746170  ....4.......stap
+         0010 73647400 34a10d00 00000000 382e3600  sdt.4.......8.6.
+         0020 00000000 00000000 00000000 6c696267  ............libg
+         0030 63630075 6e77696e 64003840 25726469  cc.unwind.8@%rdi
+         0040 20384025 72736900                     8@%rsi
+        */
+        let data: &[u8] = &[
+            0x34, 0xa1, 0x0d, 0x00, 0x00, 0x00, 0x00, 0x00, 0x38, 0x2e, 0x36, 0x00, 0x00, 0x00,
+            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x6c, 0x69, 0x62, 0x67,
+            0x63, 0x63, 0x00, 0x75, 0x6e, 0x77, 0x69, 0x6e, 0x64, 0x00, 0x38, 0x40, 0x25, 0x72,
+            0x64, 0x69, 0x20, 0x38, 0x40, 0x25, 0x72, 0x73, 0x69, 0x00,
+        ];
+        let n = UsdtNote::parse(Endianness::Little, data).unwrap();
+        assert_eq!(n.loc_addr, 0xda134);
+        assert_eq!(n.base_addr, 0x362e38);
+        assert_eq!(n.sem_addr, 0x0);
+        assert_eq!(n.provider, "libgcc");
+        assert_eq!(n.name, "unwind");
+        assert_eq!(n.args, "8@%rdi 8@%rsi");
+    }
+}

aya/src/programs/utils.rs

@@ -1,5 +1,17 @@
 //! Common functions shared between multiple eBPF program types.
-use std::{ffi::CStr, os::unix::io::RawFd};
+use libc::pid_t;
+use object::{Object, ObjectSymbol};
+use std::{
+    collections::HashMap,
+    ffi::CStr,
+    fs,
+    io::{self, BufRead, Cursor, Read},
+    mem,
+    os::{raw::c_char, unix::prelude::RawFd},
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+use thiserror::Error;
 
 use crate::{
     programs::{FdLink, Link, ProgramData, ProgramError},
@@ -22,3 +34,272 @@ pub(crate) fn attach_raw_tracepoint<T: Link + From<FdLink>>(
 
     program_data.links.insert(FdLink::new(pfd).into())
 }
+
+pub(crate) const LD_SO_CACHE_FILE: &str = "/etc/ld.so.cache";
+
+lazy_static! {
+    pub(crate) static ref LD_SO_CACHE: Result<LdSoCache, Arc<io::Error>> =
+        LdSoCache::load(LD_SO_CACHE_FILE).map_err(Arc::new);
+}
+const LD_SO_CACHE_HEADER: &str = "glibc-ld.so.cache1.1";
+#[derive(Debug)]
+pub(crate) struct CacheEntry {
+    key: String,
+    value: String,
+    _flags: i32,
+}
+
+#[derive(Debug)]
+pub(crate) struct LdSoCache {
+    entries: Vec<CacheEntry>,
+}
+
+impl LdSoCache {
+    pub fn load<T: AsRef<Path>>(path: T) -> Result<Self, io::Error> {
+        let data = fs::read(path)?;
+        Self::parse(&data)
+    }
+
+    fn parse(data: &[u8]) -> Result<Self, io::Error> {
+        let mut cursor = Cursor::new(data);
+
+        let read_u32 = |cursor: &mut Cursor<_>| -> Result<u32, io::Error> {
+            let mut buf = [0u8; mem::size_of::<u32>()];
+            cursor.read_exact(&mut buf)?;
+
+            Ok(u32::from_ne_bytes(buf))
+        };
+
+        let read_i32 = |cursor: &mut Cursor<_>| -> Result<i32, io::Error> {
+            let mut buf = [0u8; mem::size_of::<i32>()];
+            cursor.read_exact(&mut buf)?;
+
+            Ok(i32::from_ne_bytes(buf))
+        };
+
+        let mut buf = [0u8; LD_SO_CACHE_HEADER.len()];
+        cursor.read_exact(&mut buf)?;
+        let header = std::str::from_utf8(&buf).map_err(|_| {
+            io::Error::new(io::ErrorKind::InvalidData, "invalid ld.so.cache header")
+        })?;
+        if header != LD_SO_CACHE_HEADER {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "invalid ld.so.cache header",
+            ));
+        }
+
+        let num_entries = read_u32(&mut cursor)?;
+        let _str_tab_len = read_u32(&mut cursor)?;
+        cursor.consume(5 * mem::size_of::<u32>());
+
+        let mut entries = Vec::new();
+        for _ in 0..num_entries {
+            let flags = read_i32(&mut cursor)?;
+            let k_pos = read_u32(&mut cursor)? as usize;
+            let v_pos = read_u32(&mut cursor)? as usize;
+            cursor.consume(12);
+            let key =
+                unsafe { CStr::from_ptr(cursor.get_ref()[k_pos..].as_ptr() as *const c_char) }
+                    .to_string_lossy()
+                    .into_owned();
+            let value =
+                unsafe { CStr::from_ptr(cursor.get_ref()[v_pos..].as_ptr() as *const c_char) }
+                    .to_string_lossy()
+                    .into_owned();
+            entries.push(CacheEntry {
+                key,
+                value,
+                _flags: flags,
+            });
+        }
+
+        Ok(LdSoCache { entries })
+    }
+
+    pub fn resolve(&self, lib: &str) -> Option<&str> {
+        let lib = if !lib.contains(".so") {
+            lib.to_string() + ".so"
+        } else {
+            lib.to_string()
+        };
+        self.entries
+            .iter()
+            .find(|entry| entry.key.starts_with(&lib))
+            .map(|entry| entry.value.as_str())
+    }
+}
+
+#[derive(Error, Debug)]
+pub(crate) enum ResolveSymbolError {
+    #[error(transparent)]
+    Io(#[from] io::Error),
+
+    #[error("error parsing ELF")]
+    Object(#[from] object::Error),
+
+    #[error("unknown symbol `{0}`")]
+    Unknown(String),
+}
+
+pub(crate) fn resolve_symbol(path: &str, symbol: &str) -> Result<u64, ResolveSymbolError> {
+    let data = fs::read(path)?;
+    let obj = object::read::File::parse(&*data)?;
+
+    obj.dynamic_symbols()
+        .chain(obj.symbols())
+        .find(|sym| sym.name().map(|name| name == symbol).unwrap_or(false))
+        .map(|s| s.address())
+        .ok_or_else(|| ResolveSymbolError::Unknown(symbol.to_string()))
+}
+
+/// Error reading from /proc/pid/maps
+#[derive(Debug, Error)]
+pub enum ProcMapError {
+    /// An [`io::Error`]
+    #[error(transparent)]
+    IoError(io::Error),
+
+    /// Error parsing a line of /proc/pid/maps
+    #[error("proc map entry parse error")]
+    ParseError,
+}
+
+pub(crate) struct ProcMap {
+    entries: Vec<ProcMapEntry>,
+    paths: HashMap<String, String>,
+}
+
+impl ProcMap {
+    pub(crate) fn new(pid: pid_t) -> Result<Self, ProcMapError> {
+        let maps_file = format!("/proc/{}/maps", pid);
+        let data = fs::read_to_string(maps_file).map_err(ProcMapError::IoError)?;
+        let mut entries = vec![];
+        let mut paths = HashMap::new();
+        for line in data.lines() {
+            let entry = ProcMapEntry::parse(line)?;
+            if let Some(path) = &entry.path {
+                let p = PathBuf::from(path);
+                let key = p.file_name().unwrap().to_string_lossy().into_owned();
+                let value = p.to_string_lossy().to_string();
+                paths.insert(key, value);
+            }
+            entries.push(entry);
+        }
+        Ok(ProcMap { entries, paths })
+    }
+
+    pub(crate) fn find_by_name(&self, lib: &str) -> Result<Option<String>, io::Error> {
+        let ret = if lib.contains(".so") {
+            self.paths.iter().find(|(k, _)| k.as_str().starts_with(lib))
+        } else {
+            let lib = lib.to_string();
+            let lib1 = lib.clone() + ".so";
+            let lib2 = lib + "-";
+            self.paths
+                .iter()
+                .find(|(k, _)| k.starts_with(&lib1) || k.starts_with(&lib2))
+        };
+
+        Ok(ret.map(|(_, v)| v.clone()))
+    }
+
+    pub(crate) fn find_by_offset(&self, offset: u64) -> Option<&ProcMapEntry> {
+        self.entries
+            .iter()
+            .find(|&e| e.offset <= offset && offset < e.offset + (e.address_end - e.address))
+    }
+}
+
+pub(crate) struct ProcMapEntry {
+    pub address: u64,
+    pub address_end: u64,
+    _perms: String,
+    pub offset: u64,
+    _dev: String,
+    _inode: u32,
+    pub path: Option<String>,
+}
+
+impl ProcMapEntry {
+    fn parse(line: &str) -> Result<Self, ProcMapError> {
+        let parts: Vec<&str> = line.split_whitespace().collect();
+        if parts.len() < 5 {
+            return Err(ProcMapError::ParseError);
+        }
+        let addr_parts: Vec<&str> = parts[0].split('-').collect();
+        let address =
+            u64::from_str_radix(addr_parts[0], 16).map_err(|_| ProcMapError::ParseError)?;
+        let address_end =
+            u64::from_str_radix(addr_parts[1], 16).map_err(|_| ProcMapError::ParseError)?;
+        let perms = parts[1];
+        let offset = u64::from_str_radix(parts[2], 16).map_err(|_| ProcMapError::ParseError)?;
+        let dev = parts[3];
+        let inode = parts[4].parse().map_err(|_| ProcMapError::ParseError)?;
+        let path = if parts.len() == 6 {
+            if parts[5].starts_with('/') {
+                Some(parts[5].to_string())
+            } else {
+                None
+            }
+        } else {
+            None
+        };
+
+        Ok(ProcMapEntry {
+            address,
+            address_end,
+            _perms: perms.to_string(),
+            offset,
+            _dev: dev.to_string(),
+            _inode: inode,
+            path,
+        })
+    }
+}
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_parse_proc_map_entry_from_str_1() {
+        let s = "7ffd6fbea000-7ffd6fbec000 r-xp 00000000 00:00 0                          [vdso]";
+        let proc_map = ProcMapEntry::parse(s).unwrap();
+        assert_eq!(proc_map.address, 0x7ffd6fbea000);
+        assert_eq!(proc_map.address_end, 0x7ffd6fbec000);
+        assert_eq!(proc_map._perms, "r-xp");
+        assert_eq!(proc_map.offset, 0x0);
+        assert_eq!(proc_map._dev, "00:00");
+        assert_eq!(proc_map._inode, 0);
+        assert_eq!(proc_map.path, None);
+    }
+
+    #[test]
+    fn test_parse_proc_map_entry_from_str_2() {
+        let s = "7f1bca83a000-7f1bca83c000 rw-p 00036000 fd:01 2895508                    /usr/lib64/ld-linux-x86-64.so.2";
+        let proc_map = ProcMapEntry::parse(s).unwrap();
+        assert_eq!(proc_map.address, 0x7f1bca83a000);
+        assert_eq!(proc_map.address_end, 0x7f1bca83c000);
+        assert_eq!(proc_map._perms, "rw-p");
+        assert_eq!(proc_map.offset, 0x00036000);
+        assert_eq!(proc_map._dev, "fd:01");
+        assert_eq!(proc_map._inode, 2895508);
+        assert_eq!(
+            proc_map.path,
+            Some("/usr/lib64/ld-linux-x86-64.so.2".to_string())
+        );
+    }
+
+    #[test]
+    fn test_parse_proc_map_entry_from_str_3() {
+        let s = "7f1bca5f9000-7f1bca601000 rw-p 00000000 00:00 0";
+        let proc_map = ProcMapEntry::parse(s).unwrap();
+        assert_eq!(proc_map.address, 0x7f1bca5f9000);
+        assert_eq!(proc_map.address_end, 0x7f1bca601000);
+        assert_eq!(proc_map._perms, "rw-p");
+        assert_eq!(proc_map.offset, 0x0);
+        assert_eq!(proc_map._dev, "00:00");
+        assert_eq!(proc_map._inode, 0);
+        assert_eq!(proc_map.path, None);
+    }
+}

aya/src/programs/xdp.rs

@@ -109,12 +109,11 @@ impl Xdp {
 
         let k_ver = kernel_version().unwrap();
         if k_ver >= (5, 9, 0) {
-            let link_fd = bpf_link_create(prog_fd, if_index, BPF_XDP, None, flags.bits).map_err(
-                |(_, io_error)| ProgramError::SyscallError {
+            let link_fd = bpf_link_create(prog_fd, if_index, BPF_XDP, None, None, flags.bits)
+                .map_err(|(_, io_error)| ProgramError::SyscallError {
                     call: "bpf_link_create".to_owned(),
                     io_error,
-                },
-            )? as RawFd;
+                })? as RawFd;
             self.data
                 .links
                 .insert(XdpLink(XdpLinkInner::FdLink(FdLink::new(link_fd))))

aya/src/sys/bpf.rs

@@ -334,6 +334,7 @@ pub(crate) fn bpf_link_create(
     target_fd: RawFd,
     attach_type: bpf_attach_type,
     btf_id: Option<u32>,
+    cookie: Option<u64>,
     flags: u32,
 ) -> SysResult {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
@@ -345,6 +346,9 @@ pub(crate) fn bpf_link_create(
     if let Some(btf_id) = btf_id {
         attr.link_create.__bindgen_anon_2.target_btf_id = btf_id;
     }
+    if let Some(cookie) = cookie {
+        attr.link_create.__bindgen_anon_2.perf_event.bpf_cookie = cookie;
+    }
 
     sys_bpf(bpf_cmd::BPF_LINK_CREATE, &attr)
 }
@@ -571,6 +575,69 @@ pub(crate) fn is_prog_name_supported() -> bool {
     }
 }
 
+pub(crate) fn is_perf_link_supported() -> bool {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+    let u = unsafe { &mut attr.__bindgen_anon_3 };
+
+    let prog: &[u8] = &[
+        0xb7, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // mov64 r0 = 0
+        0x95, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // exit
+    ];
+
+    let gpl = b"GPL\0";
+    u.license = gpl.as_ptr() as u64;
+
+    let insns = copy_instructions(prog).unwrap();
+    u.insn_cnt = insns.len() as u32;
+    u.insns = insns.as_ptr() as u64;
+    u.prog_type = bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT as u32;
+
+    if let Ok(fd) = sys_bpf(bpf_cmd::BPF_PROG_LOAD, &attr) {
+        if let Err((code, _)) = bpf_link_create(
+            fd as i32,
+            -1,
+            bpf_attach_type::BPF_PERF_EVENT,
+            None,
+            None,
+            0,
+        ) {
+            if code == (-libc::EBADF).into() {
+                unsafe { libc::close(fd as i32) };
+                return true;
+            }
+            unsafe { libc::close(fd as i32) };
+        }
+    }
+    false
+}
+
+pub(crate) fn is_bpf_cookie_supported() -> bool {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+    let u = unsafe { &mut attr.__bindgen_anon_3 };
+
+    let prog: &[u8] = &[
+        0x85, 0x00, 0x00, 0x00, 0xae, 0x00, 0x00, 0x00, // call bpf_get_attach_cookie
+        0x95, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // exit
+    ];
+
+    let gpl = b"GPL\0";
+    u.license = gpl.as_ptr() as u64;
+
+    let insns = copy_instructions(prog).unwrap();
+    u.insn_cnt = insns.len() as u32;
+    u.insns = insns.as_ptr() as u64;
+    u.prog_type = bpf_prog_type::BPF_PROG_TYPE_KPROBE as u32;
+
+    match sys_bpf(bpf_cmd::BPF_PROG_LOAD, &attr) {
+        Ok(v) => {
+            let fd = v as RawFd;
+            unsafe { close(fd) };
+            true
+        }
+        Err(_) => false,
+    }
+}
+
 pub(crate) fn is_btf_supported() -> bool {
     let mut btf = Btf::new();
     let name_offset = btf.add_string("int".to_string());

aya/src/sys/perf_event.rs

@@ -67,6 +67,7 @@ pub(crate) fn perf_event_open_probe(
     name: &str,
     offset: u64,
     pid: Option<pid_t>,
+    ref_cnt_offset: Option<u64>,
 ) -> SysResult {
     let mut attr = unsafe { mem::zeroed::<perf_event_attr>() };
 
@@ -74,6 +75,10 @@ pub(crate) fn perf_event_open_probe(
         attr.config = 1 << ret_bit;
     }
 
+    if let Some(ref_cnt_offset) = ref_cnt_offset {
+        attr.config |= ref_cnt_offset << 32;
+    }
+
     let c_name = CString::new(name).unwrap();
 
     attr.size = mem::size_of::<perf_event_attr>() as u32;

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"
+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,129 @@
+use core::ffi::c_void;
+
+use aya_common::{
+    UsdtArgType, UsdtSpec, USDT_MAX_ARG_COUNT, USDT_MAX_IP_COUNT, USDT_MAX_SPEC_COUNT,
+};
+
+use crate::{
+    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;
+
+#[cfg(not(feature = "cookie"))]
+use crate::args::FromPtRegs;
+
+#[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,
+}
+
+/// Errors from Usdt map operations
+#[derive(Debug, Clone)]
+pub enum UsdtError {
+    MaxArgCount,
+    SpecIdNotFound,
+    ValueError,
+    IpNotFound,
+}
+
+impl UsdtContext {
+    /// Creates a new Usdtcontext.
+    pub fn new(ctx: *mut c_void) -> UsdtContext {
+        UsdtContext {
+            regs: ctx as *mut pt_regs,
+        }
+    }
+
+    /// Access the register that holds the next instruction pointer.
+    #[cfg(not(feature = "cookie"))]
+    #[inline(always)]
+    fn ip<T: FromPtRegs>(&self) -> Option<T> {
+        T::from_ip(unsafe { &*self.regs })
+    }
+
+    /// Access the spec_id from the BPF Attach Cookie.
+    #[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) }
+    }
+
+    /// Access the spec_id using the `USDT_IP_TO_SPEC_ID` map
+    #[cfg(not(feature = "cookie"))]
+    #[inline(always)]
+    fn spec_id(&self) -> Result<u32, UsdtError> {
+        let ip: i64 = self.ip().ok_or(UsdtError::IpNotFound)?;
+        let spec = unsafe {
+            USDT_IP_TO_SPEC_ID
+                .get(&ip)
+                .ok_or(UsdtError::SpecIdNotFound)?
+        };
+        Ok(*spec)
+    }
+
+    /// Returns the value of the USDT argument `n` as a u64.
+    ///
+    /// This uses the USDT_SPEC_MAP to determine the correct specification to use in order
+    /// to read the value of argument `n` from the eBPF Context.
+    #[inline(always)]
+    pub fn arg(&self, n: usize) -> Result<u64, UsdtError> {
+        if n > USDT_MAX_ARG_COUNT {
+            return Err(UsdtError::MaxArgCount);
+        }
+        let spec_id = self.spec_id().map_err(|_| UsdtError::SpecIdNotFound)?;
+        let spec = USDT_SPECS.get(spec_id).ok_or(UsdtError::SpecIdNotFound)?;
+
+        if n > (spec.arg_count as usize) {
+            return Err(UsdtError::MaxArgCount);
+        }
+
+        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(|_| UsdtError::ValueError)?
+            },
+            UsdtArgType::RegDeref => unsafe {
+                let ptr: u64 = bpf_probe_read_kernel(
+                    self.as_ptr().offset(arg_spec.reg_off as isize) as *const _,
+                )
+                .map_err(|_| UsdtError::ValueError)?;
+                let ptr = ptr as *const u64;
+                bpf_probe_read_user::<u64>(ptr.offset(arg_spec.val_off as isize))
+                    .map_err(|_| UsdtError::ValueError)?
+                // 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
+    }
+}

images/Dockerfile.codegen

@@ -0,0 +1,25 @@
+FROM ubuntu:20.04
+
+RUN apt-get update && \
+    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+        build-essential \
+        clang \
+        ca-certificates \
+        curl \
+        git \
+        ssh \
+        libssl-dev \
+        pkg-config \
+        libc6-dev \
+        libc6-dev-arm64-cross \
+        libc6-dev-armel-cross \
+        libc6-dev-riscv64-cross \
+        && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+ENV RUSTUP_HOME=/rust
+ENV CARGO_HOME=/cargo
+ENV PATH=/cargo/bin:/rust/bin:$PATH
+RUN echo "(curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal -c rustfmt --default-toolchain nightly --no-modify-path) && rustup default nightly" > /install-rust.sh && chmod 755 /install-rust.sh
+RUN ./install-rust.sh

test/integration-ebpf/src/map_test.rs

@@ -26,6 +26,7 @@ unsafe fn try_pass(_ctx: XdpContext) -> Result<u32, u32> {
     Ok(xdp_action::XDP_PASS)
 }
 
+#[cfg(target_arch = "bpf")]
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }

test/integration-ebpf/src/name_test.rs

@@ -15,6 +15,7 @@ unsafe fn try_pass(_ctx: XdpContext) -> Result<u32, u32> {
     Ok(xdp_action::XDP_PASS)
 }
 
+#[cfg(target_arch = "bpf")]
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }

test/integration-ebpf/src/pass.rs

@@ -15,6 +15,7 @@ unsafe fn try_pass(_ctx: XdpContext) -> Result<u32, u32> {
     Ok(xdp_action::XDP_PASS)
 }
 
+#[cfg(target_arch = "bpf")]
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }

test/integration-ebpf/src/test.rs

@@ -15,6 +15,7 @@ unsafe fn try_pass(_ctx: XdpContext) -> Result<u32, u32> {
     Ok(xdp_action::XDP_PASS)
 }
 
+#[cfg(target_arch = "bpf")]
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }

xtask/src/codegen/aya_common.rs

@@ -0,0 +1,72 @@
+use anyhow::anyhow;
+use std::path::PathBuf;
+
+use aya_tool::{bindgen, write_to_file};
+
+use crate::codegen::{Architecture, Options};
+
+pub fn codegen(opts: &Options) -> Result<(), anyhow::Error> {
+    codegen_bindings(opts)
+}
+
+fn codegen_bindings(opts: &Options) -> Result<(), anyhow::Error> {
+    let types = [
+        // Registers
+        "pt_regs",
+        "user_pt_regs",
+    ];
+
+    let dir = PathBuf::from("aya-common");
+    let generated = dir.join("src/generated");
+
+    let builder = || {
+        bindgen::user_builder()
+            .header(dir.join("include/linux_wrapper.h").to_string_lossy())
+            .clang_args(&[
+                "-I",
+                &*opts.libbpf_dir.join("include/uapi").to_string_lossy(),
+            ])
+            .clang_args(&["-I", &*opts.libbpf_dir.join("include").to_string_lossy()])
+    };
+
+    for arch in Architecture::supported() {
+        let mut bindgen = builder();
+
+        // Set target triple. This will set the right flags (which you can see
+        // running clang -target=X  -E - -dM </dev/null)
+        let target = match arch {
+            Architecture::X86_64 => "x86_64-unknown-linux-gnu",
+            Architecture::ARMv7 => "armv7-unknown-linux-gnu",
+            Architecture::AArch64 => "aarch64-unknown-linux-gnu",
+            Architecture::RISCV64 => "riscv64-unknown-linux-gnu",
+        };
+        bindgen = bindgen.clang_args(&["-target", target]);
+
+        // Set the sysroot. This is needed to ensure that the correct arch
+        // specific headers are imported.
+        let sysroot = match arch {
+            Architecture::X86_64 => &opts.x86_64_sysroot,
+            Architecture::ARMv7 => &opts.armv7_sysroot,
+            Architecture::AArch64 => &opts.aarch64_sysroot,
+            Architecture::RISCV64 => &opts.riscv64_sysroot,
+        };
+        bindgen = bindgen.clang_args(&["-I", &*sysroot.to_string_lossy()]);
+
+        for x in &types {
+            bindgen = bindgen.allowlist_type(x);
+        }
+
+        let bindings = bindgen
+            .generate()
+            .map_err(|_| anyhow!("bindgen failed"))?
+            .to_string();
+
+        // write the bindings, with the original helpers removed
+        write_to_file(
+            &generated.join(format!("linux_bindings_{}.rs", arch)),
+            &bindings.to_string(),
+        )?;
+    }
+
+    Ok(())
+}

xtask/src/codegen/mod.rs

@@ -1,5 +1,6 @@
 mod aya;
 mod aya_bpf_bindings;
+mod aya_common;
 mod helpers;
 
 use std::path::PathBuf;
@@ -90,7 +91,8 @@ pub fn codegen(opts: Options) -> Result<(), anyhow::Error> {
         Some(AyaBpfBindings) => aya_bpf_bindings::codegen(&opts),
         None => {
             aya::codegen(&opts)?;
-            aya_bpf_bindings::codegen(&opts)
+            aya_bpf_bindings::codegen(&opts)?;
+            aya_common::codegen(&opts)
         }
     }
 }