Browse Source

Implement FdLink conversions

Implement TryFrom functions to convert Fdlink to PerfLink/TracePointLink/
KprobeLink, and UprobeLink and vice-versa.

This allows us to pin taken links for perf programs, ultimately
ensuring the link isn't dropped when the loading process exits.

Signed-off-by: Andrew Stoycos <[email protected]>
Andrew Stoycos 1 năm trước cách đây
mục cha
commit
58895db9b4

+ 32 - 2
aya/src/programs/kprobe.rs

@@ -3,13 +3,14 @@ use std::{io, path::Path};
 use thiserror::Error;
 
 use crate::{
-    generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
+    generated::{bpf_link_type, bpf_prog_type::BPF_PROG_TYPE_KPROBE},
     programs::{
         define_link_wrapper, load_program,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{attach, ProbeKind},
-        ProgramData, ProgramError,
+        FdLink, LinkError, ProgramData, ProgramError,
     },
+    sys::bpf_link_get_info_by_fd,
     VerifierLogLevel,
 };
 
@@ -119,3 +120,32 @@ pub enum KProbeError {
         io_error: io::Error,
     },
 }
+
+impl TryFrom<KProbeLink> for FdLink {
+    type Error = LinkError;
+
+    fn try_from(value: KProbeLink) -> Result<Self, Self::Error> {
+        if let PerfLinkInner::FdLink(fd) = value.into_inner() {
+            Ok(fd)
+        } else {
+            Err(LinkError::InvalidLink)
+        }
+    }
+}
+
+impl TryFrom<FdLink> for KProbeLink {
+    type Error = LinkError;
+
+    fn try_from(fd_link: FdLink) -> Result<Self, Self::Error> {
+        let info =
+            bpf_link_get_info_by_fd(fd_link.fd).map_err(|io_error| LinkError::SyscallError {
+                call: "BPF_OBJ_GET_INFO_BY_FD",
+                code: 0,
+                io_error,
+            })?;
+        if info.type_ == (bpf_link_type::BPF_LINK_TYPE_KPROBE_MULTI as u32) {
+            return Ok(KProbeLink::new(PerfLinkInner::FdLink(fd_link)));
+        }
+        Err(LinkError::InvalidLink)
+    }
+}

+ 32 - 2
aya/src/programs/perf_event.rs

