aya: add a safe wrapper for 'BPF_TASK_FD_QUERY'

This wraps the logic for `BPF_TASK_FD_QUERY` in a safe helper.

aya/src/programs/mod.rs

@@ -66,7 +66,7 @@ pub mod xdp;
 
 use libc::ENOSPC;
 use std::{
-    ffi::CString,
+    ffi::{CStr, CString},
     io,
     os::unix::io::{AsRawFd, RawFd},
     path::Path,
@@ -103,13 +103,14 @@ pub use uprobe::{UProbe, UProbeError};
 pub use xdp::{Xdp, XdpError, XdpFlags};
 
 use crate::{
-    generated::{bpf_attach_type, bpf_prog_info, bpf_prog_type},
+    generated::{bpf_attach_type, bpf_prog_info, bpf_prog_type, bpf_task_fd_type},
     maps::MapError,
     obj::{self, btf::BtfError, Function, KernelVersion},
     pin::PinError,
     sys::{
         bpf_get_object, bpf_load_program, bpf_pin_object, bpf_prog_get_fd_by_id,
-        bpf_prog_get_info_by_fd, bpf_prog_query, retry_with_verifier_logs, BpfLoadProgramAttrs,
+        bpf_prog_get_info_by_fd, bpf_prog_query, bpf_task_fd_query, retry_with_verifier_logs,
+        BpfLoadProgramAttrs,
     },
     util::VerifierLog,
 };
@@ -849,3 +850,111 @@ impl ProgramInfo {
         Ok(ProgramInfo(info))
     }
 }
+
+/// Kind of a probe program from a process FD.
+#[repr(u32)]
+#[non_exhaustive]
+#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
+pub enum TaskFdKind {
+    /// Raw tracepoint.
+    RawTracePoint = bpf_task_fd_type::BPF_FD_TYPE_RAW_TRACEPOINT as u32,
+    /// Tracepoint.
+    TracePoint = bpf_task_fd_type::BPF_FD_TYPE_TRACEPOINT as u32,
+    /// Kernel probe.
+    KProbe = bpf_task_fd_type::BPF_FD_TYPE_KPROBE as u32,
+    /// Kernel return probe.
+    KRetProbe = bpf_task_fd_type::BPF_FD_TYPE_KRETPROBE as u32,
+    /// User space probe.
+    UProbe = bpf_task_fd_type::BPF_FD_TYPE_UPROBE as u32,
+    /// User space return probe.
+    URetProbe = bpf_task_fd_type::BPF_FD_TYPE_URETPROBE as u32,
+}
+
+impl TryFrom<u32> for TaskFdKind {
+    type Error = ProgramError;
+
+    fn try_from(v: u32) -> Result<Self, Self::Error> {
+        use bpf_task_fd_type::*;
+        use TaskFdKind::*;
+        Ok(match v {
+            x if x == BPF_FD_TYPE_RAW_TRACEPOINT as u32 => RawTracePoint,
+            x if x == BPF_FD_TYPE_TRACEPOINT as u32 => TracePoint,
+            x if x == BPF_FD_TYPE_KPROBE as u32 => KProbe,
+            x if x == BPF_FD_TYPE_KRETPROBE as u32 => KRetProbe,
+            x if x == BPF_FD_TYPE_UPROBE as u32 => UProbe,
+            x if x == BPF_FD_TYPE_URETPROBE as u32 => URetProbe,
+            _ => return Err(ProgramError::UnexpectedProgramType),
+        })
+    }
+}
+
+/// Details of a probe program from a process FD.
+#[derive(Debug)]
+pub struct TaskFdDetails<'buf> {
+    /// Program ID.
+    id: u32,
+    /// Program kind.
+    kind: TaskFdKind,
+    /// Program name.
+    name: Option<&'buf CStr>,
+    /// Optional program probe address.
+    probe_addr: Option<u64>,
+    /// Optional program probe offset.
+    probe_offset: Option<u64>,
+}
+
+impl<'buf> TaskFdDetails<'buf> {
+    /// Query the details of an FD in a process (task), looking for probe programs.
+    ///
+    /// `Ok(None)` is returned if the target FD does not exist or is not a probe program.
+    pub fn query_task_fd(
+        pid: u32,
+        target_fd: RawFd,
+        out_name_buf: Option<&mut [u8]>,
+    ) -> Result<Option<TaskFdDetails<'buf>>, ProgramError> {
+        let out = match bpf_task_fd_query(pid, target_fd, out_name_buf) {
+            Ok(v) => v,
+            // ENOTSUPP (errno 95) means that the target FD is not a perf program.
+            Err(e) if e.raw_os_error() == Some(95) => return Ok(None),
+            // ENOENT (errno 2) means that the target FD or PID does not exist.
+            Err(e) if e.raw_os_error() == Some(2) => return Ok(None),
+            Err(e) => {
+                return Err(ProgramError::SyscallError {
+                    call: "bpf_task_fd_query".to_owned(),
+                    io_error: e,
+                })
+            }
+        };
+        let kind = TaskFdKind::try_from(out.fd_type as u32)?;
+        Ok(Some(TaskFdDetails {
+            id: out.prog_id,
+            kind,
+            name: out.name,
+            probe_addr: out.probe_addr,
+            probe_offset: out.probe_offset,
+        }))
+    }
+
+    /// Return the program ID.
+    pub fn id(&self) -> u32 {
+        self.id
+    }
+
+    /// Return the program kind.
+    pub fn kind(&self) -> TaskFdKind {
+        self.kind
+    }
+
+    /// Return the program name.
+    pub fn name(&self) -> Option<&CStr> {
+        self.name
+    }
+
+    /// Return the program probe address and offset.
+    pub fn address_and_offset(&self) -> Option<(u64, u64)> {
+        match (self.probe_addr, self.probe_offset) {
+            (Some(a), Some(o)) => Some((a, o)),
+            _ => None,
+        }
+    }
+}

