aya: ProbeKind is Entry/Return

We already have separate types for KProbe and UProbe.

aya/src/bpf.rs

@@ -598,11 +598,11 @@ impl<'a> EbpfLoader<'a> {
                     match &section {
                         ProgramSection::KProbe => Program::KProbe(KProbe {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            kind: ProbeKind::KProbe,
+                            kind: ProbeKind::Entry,
                         }),
                         ProgramSection::KRetProbe => Program::KProbe(KProbe {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            kind: ProbeKind::KRetProbe,
+                            kind: ProbeKind::Return,
                         }),
                         ProgramSection::UProbe { sleepable } => {
                             let mut data =
@@ -612,7 +612,7 @@ impl<'a> EbpfLoader<'a> {
                             }
                             Program::UProbe(UProbe {
                                 data,
-                                kind: ProbeKind::UProbe,
+                                kind: ProbeKind::Entry,
                             })
                         }
                         ProgramSection::URetProbe { sleepable } => {
@@ -623,7 +623,7 @@ impl<'a> EbpfLoader<'a> {
                             }
                             Program::UProbe(UProbe {
                                 data,
-                                kind: ProbeKind::URetProbe,
+                                kind: ProbeKind::Return,
                             })
                         }
                         ProgramSection::TracePoint => Program::TracePoint(TracePoint {

aya/src/programs/kprobe.rs

@@ -1,6 +1,7 @@
 //! Kernel space probes.
 use std::{
     ffi::OsStr,
+    fmt::{self, Write},
     io,
     os::fd::AsFd as _,
     path::{Path, PathBuf},
@@ -15,7 +16,7 @@ use crate::{
         FdLink, LinkError, ProgramData, ProgramError, ProgramType, define_link_wrapper,
         impl_try_into_fdlink, load_program,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
-        probe::{ProbeKind, attach},
+        probe::{Probe, ProbeKind, attach},
     },
     sys::bpf_link_get_info_by_fd,
 };
@@ -59,8 +60,8 @@ impl KProbe {
         load_program(BPF_PROG_TYPE_KPROBE, &mut self.data)
     }
 
-    /// Returns `KProbe` if the program is a `kprobe`, or `KRetProbe` if the
-    /// program is a `kretprobe`.
+    /// Returns [`ProbeKind::Entry`] if the program is a `kprobe`, or
+    /// [`ProbeKind::Return`] if the program is a `kretprobe`.
     pub fn kind(&self) -> ProbeKind {
         self.kind
     }
@@ -81,9 +82,10 @@ impl KProbe {
         fn_name: T,
         offset: u64,
     ) -> Result<KProbeLinkId, ProgramError> {
-        attach(
-            &mut self.data,
-            self.kind,
+        let Self { data, kind } = self;
+        attach::<Self, _>(
+            data,
+            *kind,
             fn_name.as_ref(),
             offset,
             None, // pid
@@ -103,6 +105,23 @@ impl KProbe {
     }
 }
 
+impl Probe for KProbe {
+    const PMU: &'static str = "kprobe";
+
+    type Error = KProbeError;
+
+    fn file_error(filename: PathBuf, io_error: io::Error) -> Self::Error {
+        KProbeError::FileError { filename, io_error }
+    }
+
+    fn write_offset<W: Write>(w: &mut W, kind: ProbeKind, offset: u64) -> fmt::Result {
+        match kind {
+            ProbeKind::Entry => write!(w, "+{offset}"),
+            ProbeKind::Return => Ok(()),
+        }
+    }
+}
+
 define_link_wrapper!(
     KProbeLink,
     KProbeLinkId,

aya/src/programs/perf_attach.rs

@@ -8,10 +8,7 @@ use aya_obj::generated::bpf_attach_type::BPF_PERF_EVENT;
 
 use crate::{
     FEATURES,
-    programs::{
-        FdLink, Link, ProgramError, id_as_key,
-        probe::{ProbeEvent, detach_debug_fs},
-    },
+    programs::{FdLink, Link, ProgramError, id_as_key, probe::ProbeEvent},
     sys::{
         BpfLinkCreateArgs, LinkTarget, PerfEventIoctlRequest, SyscallError, bpf_link_create,
         is_bpf_cookie_supported, perf_event_ioctl,
@@ -72,7 +69,7 @@ impl Link for PerfLink {
         let Self { perf_fd, event } = self;
         let _: io::Result<()> = perf_event_ioctl(perf_fd.as_fd(), PerfEventIoctlRequest::Disable);
         if let Some(event) = event {
-            let _: Result<_, _> = detach_debug_fs(event);
+            let _: Result<(), ProgramError> = event.detach();
         }
 
         Ok(())

aya/src/programs/probe.rs

@@ -1,6 +1,6 @@
 use std::{
     ffi::{OsStr, OsString},
-    fmt::Write as _,
+    fmt::{self, Write},
     fs::{self, OpenOptions},
     io::{self, Write as _},
     os::fd::AsFd as _,
@@ -11,9 +11,8 @@ use std::{
 
 use crate::{
     programs::{
-        Link, ProgramData, ProgramError, 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, perf_attach, perf_attach::PerfLinkInner,
+        perf_attach_debugfs, trace_point::read_sys_fs_trace_point_id, utils::find_tracefs_path,
     },
     sys::{SyscallError, perf_event_open_probe, perf_event_open_trace_point},
     util::KernelVersion,
@@ -24,23 +23,20 @@ static PROBE_NAME_INDEX: AtomicUsize = AtomicUsize::new(0);
 /// Kind of probe program
 #[derive(Debug, Copy, Clone)]
 pub enum ProbeKind {
-    /// Kernel probe
-    KProbe,
-    /// Kernel return probe
-    KRetProbe,
-    /// User space probe
-    UProbe,
-    /// User space return probe
-    URetProbe,
+    /// Probe the entry of the function
+    Entry,
+    /// Probe the return of the function
+    Return,
 }
 
-impl ProbeKind {
-    fn pmu(&self) -> &'static str {
-        match *self {
-            Self::KProbe | Self::KRetProbe => "kprobe",
-            Self::UProbe | Self::URetProbe => "uprobe",
-        }
-    }
+pub(crate) trait Probe {
+    const PMU: &'static str;
+
+    type Error: Into<ProgramError>;
+
+    fn file_error(filename: PathBuf, io_error: io::Error) -> Self::Error;
+
+    fn write_offset<W: Write>(w: &mut W, kind: ProbeKind, offset: u64) -> fmt::Result;
 }
 
 pub(crate) fn lines(bytes: &[u8]) -> impl Iterator<Item = &OsStr> {
@@ -94,11 +90,21 @@ impl OsStringExt for OsStr {
 
 #[derive(Debug)]
 pub(crate) struct ProbeEvent {
-    kind: ProbeKind,
     event_alias: OsString,
+    detach_debug_fs: fn(&OsStr) -> Result<(), ProgramError>,
 }
 
-pub(crate) fn attach<T: Link + From<PerfLinkInner>>(
+impl ProbeEvent {
+    pub(crate) fn detach(self) -> Result<(), ProgramError> {
+        let Self {
+            event_alias,
+            detach_debug_fs,
+        } = self;
+        detach_debug_fs(&event_alias)
+    }
+}
+
+pub(crate) fn attach<P: Probe, T: Link + From<PerfLinkInner>>(
     program_data: &mut ProgramData<T>,
     kind: ProbeKind,
     // NB: the meaning of this argument is different for kprobe/kretprobe and uprobe/uretprobe; in
@@ -120,58 +126,44 @@ pub(crate) fn attach<T: Link + From<PerfLinkInner>>(
         if cookie.is_some() {
             return Err(ProgramError::AttachCookieNotSupported);
         }
-        let (perf_fd, event_alias) = create_as_trace_point(kind, fn_name, offset, pid)?;
-        perf_attach_debugfs(prog_fd, perf_fd, ProbeEvent { kind, event_alias })
+        let (perf_fd, event_alias) = create_as_trace_point::<P>(kind, fn_name, offset, pid)?;
+        perf_attach_debugfs(
+            prog_fd,
+            perf_fd,
+            ProbeEvent {
+                event_alias,
+                detach_debug_fs: detach_debug_fs::<P>,
+            },
+        )
     } else {
-        let perf_fd = create_as_probe(kind, fn_name, offset, pid)?;
+        let perf_fd = create_as_probe::<P>(kind, fn_name, offset, pid)?;
         perf_attach(prog_fd, perf_fd, cookie)
     }?;
     program_data.links.insert(T::from(link))
 }
 
-pub(crate) fn detach_debug_fs(event: ProbeEvent) -> Result<(), ProgramError> {
+fn detach_debug_fs<P: Probe>(event_alias: &OsStr) -> Result<(), ProgramError> {
     let tracefs = find_tracefs_path()?;
 
-    let ProbeEvent {
-        kind,
-        event_alias: _,
-    } = &event;
-    let kind = *kind;
-    let result = delete_probe_event(tracefs, event);
-
-    result.map_err(|(filename, io_error)| match kind {
-        ProbeKind::KProbe | ProbeKind::KRetProbe => {
-            KProbeError::FileError { filename, io_error }.into()
-        }
-        ProbeKind::UProbe | ProbeKind::URetProbe => {
-            UProbeError::FileError { filename, io_error }.into()
-        }
-    })
+    delete_probe_event(tracefs, P::PMU, event_alias)
+        .map_err(|(filename, io_error)| P::file_error(filename, io_error).into())
 }
 
-fn create_as_probe(
+fn create_as_probe<P: Probe>(
     kind: ProbeKind,
     fn_name: &OsStr,
     offset: u64,
     pid: Option<u32>,
 ) -> Result<crate::MockableFd, ProgramError> {
-    let perf_ty = match kind {
-        ProbeKind::KProbe | ProbeKind::KRetProbe => read_sys_fs_perf_type(kind.pmu())
-            .map_err(|(filename, io_error)| KProbeError::FileError { filename, io_error })?,
-        ProbeKind::UProbe | ProbeKind::URetProbe => read_sys_fs_perf_type(kind.pmu())
-            .map_err(|(filename, io_error)| UProbeError::FileError { filename, io_error })?,
-    };
+    let perf_ty = read_sys_fs_perf_type(P::PMU)
+        .map_err(|(filename, io_error)| P::file_error(filename, io_error).into())?;
 
     let ret_bit = match kind {
-        ProbeKind::KRetProbe => Some(
-            read_sys_fs_perf_ret_probe(kind.pmu())
-                .map_err(|(filename, io_error)| KProbeError::FileError { filename, io_error })?,
-        ),
-        ProbeKind::URetProbe => Some(
-            read_sys_fs_perf_ret_probe(kind.pmu())
-                .map_err(|(filename, io_error)| UProbeError::FileError { filename, io_error })?,
+        ProbeKind::Return => Some(
+            read_sys_fs_perf_ret_probe(P::PMU)
+                .map_err(|(filename, io_error)| P::file_error(filename, io_error).into())?,
         ),
-        ProbeKind::UProbe | ProbeKind::KProbe => None,
+        ProbeKind::Entry => None,
     };
 
     perf_event_open_probe(perf_ty, ret_bit, fn_name, offset, pid)
@@ -182,7 +174,7 @@ fn create_as_probe(
         .map_err(Into::into)
 }
 
-fn create_as_trace_point(
+fn create_as_trace_point<P: Probe>(
     kind: ProbeKind,
     name: &OsStr,
     offset: u64,
@@ -190,14 +182,10 @@ fn create_as_trace_point(
 ) -> Result<(crate::MockableFd, OsString), ProgramError> {
     let tracefs = find_tracefs_path()?;
 
-    let event_alias = match kind {
-        ProbeKind::KProbe | ProbeKind::KRetProbe => create_probe_event(tracefs, kind, name, offset)
-            .map_err(|(filename, io_error)| KProbeError::FileError { filename, io_error })?,
-        ProbeKind::UProbe | ProbeKind::URetProbe => create_probe_event(tracefs, kind, name, offset)
-            .map_err(|(filename, io_error)| UProbeError::FileError { filename, io_error })?,
-    };
+    let event_alias = create_probe_event::<P>(tracefs, kind, name, offset)
+        .map_err(|(filename, io_error)| P::file_error(filename, io_error).into())?;
 
-    let category = format!("{}s", kind.pmu());
+    let category = format!("{}s", P::PMU);
     let tpid = read_sys_fs_trace_point_id(tracefs, &category, event_alias.as_ref())?;
     let perf_fd = perf_event_open_trace_point(tpid, pid).map_err(|io_error| SyscallError {
         call: "perf_event_open",
@@ -207,7 +195,7 @@ fn create_as_trace_point(
     Ok((perf_fd, event_alias))
 }
 
-fn create_probe_event(
+fn create_probe_event<P: Probe>(
     tracefs: &Path,
     kind: ProbeKind,
     fn_name: &OsStr,
@@ -216,8 +204,8 @@ fn create_probe_event(
     use std::os::unix::ffi::OsStrExt as _;
 
     let probe_type_prefix = match kind {
-        ProbeKind::KProbe | ProbeKind::UProbe => 'p',
-        ProbeKind::KRetProbe | ProbeKind::URetProbe => 'r',
+        ProbeKind::Entry => 'p',
+        ProbeKind::Return => 'r',
     };
 
     let mut event_alias = OsString::new();
@@ -244,18 +232,14 @@ fn create_probe_event(
     .unwrap();
 
     let mut probe = OsString::new();
-    write!(&mut probe, "{}:{}s/", probe_type_prefix, kind.pmu()).unwrap();
+    write!(&mut probe, "{}:{}s/", probe_type_prefix, P::PMU).unwrap();
     probe.push(&event_alias);
     probe.push(" ");
     probe.push(fn_name);
-    match kind {
-        ProbeKind::KProbe => write!(&mut probe, "+{offset}").unwrap(),
-        ProbeKind::UProbe | ProbeKind::URetProbe => write!(&mut probe, ":{offset:#x}").unwrap(),
-        ProbeKind::KRetProbe => {}
-    };
+    P::write_offset(&mut probe, kind, offset).unwrap();
     probe.push("\n");
 
-    let events_file_name = tracefs.join(format!("{}_events", kind.pmu()));
+    let events_file_name = tracefs.join(format!("{}_events", P::PMU));
     OpenOptions::new()
         .append(true)
         .open(&events_file_name)
@@ -265,11 +249,14 @@ fn create_probe_event(
     Ok(event_alias)
 }
 
-fn delete_probe_event(tracefs: &Path, event: ProbeEvent) -> Result<(), (PathBuf, io::Error)> {
+fn delete_probe_event(
+    tracefs: &Path,
+    pmu: &str,
+    event_alias: &OsStr,
+) -> Result<(), (PathBuf, io::Error)> {
     use std::os::unix::ffi::OsStrExt as _;
 
-    let ProbeEvent { kind, event_alias } = event;
-    let events_file_name = tracefs.join(format!("{}_events", kind.pmu()));
+    let events_file_name = tracefs.join(format!("{}_events", pmu));
 
     fs::read(&events_file_name)
         .and_then(|events| {

aya/src/programs/uprobe.rs

@@ -2,6 +2,7 @@
 use std::{
     error::Error,
     ffi::{CStr, OsStr, OsString},
+    fmt::{self, Write},
     fs,
     io::{self, BufRead as _, Cursor, Read as _},
     mem,
@@ -20,7 +21,7 @@ use crate::{
         FdLink, LinkError, ProgramData, ProgramError, ProgramType, define_link_wrapper,
         impl_try_into_fdlink, load_program,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
-        probe::{OsStringExt as _, ProbeKind, attach},
+        probe::{OsStringExt as _, Probe, ProbeKind, attach},
     },
     sys::bpf_link_get_info_by_fd,
     util::MMap,
@@ -80,8 +81,8 @@ impl UProbe {
         load_program(BPF_PROG_TYPE_KPROBE, &mut self.data)
     }
 
-    /// Returns `UProbe` if the program is a `uprobe`, or `URetProbe` if the
-    /// program is a `uretprobe`.
+    /// Returns [`ProbeKind::Entry`] if the program is a `uprobe`, or
+    /// [`ProbeKind::Return`] if the program is a `uretprobe`.
     pub fn kind(&self) -> ProbeKind {
         self.kind
     }
@@ -129,8 +130,9 @@ impl UProbe {
             offset
         };
 
+        let Self { data, kind } = self;
         let path = path.as_os_str();
-        attach(&mut self.data, self.kind, path, offset, pid, cookie)
+        attach::<Self, _>(data, *kind, path, offset, pid, cookie)
     }
 
     /// Creates a program from a pinned entry on a bpffs.
@@ -145,6 +147,20 @@ impl UProbe {
     }
 }
 
+impl Probe for UProbe {
+    const PMU: &'static str = "uprobe";
+
+    type Error = UProbeError;
+
+    fn file_error(filename: PathBuf, io_error: io::Error) -> Self::Error {
+        UProbeError::FileError { filename, io_error }
+    }
+
+    fn write_offset<W: Write>(w: &mut W, _: ProbeKind, offset: u64) -> fmt::Result {
+        write!(w, ":{offset:#x}")
+    }
+}
+
 fn resolve_attach_path<'a, 'b, 'c, T>(
     target: &'a Path,
     proc_map: Option<&'b ProcMap<T>>,

test/integration-test/src/tests/info.rs

@@ -60,7 +60,7 @@ fn test_loaded_programs() {
         .unwrap();
 
     let mut p: UProbe = unsafe {
-        UProbe::from_program_info(info, "test_uprobe".into(), aya::programs::ProbeKind::UProbe)
+        UProbe::from_program_info(info, "test_uprobe".into(), aya::programs::ProbeKind::Entry)
             .unwrap()
     };
 

test/integration-test/src/tests/load.rs

@@ -557,7 +557,7 @@ fn pin_lifecycle_kprobe() {
     let attach = |prog: &mut P| prog.attach("try_to_wake_up", 0).unwrap();
     let program_pin = "/sys/fs/bpf/aya-kprobe-test-prog";
     let link_pin = "/sys/fs/bpf/aya-kprobe-test-try-to-wake-up";
-    let from_pin = |program_pin: &str| P::from_pin(program_pin, ProbeKind::KProbe).unwrap();
+    let from_pin = |program_pin: &str| P::from_pin(program_pin, ProbeKind::Entry).unwrap();
     run_pin_program_lifecycle_test(
         crate::TEST,
         program_name,
@@ -587,7 +587,7 @@ fn pin_lifecycle_uprobe() {
     };
     let program_pin = "/sys/fs/bpf/aya-uprobe-test-prog";
     let link_pin = "/sys/fs/bpf/aya-uprobe-test-uprobe-function";
-    let from_pin = |program_pin: &str| P::from_pin(program_pin, ProbeKind::UProbe).unwrap();
+    let from_pin = |program_pin: &str| P::from_pin(program_pin, ProbeKind::Entry).unwrap();
     run_pin_program_lifecycle_test(
         crate::TEST,
         program_name,

xtask/public-api/aya.txt

@@ -8014,10 +8014,8 @@ pub unsafe fn aya::programs::LsmAttachType::clone_to_uninit(&self, dest: *mut u8
 impl<T> core::convert::From<T> for aya::programs::LsmAttachType
 pub fn aya::programs::LsmAttachType::from(t: T) -> T
 pub enum aya::programs::ProbeKind
-pub aya::programs::ProbeKind::KProbe
-pub aya::programs::ProbeKind::KRetProbe
-pub aya::programs::ProbeKind::UProbe
-pub aya::programs::ProbeKind::URetProbe
+pub aya::programs::ProbeKind::Entry
+pub aya::programs::ProbeKind::Return
 impl core::clone::Clone for aya::programs::ProbeKind
 pub fn aya::programs::ProbeKind::clone(&self) -> aya::programs::ProbeKind
 impl core::fmt::Debug for aya::programs::ProbeKind