aya: add a safe wrapper for 'BPF_TASK_FD_QUERY' logic

aya/src/programs/mod.rs

@@ -71,7 +71,7 @@ pub mod uprobe;
 pub mod xdp;
 
 use std::{
-    ffi::CString,
+    ffi::{CStr, CString},
     io,
     os::fd::{AsFd, BorrowedFd},
     path::{Path, PathBuf},
@@ -81,7 +81,7 @@ use std::{
 use aya_obj::{
     VerifierLog,
     btf::BtfError,
-    generated::{bpf_attach_type, bpf_link_info, bpf_prog_info, bpf_prog_type},
+    generated::{bpf_attach_type, bpf_link_info, bpf_prog_info, bpf_prog_type, bpf_task_fd_type},
 };
 use info::impl_info;
 pub use info::{ProgramInfo, ProgramType, loaded_programs};
@@ -127,7 +127,7 @@ use crate::{
     sys::{
         EbpfLoadProgramAttrs, NetlinkError, ProgQueryTarget, SyscallError, bpf_btf_get_fd_by_id,
         bpf_get_object, bpf_link_get_fd_by_id, bpf_link_get_info_by_fd, bpf_load_program,
-        bpf_pin_object, bpf_prog_get_fd_by_id, bpf_prog_query, iter_link_ids,
+        bpf_pin_object, bpf_prog_get_fd_by_id, bpf_prog_query, bpf_task_fd_query, iter_link_ids,
         retry_with_verifier_logs,
     },
     util::KernelVersion,
@@ -1089,3 +1089,114 @@ pub fn loaded_links() -> impl Iterator<Item = Result<bpf_link_info, ProgramError
         })
         .map(|result| result.map_err(Into::into))
 }
+
+/// 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 TaskFdKind::*;
+        use bpf_task_fd_type::*;
+        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<'a> {
+    /// Program ID.
+    id: u32,
+    /// Program kind.
+    kind: TaskFdKind,
+    /// Program name.
+    name: Option<&'a CStr>,
+    /// Optional program probe address.
+    probe_addr: Option<u64>,
+    /// Optional program probe offset.
+    probe_offset: Option<u64>,
+}
+
+impl TaskFdDetails<'_> {
+    /// 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: BorrowedFd<'_>,
+        out_name_buf: Option<&mut [u8]>,
+    ) -> Result<Option<Self>, ProgramError> {
+        let out = match bpf_task_fd_query(pid, target_fd, out_name_buf) {
+            Ok(v) => v,
+            Err(io_error) => {
+                return match io_error.raw_os_error() {
+                    // ENOTSUPP (errno 95) means that the target FD is not a perf program.
+                    Some(95) => Ok(None),
+                    // ENOENT (errno 2) means that the target FD or PID does not exist.
+                    Some(2) => Ok(None),
+                    _ => Err(SyscallError {
+                        call: "bpf_task_fd_query",
+                        io_error,
+                    }
+                    .into()),
+                };
+            }
+        };
+        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

@@ -18,7 +18,7 @@ use aya_obj::{
         BPF_ALU64, BPF_CALL, BPF_DW, BPF_EXIT, BPF_F_REPLACE, BPF_IMM, BPF_JMP, BPF_K, BPF_LD,
         BPF_MEM, BPF_MOV, BPF_PSEUDO_MAP_VALUE, BPF_ST, BPF_SUB, BPF_X, bpf_attach_type, bpf_attr,
         bpf_btf_info, bpf_cmd, bpf_func_id::*, bpf_insn, bpf_link_info, bpf_map_info, bpf_map_type,
-        bpf_prog_info, bpf_prog_type, bpf_stats_type,
+        bpf_prog_info, bpf_prog_type, bpf_stats_type, bpf_task_fd_type,
     },
     maps::{LegacyMap, bpf_map_def},
 };
@@ -601,6 +601,69 @@ pub(crate) fn bpf_prog_get_info_by_fd(
     })
 }
 
+#[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: BorrowedFd<'_>,
+    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_raw_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(io_error) = unit_sys_bpf(bpf_cmd::BPF_TASK_FD_QUERY, &mut 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 io_error.raw_os_error() {
+            Some(524) => Err(io::Error::from_raw_os_error(95)),
+            _ => Err(io_error),
+        };
+    }
+
+    let fd_type =
+        unsafe { std::mem::transmute::<u32, bpf_task_fd_type>(attr.task_fd_query.fd_type) };
+    let name = out_name_buf.map(|buf| unsafe {
+        CStr::from_bytes_with_nul_unchecked(std::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,
+    })
+}
+
 /// Introduced in kernel v4.13.
 pub(crate) fn bpf_map_get_fd_by_id(map_id: u32) -> Result<crate::MockableFd, SyscallError> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };