aya-common: Add a new package

This commit adds aya-common, which is for code that is shared between
aya-bpf and aya. The initial use case is for USDT probes, which require
a lot of co-operation between aya and aya-bpf to provide a good UX.

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

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,297 @@
+#![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 {
+    Const,
+    Reg,
+    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 {
+    /// scalar interpreted depending on arg_type
+    pub val_off: u64,
+    /// arg location case
+    pub arg_type: UsdtArgType,
+    /// offset of referenced register within struct pt_regs
+    pub reg_off: i16,
+    /// whether arg should be interpreted as signed value
+    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,
+}
+
+#[repr(C)]
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "user", derive(Debug))]
+pub struct UsdtSpec {
+    pub args: [UsdtArgSpec; USDT_MAX_ARG_COUNT],
+    pub cookie: u64,
+    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
+                }
+            );
+        }
+    }
+}

xtask/src/codegen/aya_common.rs

@@ -0,0 +1,72 @@
+use anyhow::anyhow;
+use std::path::PathBuf;
+
+use aya_gen::{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)
         }
     }
 }