aya: Enable bpf_link for perf_attach programs

This adds support for bpf_link to PerfEvent, Tracepoint, Kprobe and
Uprobe programs.

Signed-off-by: Dave Tucker <[email protected]>

aya/src/programs/kprobe.rs

@@ -6,7 +6,7 @@ use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
     programs::{
         define_link_wrapper, load_program,
-        perf_attach::{PerfLink, PerfLinkId},
+        perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{attach, ProbeKind},
         ProgramData, ProgramError,
     },
@@ -101,8 +101,8 @@ define_link_wrapper!(
     KProbeLink,
     /// The type returned by [KProbe::attach]. Can be passed to [KProbe::detach].
     KProbeLinkId,
-    PerfLink,
-    PerfLinkId
+    PerfLinkInner,
+    PerfLinkIdInner
 );
 
 /// The type returned when attaching a [`KProbe`] fails.

aya/src/programs/perf_attach.rs

@@ -3,11 +3,42 @@ use libc::close;
 use std::os::unix::io::RawFd;
 
 use crate::{
-    programs::{probe::detach_debug_fs, Link, ProbeKind, ProgramData, ProgramError},
-    sys::perf_event_ioctl,
-    PERF_EVENT_IOC_DISABLE, PERF_EVENT_IOC_ENABLE, PERF_EVENT_IOC_SET_BPF,
+    generated::bpf_attach_type::BPF_PERF_EVENT,
+    programs::{probe::detach_debug_fs, FdLink, Link, ProbeKind, ProgramError},
+    sys::{bpf_link_create, perf_event_ioctl},
+    FEATURES, PERF_EVENT_IOC_DISABLE, PERF_EVENT_IOC_ENABLE, PERF_EVENT_IOC_SET_BPF,
 };
 
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub(crate) enum PerfLinkIdInner {
+    FdLinkId(<FdLink as Link>::Id),
+    PerfLinkId(<PerfLink as Link>::Id),
+}
+
+#[derive(Debug)]
+pub(crate) enum PerfLinkInner {
+    FdLink(FdLink),
+    PerfLink(PerfLink),
+}
+
+impl Link for PerfLinkInner {
+    type Id = PerfLinkIdInner;
+
+    fn id(&self) -> Self::Id {
+        match self {
+            PerfLinkInner::FdLink(link) => PerfLinkIdInner::FdLinkId(link.id()),
+            PerfLinkInner::PerfLink(link) => PerfLinkIdInner::PerfLinkId(link.id()),
+        }
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        match self {
+            PerfLinkInner::FdLink(link) => link.detach(),
+            PerfLinkInner::PerfLink(link) => link.detach(),
+        }
+    }
+}
+
 /// The identifer of a PerfLink.
 #[derive(Debug, Hash, Eq, PartialEq)]
 pub struct PerfLinkId(RawFd);
@@ -41,29 +72,36 @@ impl Link for PerfLink {
     }
 }
 