@@ -6,6 +6,7 @@ pub use crate::generated::{
 
 use crate::{
     generated::{
+        bpf_link_type,
         bpf_prog_type::BPF_PROG_TYPE_PERF_EVENT,
         perf_type_id::{
             PERF_TYPE_BREAKPOINT, PERF_TYPE_HARDWARE, PERF_TYPE_HW_CACHE, PERF_TYPE_RAW,
@@ -16,9 +17,9 @@ use crate::{
         links::define_link_wrapper,
         load_program, perf_attach,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
-        ProgramData, ProgramError,
+        FdLink, LinkError, ProgramData, ProgramError,
     },
-    sys::perf_event_open,
+    sys::{bpf_link_get_info_by_fd, perf_event_open},
 };
 
 /// The type of perf event
@@ -189,6 +190,35 @@ impl PerfEvent {
     }
 }
 
+impl TryFrom<PerfEventLink> for FdLink {
+    type Error = LinkError;
+
+    fn try_from(value: PerfEventLink) -> Result<Self, Self::Error> {
+        if let PerfLinkInner::FdLink(fd) = value.into_inner() {
+            Ok(fd)
+        } else {
+            Err(LinkError::InvalidLink)
+        }
+    }
+}
+
+impl TryFrom<FdLink> for PerfEventLink {
+    type Error = LinkError;
+
+    fn try_from(fd_link: FdLink) -> Result<Self, Self::Error> {
+        let info =
+            bpf_link_get_info_by_fd(fd_link.fd).map_err(|io_error| LinkError::SyscallError {
+                call: "BPF_OBJ_GET_INFO_BY_FD",
+                code: 0,
+                io_error,
+            })?;
+        if info.type_ == (bpf_link_type::BPF_LINK_TYPE_PERF_EVENT as u32) {
+            return Ok(PerfEventLink::new(PerfLinkInner::FdLink(fd_link)));
+        }
+        Err(LinkError::InvalidLink)
+    }
+}
+
 define_link_wrapper!(
     /// The link used by [PerfEvent] programs.
     PerfEventLink,

+ 32 - 3
aya/src/programs/trace_point.rs

@@ -3,14 +3,14 @@ use std::{fs, io, path::Path};
 use thiserror::Error;
 
 use crate::{
-    generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT,
+    generated::{bpf_link_type, bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT},
     programs::{
         define_link_wrapper, load_program,
         perf_attach::{perf_attach, PerfLinkIdInner, PerfLinkInner},
         utils::find_tracefs_path,
-        ProgramData, ProgramError,
+        FdLink, LinkError, ProgramData, ProgramError,
     },
-    sys::perf_event_open_trace_point,
+    sys::{bpf_link_get_info_by_fd, perf_event_open_trace_point},
 };
 
 /// The type returned when attaching a [`TracePoint`] fails.
@@ -116,6 +116,35 @@ define_link_wrapper!(
     PerfLinkIdInner
 );
 
+impl TryFrom<TracePointLink> for FdLink {
+    type Error = LinkError;
+
+    fn try_from(value: TracePointLink) -> Result<Self, Self::Error> {
+        if let PerfLinkInner::FdLink(fd) = value.into_inner() {
+            Ok(fd)
+        } else {
+            Err(LinkError::InvalidLink)
+        }
+    }
+}
+
+impl TryFrom<FdLink> for TracePointLink {
+    type Error = LinkError;
+
+    fn try_from(fd_link: FdLink) -> Result<Self, Self::Error> {
+        let info =
+            bpf_link_get_info_by_fd(fd_link.fd).map_err(|io_error| LinkError::SyscallError {
+                call: "BPF_OBJ_GET_INFO_BY_FD",
+                code: 0,
+                io_error,
+            })?;
+        if info.type_ == (bpf_link_type::BPF_LINK_TYPE_TRACING as u32) {
+            return Ok(TracePointLink::new(PerfLinkInner::FdLink(fd_link)));
+        }
+        Err(LinkError::InvalidLink)
+    }
+}
+
 pub(crate) fn read_sys_fs_trace_point_id(
     tracefs: &Path,
     category: &str,

+ 32 - 2
aya/src/programs/uprobe.rs

@@ -14,13 +14,14 @@ use std::{
 use thiserror::Error;
 
 use crate::{
-    generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
+    generated::{bpf_link_type, bpf_prog_type::BPF_PROG_TYPE_KPROBE},
     programs::{
         define_link_wrapper, load_program,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{attach, ProbeKind},
-        ProgramData, ProgramError,
+        FdLink, LinkError, ProgramData, ProgramError,
     },
+    sys::bpf_link_get_info_by_fd,
     VerifierLogLevel,
 };
 
@@ -160,6 +161,35 @@ define_link_wrapper!(
     PerfLinkIdInner
 );
 
+impl TryFrom<UProbeLink> for FdLink {
+    type Error = LinkError;
+
+    fn try_from(value: UProbeLink) -> Result<Self, Self::Error> {
+        if let PerfLinkInner::FdLink(fd) = value.into_inner() {
+            Ok(fd)
+        } else {
+            Err(LinkError::InvalidLink)
+        }
+    }
+}
+
+impl TryFrom<FdLink> for UProbeLink {
+    type Error = LinkError;
+
+    fn try_from(fd_link: FdLink) -> Result<Self, Self::Error> {
+        let info =
+            bpf_link_get_info_by_fd(fd_link.fd).map_err(|io_error| LinkError::SyscallError {
+                call: "BPF_OBJ_GET_INFO_BY_FD",
+                code: 0,
+                io_error,
+            })?;
+        if info.type_ == (bpf_link_type::BPF_LINK_TYPE_TRACING as u32) {
+            return Ok(UProbeLink::new(PerfLinkInner::FdLink(fd_link)));
+        }
+        Err(LinkError::InvalidLink)
+    }
+}
+
 /// The type returned when attaching an [`UProbe`] fails.
 #[derive(Debug, Error)]
 pub enum UProbeError {