nordisk/edda

crates/nornir-trace/src/ebpf.rs

//! **Live: eBPF uprobes on chosen functions of a running program** (`bpftrace`).
//!
//! The third source, after XRay (every call, rebuilt) and perf (sampled, no
//! rebuild): a few NAMED functions of a program that is already running, each
//! entry and return reported as it happens, so the map lights while the program
//! works. No rebuild, no restart.
//!
//! It is privileged, and this crate never asks for privileges: attaching a
//! uprobe needs `CAP_BPF` + `CAP_PERFMON` (kernel ≥ 5.8) or root on the box, and
//! `kernel.perf_event_paranoid` does not apply to it. The forge runs it only
//! where its operator allowed it (runa `LiveTrace`, opt-in); this module only
//! writes the probe script and reads what it prints.
//!
//! ```text
//! bpftrace -e '<script>'        # one uprobe + one uretprobe per function
//! E <tid> <nsecs> <probe>       # what each probe prints, one line per event
//! X <tid> <nsecs> <probe>
//! ```

use anyhow::{bail, Context, Result};
use object::{Object, ObjectSymbol};

use crate::xray::demangle;
use crate::EventKind;

/// The most probes one live trace attaches (each is a uprobe + a uretprobe).
pub const MAX_PROBES: usize = 64;

/// One function to probe: its linker symbol and its demangled name.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Probe {
    pub symbol: String,
    pub demangled: String,
}

/// Does a demangled name match a pattern (`a::b::c`, or `a::b::*` = every
/// function below `a::b`)?
pub fn matches(pattern: &str, demangled: &str) -> bool {
    match pattern.strip_suffix('*') {
        Some(prefix) => demangled.starts_with(prefix),
        None => demangled == pattern,
    }
}

/// **The functions of a program the patterns name** (function symbols of its
/// symbol table, demangled without hash), at most [`MAX_PROBES`]; `Err` when a
/// pattern names none or there are too many.
pub fn resolve_probes(elf: &[u8], patterns: &[String]) -> Result<Vec<Probe>> {
    let file = object::File::parse(elf).context("parse the program to probe")?;
    let mut out: Vec<Probe> = Vec::new();
    for p in patterns {
        let before = out.len();
        for s in file.symbols() {
            if s.kind() != object::SymbolKind::Text || !s.is_definition() {
                continue;
            }
            let Ok(name) = s.name() else { continue };
            let d = demangle(name);
            if matches(p, &d) && !out.iter().any(|x| x.symbol == name) {
                out.push(Probe { symbol: name.to_string(), demangled: d });
            }
        }
        if out.len() == before {
            bail!("`{p}` names no function of the program (inlined away, or not linked in)");
        }
    }
    if out.len() > MAX_PROBES {
        bail!("{} functions match: at most {MAX_PROBES} are probed (name them more narrowly)", out.len());
    }
    Ok(out)
}

/// **The bpftrace script**: for probe `i`, its entry prints `E tid nsecs i` and
/// its return `X tid nsecs i`. The program path and symbols are quoted; a
/// symbol is a linker name (`[A-Za-z0-9_$.]`), a path is refused if it holds a
/// quote or a newline.
pub fn script(program: &str, probes: &[Probe]) -> Result<String> {
    if program.contains(['"', '\n', '\\']) {
        bail!("a program path with a quote, backslash or newline is not probed: {program:?}");
    }
    let mut s = String::new();
    for (i, p) in probes.iter().enumerate() {
        if !p.symbol.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'$' || b == b'.') {
            bail!("`{}` is no linker symbol name", p.symbol);
        }
        s.push_str(&format!("uprobe:\"{program}\":\"{}\" {{ printf(\"E %d %llu {i}\\n\", tid, nsecs); }}\n", p.symbol));
        s.push_str(&format!("uretprobe:\"{program}\":\"{}\" {{ printf(\"X %d %llu {i}\\n\", tid, nsecs); }}\n", p.symbol));
    }
    Ok(s)
}

/// One line of what the script printed: `(thread, nsecs, probe, kind)`; any
/// other line (bpftrace's `Attaching N probes...`) is `None`.
pub fn parse_line(line: &str) -> Option<(u32, u64, u32, EventKind)> {
    let mut t = line.split_whitespace();
    let kind = match t.next()? {
        "E" => EventKind::Enter,
        "X" => EventKind::Exit,
        _ => return None,
    };
    let tid = t.next()?.parse().ok()?;
    let ns = t.next()?.parse().ok()?;
    let probe = t.next()?.parse().ok()?;
    t.next().is_none().then_some((tid, ns, probe, kind))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn patterns_match_paths_and_module_wildcards() {
        assert!(matches("dwarves::serve::*", "dwarves::serve::accept"));
        assert!(matches("dwarves::main", "dwarves::main"));
        assert!(!matches("dwarves::main", "dwarves::main_loop"));
        assert!(!matches("dwarves::serve::*", "dwarves::server::x"));
    }

    #[test]
    fn the_script_probes_each_function_twice_and_refuses_what_could_escape() {
        let probes = vec![
            Probe { symbol: "_ZN2xt3top17h0123456789abcdefE".into(), demangled: "xt::top".into() },
            Probe { symbol: "_RNvCs1_2xt6middle".into(), demangled: "xt::middle".into() },
        ];
        let s = script("/w/target/debug/xt", &probes).unwrap();
        assert_eq!(s.lines().count(), 4);
        assert!(s.contains(r#"uprobe:"/w/target/debug/xt":"_ZN2xt3top17h0123456789abcdefE" { printf("E %d %llu 0\n", tid, nsecs); }"#), "{s}");
        assert!(s.contains(r#"uretprobe:"/w/target/debug/xt":"_RNvCs1_2xt6middle" { printf("X %d %llu 1\n", tid, nsecs); }"#), "{s}");
        assert!(script("/w/x\" }; system(\"id", &probes).is_err());
        assert!(script("/w/x", &[Probe { symbol: "a\"b".into(), demangled: "a".into() }]).is_err());
    }

    #[test]
    fn printed_lines_read_as_events() {
        assert_eq!(parse_line("E 4242 123456789 3"), Some((4242, 123456789, 3, EventKind::Enter)));
        assert_eq!(parse_line("X 4242 123456999 3"), Some((4242, 123456999, 3, EventKind::Exit)));
        assert_eq!(parse_line("Attaching 4 probes..."), None);
        assert_eq!(parse_line("E 1 2 3 4"), None);
    }
}