aya rust bpf userprobe tracing

bpf_aya.rs

#![no_std]
#![no_main]
use aya_bpf::helpers::bpf_get_current_pid_tgid;
use aya_bpf::helpers::gen::bpf_ktime_get_ns;
use aya_bpf::helpers::bpf_probe_read_user_str_bytes;
use aya_bpf::helpers::bpf_probe_read_kernel_str_bytes;
use aya_bpf::BpfContext;
use aya_bpf::{
    macros::map,
    macros::{uprobe, uretprobe},
    maps::HashMap,
    programs::ProbeContext,
};
use aya_log_ebpf::info;
use core::ffi::{ c_char, CStr};

use tracing_first_ebpf::generated::addrinfo;

#[map(name = "LATENCY")]
static mut LATENCY: HashMap<u32, u64> = HashMap::with_max_entries(1024, 0);

#[uprobe]
pub fn tracing_first(ctx: ProbeContext) -> u32 {
    match try_tracing_first(ctx) {
        Ok(ret) => ret,
        Err(ret) => ret,
    }
}

fn try_tracing_first(ctx: ProbeContext) -> Result<u32, u32> {
    unsafe {
        let node: *const c_char = ctx.arg(0).unwrap();
        let mut buf = [0u8; 16];
        unsafe { bpf_probe_read_user_str_bytes(node as *const u8, &mut buf).map_err(|e| e as u32)? };
        let node = core::str::from_utf8_unchecked(&buf[..]);

        let service: *const c_char = ctx.arg(1).unwrap();
        let mut buf = [0u8; 16];
        unsafe { bpf_probe_read_user_str_bytes(service as *const u8, &mut buf).map_err(|e| e as u32)? };
        let service = core::str::from_utf8_unchecked(&buf[..]);


        let tid: u32 = bpf_get_current_pid_tgid() as u32;
        let now = bpf_ktime_get_ns();

        let Ok(()) = LATENCY.insert(&tid, &now, 0_u64) else {
            info!(&ctx, "failed to insert latency");
            return Err(0)
        };

        let command = ctx.command().unwrap_or_default();
        if command.len() > 16 {
            return Err(2);
        };
        let comm = core::str::from_utf8_unchecked(&command[..]);
        info!(&ctx, "function getaddrinfo called from command: {} node {} service: {} ", comm, node, service);
    };
    Ok(0)
}

#[uretprobe]
pub fn tracing_uretprobe(ctx: ProbeContext) -> u32 {
    match try_tracing_uretprobe(ctx) {
        Ok(ret) => ret,
        Err(ret) => ret,
    }
}

fn try_tracing_uretprobe(ctx: ProbeContext) -> Result<u32, u32> {
    unsafe {
        let tid: u32 = bpf_get_current_pid_tgid() as u32;
        let now = bpf_ktime_get_ns();

        let result: *const *const addrinfo = ctx.arg(3).ok_or(1u32)?;
        let mut buf = [0u8; 32];
        
        let canonname = unsafe {
            let canonname: *const c_char = (**result).ai_canonname;
            if !canonname.is_null() {
                core::str::from_utf8_unchecked(
                    bpf_probe_read_user_str_bytes(canonname as *const u8, &mut buf)
                        .map_err(|e| e as u32)?,
                )
            } else {
                ""
            }
        };

        if canonname.len() > 32 {
            return Err(2);
        }


        if let Some(start) = LATENCY.get(&tid) {
            let latency = now - start;
            let ret = LATENCY.remove(&tid);
            if ret.is_err() {
                info!(&ctx, "failed to delete latency");
                return Err(0)
            }
            info!(&ctx, "function getaddrinfo retProbe called by with latency: {}nanos and canonname: {}", latency, canonname);
        } else {
            info!(&ctx, "failed to get LATENCY element: {}", tid);
            return Err(1);
        };
    };

    Ok(0)
}

#[panic_handler]
fn panic(_info: &core::panic::PanicInfo) -> ! {
    unsafe { core::hint::unreachable_unchecked() }
}