-pub(crate) fn perf_attach<T: Link + From<PerfLink>>(
-    data: &mut ProgramData<T>,
-    fd: RawFd,
-) -> Result<T::Id, ProgramError> {
-    perf_attach_either(data, fd, None, None)
+pub(crate) fn perf_attach(prog_fd: RawFd, fd: RawFd) -> Result<PerfLinkInner, ProgramError> {
+    if FEATURES.bpf_perf_link {
+        let link_fd =
+            bpf_link_create(prog_fd, fd, BPF_PERF_EVENT, None, 0).map_err(|(_, io_error)| {
+                ProgramError::SyscallError {
+                    call: "bpf_link_create".to_owned(),
+                    io_error,
+                }
+            })? as RawFd;
+        Ok(PerfLinkInner::FdLink(FdLink::new(link_fd)))
+    } else {
+        perf_attach_either(prog_fd, fd, None, None)
+    }
 }
 
-pub(crate) fn perf_attach_debugfs<T: Link + From<PerfLink>>(
-    data: &mut ProgramData<T>,
+pub(crate) fn perf_attach_debugfs(
+    prog_fd: RawFd,
     fd: RawFd,
     probe_kind: ProbeKind,
     event_alias: String,
-) -> Result<T::Id, ProgramError> {
-    perf_attach_either(data, fd, Some(probe_kind), Some(event_alias))
+) -> Result<PerfLinkInner, ProgramError> {
+    perf_attach_either(prog_fd, fd, Some(probe_kind), Some(event_alias))
 }
 
-fn perf_attach_either<T: Link + From<PerfLink>>(
-    data: &mut ProgramData<T>,
+fn perf_attach_either(
+    prog_fd: RawFd,
     fd: RawFd,
     probe_kind: Option<ProbeKind>,
     event_alias: Option<String>,
-) -> Result<T::Id, ProgramError> {
-    let prog_fd = data.fd_or_err()?;
+) -> Result<PerfLinkInner, ProgramError> {
     perf_event_ioctl(fd, PERF_EVENT_IOC_SET_BPF, prog_fd).map_err(|(_, io_error)| {
         ProgramError::SyscallError {
             call: "PERF_EVENT_IOC_SET_BPF".to_owned(),
@@ -77,12 +115,9 @@ fn perf_attach_either<T: Link + From<PerfLink>>(
         }
     })?;
 
-    data.links.insert(
-        PerfLink {
-            perf_fd: fd,
-            probe_kind,
-            event_alias,
-        }
-        .into(),
-    )
+    Ok(PerfLinkInner::PerfLink(PerfLink {
+        perf_fd: fd,
+        probe_kind,
+        event_alias,
+    }))
 }

aya/src/programs/perf_event.rs

@@ -13,8 +13,9 @@ use crate::{
         },
     },
     programs::{
+        links::define_link_wrapper,
         load_program, perf_attach,
-        perf_attach::{PerfLink, PerfLinkId},
+        perf_attach::{PerfLinkIdInner, PerfLinkInner},
         ProgramData, ProgramError,
     },
     sys::perf_event_open,
@@ -119,7 +120,7 @@ pub enum PerfEventScope {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_PERF_EVENT")]
 pub struct PerfEvent {
-    pub(crate) data: ProgramData<PerfLink>,
+    pub(crate) data: ProgramData<PerfEventLink>,
 }
 
 impl PerfEvent {
@@ -141,7 +142,7 @@ impl PerfEvent {
         config: u64,
         scope: PerfEventScope,
         sample_policy: SamplePolicy,
-    ) -> Result<PerfLinkId, ProgramError> {
+    ) -> Result<PerfEventLinkId, ProgramError> {
         let (sample_period, sample_frequency) = match sample_policy {
             SamplePolicy::Period(period) => (period, None),
             SamplePolicy::Frequency(frequency) => (0, Some(frequency)),
@@ -168,13 +169,14 @@ impl PerfEvent {
             io_error,
         })? as i32;
 
-        perf_attach(&mut self.data, fd)
+        let link = perf_attach(self.data.fd_or_err()?, fd)?;
+        self.data.links.insert(PerfEventLink::new(link))
     }
 
     /// Detaches the program.
     ///
     /// See [PerfEvent::attach].
-    pub fn detach(&mut self, link_id: PerfLinkId) -> Result<(), ProgramError> {
+    pub fn detach(&mut self, link_id: PerfEventLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
 
@@ -182,7 +184,16 @@ impl PerfEvent {
     ///
     /// The link will be detached on `Drop` and the caller is now responsible
     /// for managing its lifetime.
-    pub fn take_link(&mut self, link_id: PerfLinkId) -> Result<PerfLink, ProgramError> {
+    pub fn take_link(&mut self, link_id: PerfEventLinkId) -> Result<PerfEventLink, ProgramError> {
         self.data.take_link(link_id)
     }
 }
+
+define_link_wrapper!(
+    /// The link used by [PerfEvent] programs.
+    PerfEventLink,
+    /// The type returned by [PerfEvent::attach]. Can be passed to [PerfEvent::detach].
+    PerfEventLinkId,
+    PerfLinkInner,
+    PerfLinkIdInner
+);

aya/src/programs/probe.rs

@@ -8,7 +8,7 @@ use std::{
 
 use crate::{
     programs::{
-        kprobe::KProbeError, perf_attach, perf_attach::PerfLink, perf_attach_debugfs,
+        kprobe::KProbeError, perf_attach, perf_attach::PerfLinkInner, perf_attach_debugfs,
         trace_point::read_sys_fs_trace_point_id, uprobe::UProbeError, utils::find_tracefs_path,
         Link, ProgramData, ProgramError,
     },
@@ -37,7 +37,7 @@ impl ProbeKind {
     }
 }
 
-pub(crate) fn attach<T: Link + From<PerfLink>>(
+pub(crate) fn attach<T: Link + From<PerfLinkInner>>(
     program_data: &mut ProgramData<T>,
     kind: ProbeKind,
     fn_name: &str,
@@ -49,13 +49,18 @@ pub(crate) fn attach<T: Link + From<PerfLink>>(
     let k_ver = kernel_version().unwrap();
     if k_ver < (4, 17, 0) {
         let (fd, event_alias) = create_as_trace_point(kind, fn_name, offset, pid)?;
-
-        return perf_attach_debugfs(program_data, fd, kind, event_alias);
+        let link = T::from(perf_attach_debugfs(
+            program_data.fd_or_err()?,
+            fd,
+            kind,
+            event_alias,
+        )?);
+        return program_data.links.insert(link);
     };
 
     let fd = create_as_probe(kind, fn_name, offset, pid)?;
-
-    perf_attach(program_data, fd)
+    let link = T::from(perf_attach(program_data.fd_or_err()?, fd)?);
+    program_data.links.insert(link)
 }
 
 pub(crate) fn detach_debug_fs(kind: ProbeKind, event_alias: &str) -> Result<(), ProgramError> {

aya/src/programs/trace_point.rs

@@ -6,7 +6,7 @@ use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT,
     programs::{
         define_link_wrapper, load_program,
-        perf_attach::{perf_attach, PerfLink, PerfLinkId},
+        perf_attach::{perf_attach, PerfLinkIdInner, PerfLinkInner},
         utils::find_tracefs_path,
         ProgramData, ProgramError,
     },
@@ -87,7 +87,8 @@ impl TracePoint {
             }
         })? as i32;
 
-        perf_attach(&mut self.data, fd)
+        let link = perf_attach(self.data.fd_or_err()?, fd)?;
+        self.data.links.insert(TracePointLink::new(link))
     }
 
     /// Detaches from a trace point.
@@ -111,8 +112,8 @@ define_link_wrapper!(
     TracePointLink,
     /// The type returned by [TracePoint::attach]. Can be passed to [TracePoint::detach].
     TracePointLinkId,
-    PerfLink,
-    PerfLinkId
+    PerfLinkInner,
+    PerfLinkIdInner
 );
 
 pub(crate) fn read_sys_fs_trace_point_id(

aya/src/programs/uprobe.rs

@@ -17,7 +17,7 @@ use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
     programs::{
         define_link_wrapper, load_program,
-        perf_attach::{PerfLink, PerfLinkId},
+        perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{attach, ProbeKind},
         ProgramData, ProgramError,
     },
@@ -154,8 +154,8 @@ define_link_wrapper!(
     UProbeLink,
     /// The type returned by [UProbe::attach]. Can be passed to [UProbe::detach].
     UProbeLinkId,
-    PerfLink,
-    PerfLinkId
+    PerfLinkInner,
+    PerfLinkIdInner
 );
 
 /// The type returned when attaching an [`UProbe`] fails.