aya/src/sys/bpf.rs

@@ -12,7 +12,7 @@ use libc::{c_char, c_long, close, ENOENT, ENOSPC};
 use crate::{
     generated::{
         bpf_attach_type, bpf_attr, bpf_btf_info, bpf_cmd, bpf_insn, bpf_link_info, bpf_map_info,
-        bpf_prog_info, bpf_prog_type, BPF_F_REPLACE,
+        bpf_prog_info, bpf_prog_type, bpf_task_fd_type, BPF_F_REPLACE,
     },
     maps::PerCpuValues,
     obj::{
@@ -451,6 +451,68 @@ pub(crate) fn bpf_prog_get_info_by_fd(prog_fd: RawFd) -> Result<bpf_prog_info, i
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct TaskFdQueryOutput<'buf> {
+    pub(crate) prog_id: u32,
+    pub(crate) fd_type: bpf_task_fd_type,
+    pub(crate) name: Option<&'buf CStr>,
+    pub(crate) probe_offset: Option<u64>,
+    pub(crate) probe_addr: Option<u64>,
+}
+
+pub(crate) fn bpf_task_fd_query<'buf>(
+    pid: u32,
+    target_fd: RawFd,
+    out_name_buf: Option<&mut [u8]>,
+) -> Result<TaskFdQueryOutput<'buf>, io::Error> {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.task_fd_query.pid = pid;
+    attr.task_fd_query.fd = target_fd as u32;
+    let mut out_name_buf = out_name_buf;
+    if let Some(buf) = &mut out_name_buf {
+        attr.task_fd_query.buf = buf.as_mut_ptr() as u64;
+        attr.task_fd_query.buf_len = buf.len() as u32;
+    };
+
+    if let Err((_, err)) = sys_bpf(bpf_cmd::BPF_TASK_FD_QUERY, &attr) {
+        // The kernel here may leak an internal ENOTSUPP code (524), so
+        // this needs to translate it back to POSIX-defined ENOTSUPP (95).
+        return match err.raw_os_error() {
+            Some(524) => Err(io::Error::from_raw_os_error(95)),
+            _ => Err(err),
+        };
+    }
+
+    let fd_type = unsafe { std::mem::transmute(attr.task_fd_query.fd_type) };
+    let name = out_name_buf.map(|buf| unsafe {
+        CStr::from_bytes_with_nul_unchecked(slice::from_raw_parts(
+            buf.as_ptr(),
+            attr.task_fd_query.buf_len as usize + 1,
+        ))
+    });
+    let (probe_offset, probe_addr) = match fd_type {
+        bpf_task_fd_type::BPF_FD_TYPE_KPROBE
+        | bpf_task_fd_type::BPF_FD_TYPE_KRETPROBE
+        | bpf_task_fd_type::BPF_FD_TYPE_UPROBE
+        | bpf_task_fd_type::BPF_FD_TYPE_URETPROBE => unsafe {
+            (
+                Some(attr.task_fd_query.probe_offset),
+                Some(attr.task_fd_query.probe_addr),
+            )
+        },
+        _ => (None, None),
+    };
+
+    Ok(TaskFdQueryOutput {
+        prog_id: unsafe { attr.task_fd_query.prog_id },
+        fd_type,
+        name,
+        probe_offset,
+        probe_addr,
+    })
+}
+
 pub(crate) fn bpf_map_get_info_by_fd(prog_fd: RawFd) -> Result<bpf_map_info, io::Error> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
     // info gets entirely populated by the kernel