feat: Allow conversions to Program from ProgramInfo

dave-tucker · dave-tucker · commit dc12a1008de2 · 2025-03-06T17:31:07.000Z
Allow for a ProgramInfo to be converted into one of the program types
that we support. This allows for a user of Aya access to reattach,
pin or unload a program that was either, previously loaded, or was
loaded by another process.

Signed-off-by: Dave Tucker &lt;dave@dtucker.co.uk&gt;
diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
@@ -613,15 +613,15 @@ impl<'a> EbpfLoader<'a> {
                         }
                         ProgramSection::CgroupSkb => Program::CgroupSkb(CgroupSkb {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            expected_attach_type: None,
+                            attach_type: None,
                         }),
                         ProgramSection::CgroupSkbIngress => Program::CgroupSkb(CgroupSkb {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            expected_attach_type: Some(CgroupSkbAttachType::Ingress),
+                            attach_type: Some(CgroupSkbAttachType::Ingress),
                         }),
                         ProgramSection::CgroupSkbEgress => Program::CgroupSkb(CgroupSkb {
                             data: ProgramData::new(prog_name, obj, btf_fd, *verifier_log_level),
-                            expected_attach_type: Some(CgroupSkbAttachType::Egress),
+                            attach_type: Some(CgroupSkbAttachType::Egress),
                         }),
                         ProgramSection::CgroupSockAddr { attach_type, .. } => {
                             Program::CgroupSockAddr(CgroupSockAddr {
diff --git a/aya/src/programs/cgroup_skb.rs b/aya/src/programs/cgroup_skb.rs
@@ -57,18 +57,16 @@ use crate::{
 #[doc(alias = "BPF_PROG_TYPE_CGROUP_SKB")]
 pub struct CgroupSkb {
     pub(crate) data: ProgramData<CgroupSkbLink>,
-    pub(crate) expected_attach_type: Option<CgroupSkbAttachType>,
+    pub(crate) attach_type: Option<CgroupSkbAttachType>,
 }
 
 impl CgroupSkb {
     /// Loads the program inside the kernel.
     pub fn load(&mut self) -> Result<(), ProgramError> {
-        self.data.expected_attach_type =
-            self.expected_attach_type
-                .map(|attach_type| match attach_type {
-                    CgroupSkbAttachType::Ingress => BPF_CGROUP_INET_INGRESS,
-                    CgroupSkbAttachType::Egress => BPF_CGROUP_INET_EGRESS,
-                });
+        self.data.expected_attach_type = self.attach_type.map(|attach_type| match attach_type {
+            CgroupSkbAttachType::Ingress => BPF_CGROUP_INET_INGRESS,
+            CgroupSkbAttachType::Egress => BPF_CGROUP_INET_EGRESS,
+        });
         load_program(BPF_PROG_TYPE_CGROUP_SKB, &mut self.data)
     }
 
@@ -79,7 +77,7 @@ impl CgroupSkb {
     /// method returns `None` for programs defined with the generic section
     /// `cgroup/skb`.
     pub fn expected_attach_type(&self) -> &Option<CgroupSkbAttachType> {
-        &self.expected_attach_type
+        &self.attach_type
     }
 
     /// Attaches the program to the given cgroup.
@@ -138,7 +136,7 @@ impl CgroupSkb {
         let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
         Ok(Self {
             data,
-            expected_attach_type: Some(expected_attach_type),
+            attach_type: Some(expected_attach_type),
         })
     }
 }
diff --git a/aya/src/programs/kprobe.rs b/aya/src/programs/kprobe.rs
@@ -9,10 +9,12 @@ use std::{
 use aya_obj::generated::{bpf_link_type, bpf_prog_type::BPF_PROG_TYPE_KPROBE};
 use thiserror::Error;
 
+use super::ProgramType;
 use crate::{
     VerifierLogLevel,
     programs::{
-        FdLink, LinkError, ProgramData, ProgramError, define_link_wrapper, load_program,
+        FdLink, LinkError, ProgramData, ProgramError, ProgramInfo, define_link_wrapper,
+        load_program,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{ProbeKind, attach},
     },
@@ -97,6 +99,62 @@ impl KProbe {
         let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
         Ok(Self { data, kind })
     }
+
+    /// Constructs an instance of [KProbe] from a [ProgramInfo].
+    ///
+    /// This allows you to get a handle to an already loaded program
+    /// from the kernel without having to load it again.
+    ///
+    /// # Safety
+    ///
+    /// The type of program in the [ProgramInfo] may not be of
+    /// [`ProgramType::KProbe`]. This function is unable to
+    /// check since the kernel doesn't expose this information.
+    /// The caller is responsible for ensuring that the program
+    /// type is correct (i.e they previously loaded it) otherwise
+    /// the behavior is undefined.
+    ///
+    /// # Errors
+    ///
+    /// - If the program type is not  [`ProgramType::KProbe`]
+    /// - If the file descriptor of the program cannot be cloned
+    ///
+    /// # Example
+    /// ```no_run
+    /// # use aya::programs::loaded_programs;
+    /// # use aya::programs::{KProbe, ProbeKind, ProgramInfo, ProgramType};
+    ///
+    /// for program in loaded_programs() {
+    ///    let program = program?;
+    ///    if program.program_type()? == ProgramType::KProbe {
+    ///       let p = unsafe { KProbe::from_program_info(None, program, ProbeKind::KProbe)? };
+    ///       // do something with the program
+    ///     }
+    /// }
+    /// # Ok::<(), aya::programs::ProgramError>(())
+    /// ```
+    pub unsafe fn from_program_info(
+        name: Option<String>,
+        info: ProgramInfo,
+        kind: ProbeKind,
+    ) -> Result<Self, ProgramError> {
+        if info.program_type()? != ProgramType::KProbe {
+            return Err(ProgramError::UnexpectedProgramType {});
+        }
+        if kind != ProbeKind::KProbe && kind != ProbeKind::KRetProbe {
+            return Err(ProgramError::UnexpectedProgramType {});
+        }
+        Ok(Self {
+            data: ProgramData::from_bpf_prog_info(
+                name,
+                crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
+                Path::new(""),
+                info.0,
+                VerifierLogLevel::default(),
+            )?,
+            kind,
+        })
+    }
 }
 
 define_link_wrapper!(
diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs
@@ -990,6 +990,157 @@ impl_from_pin!(
     Iter,
 );
 
+macro_rules! impl_from_prog_info {
+    ($($struct_name:ident),+ $(,)?) => {
+        $(
+            impl $struct_name {
+                /// Consructs a program from a [ProgramInfo].
+                ///
+                /// This allows you to get a handle to an already loaded program
+                /// from the kernel without having to load it again.
+                ///
+                /// # Errors
+                ///
+                /// - If the program type is not a match
+                /// - If the file descriptor of the program cannot be cloned
+                pub fn from_program_info(
+                    name: Option<String>,
+                    info: ProgramInfo,
+                ) -> Result<$struct_name, ProgramError> {
+                    if info.program_type()? != ProgramType::$struct_name {
+                        return Err(ProgramError::UnexpectedProgramType {});
+                    }
+                    Ok($struct_name {
+                        data: ProgramData::from_bpf_prog_info(
+                            name,
+                            crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
+                            Path::new(""),
+                            info.0,
+                            VerifierLogLevel::default(),
+                        )?,
+                    })
+                }
+            }
+        )+
+    }
+}
+
+macro_rules! impl_from_prog_info_attach_type {
+    ($($struct_name:ident, $attach_type:ty),+ $(,)?) => {
+        $(
+            impl $struct_name {
+
+                /// Consructs a program from a [ProgramInfo].
+                ///
+                /// This allows you to get a handle to an already loaded program
+                /// from the kernel without having to load it again.
+                ///
+                /// # Errors
+                ///
+                /// - If the program type is not a match
+                /// - If the file descriptor of the program cannot be cloned
+                pub fn from_program_info(
+                    name: Option<String>,
+                    info: ProgramInfo,
+                    attach_type: $attach_type,
+                ) -> Result<$struct_name, ProgramError> {
+                    if info.program_type()? != ProgramType::$struct_name {
+                        return Err(ProgramError::UnexpectedProgramType {});
+                    }
+                    Ok($struct_name {
+                        data: ProgramData::from_bpf_prog_info(
+                            name,
+                            crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
+                            Path::new(""),
+                            info.0,
+                            VerifierLogLevel::default(),
+                        )?,
+                        attach_type,
+                    })
+                }
+            }
+        )+
+    }
+}
+
+macro_rules! impl_from_prog_info_unsafe {
+    ($($struct_name:ident, $prog_type:expr),+ $(,)?) => {
+        $(
+            impl $struct_name {
+
+                /// Consructs a program from a [ProgramInfo].
+                ///
+                /// This allows you to get a handle to an already loaded program
+                /// from the kernel without having to load it again.
+                ///
+                /// # Errors
+                ///
+                /// - If the program type is not a match
+                /// - If the file descriptor of the program cannot be cloned
+                ///
+                /// # Safety
+                ///
+                /// The type of program in the [ProgramInfo] may be
+                /// ambiguous due to missing information in the kernel.
+                /// The caller is responsible for ensuring that the program
+                /// type is correct otherwise the behavior is undefined.
+                pub unsafe fn from_program_info(
+                    name: Option<String>,
+                    info: ProgramInfo,
+                ) -> Result<$struct_name, ProgramError> {
+                    if info.program_type()? != $prog_type {
+                        return Err(ProgramError::UnexpectedProgramType {});
+                    }
+                    Ok($struct_name {
+                        data: ProgramData::from_bpf_prog_info(
+                            name,
+                            crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
+                            Path::new(""),
+                            info.0,
+                            VerifierLogLevel::default(),
+                        )?,
+                    })
+                }
+            }
+        )+
+    }
+}
+
+impl_from_prog_info_unsafe!(
+    BtfTracePoint,
+    ProgramType::Tracing,
+    FEntry,
+    ProgramType::Tracing,
+    FExit,
+    ProgramType::Tracing,
+);
+
+impl_from_prog_info!(
+    CgroupDevice,
+    CgroupSysctl,
+    Extension,
+    LircMode2,
+    Lsm,
+    PerfEvent,
+    RawTracePoint,
+    SkLookup,
+    SkMsg,
+    SockOps,
+    SocketFilter,
+    SchedClassifier,
+);
+
+impl_from_prog_info_attach_type!(
+    CgroupSkb,
+    Option<CgroupSkbAttachType>,
+    CgroupSockAddr,
+    CgroupSockAddrAttachType,
+    CgroupSock,
+    CgroupSockAttachType,
+    CgroupSockopt,
+    CgroupSockoptAttachType,
+);
+
 macro_rules! impl_try_from_program {
     ($($ty:ident),+ $(,)?) => {
         $(
diff --git a/aya/src/programs/probe.rs b/aya/src/programs/probe.rs
@@ -24,7 +24,7 @@ use crate::{
 static PROBE_NAME_INDEX: AtomicUsize = AtomicUsize::new(0);
 
 /// Kind of probe program
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq)]
 pub enum ProbeKind {
     /// Kernel probe
     KProbe,
diff --git a/aya/src/programs/uprobe.rs b/aya/src/programs/uprobe.rs
@@ -18,7 +18,8 @@ use thiserror::Error;
 use crate::{
     VerifierLogLevel,
     programs::{
-        FdLink, LinkError, ProgramData, ProgramError, define_link_wrapper, load_program,
+        FdLink, LinkError, ProgramData, ProgramError, ProgramInfo, ProgramType,
+        define_link_wrapper, load_program,
         perf_attach::{PerfLinkIdInner, PerfLinkInner},
         probe::{OsStringExt as _, ProbeKind, attach},
     },
@@ -139,6 +140,62 @@ impl UProbe {
         let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
         Ok(Self { data, kind })
     }
+
+    /// Constructs an instance of [UProbe] from a [ProgramInfo].
+    ///
+    /// This allows you to get a handle to an already loaded program
+    /// from the kernel without having to load it again.
+    ///
+    /// # Safety
+    ///
+    /// The type of program in the [ProgramInfo] may not be of
+    /// [`ProgramType::KProbe`]. This function is unable to
+    /// check since the kernel doesn't expose this information.
+    /// The caller is responsible for ensuring that the program
+    /// type is correct (i.e they previously loaded it) otherwise
+    /// the behavior is undefined.
+    ///
+    /// # Errors
+    ///
+    /// - If the program type is not  [`ProgramType::KProbe`]
+    /// - If the file descriptor of the program cannot be cloned
+    ///
+    /// # Example
+    /// ```no_run
+    /// # use aya::programs::loaded_programs;
+    /// # use aya::programs::{UProbe, ProbeKind, ProgramInfo, ProgramType};
+    ///
+    /// for program in loaded_programs() {
+    ///    let program = program?;
+    ///    if program.program_type()? == ProgramType::KProbe {
+    ///       let p = unsafe { UProbe::from_program_info(None, program, ProbeKind::UProbe)? };
+    ///       // do something with the program
+    ///     }
+    /// }
+    /// # Ok::<(), aya::programs::ProgramError>(())
+    /// ```
+    pub unsafe fn from_program_info(
+        name: Option<String>,
+        info: ProgramInfo,
+        kind: ProbeKind,
+    ) -> Result<Self, ProgramError> {
+        if info.program_type()? != ProgramType::KProbe {
+            return Err(ProgramError::UnexpectedProgramType {});
+        }
+        if kind != ProbeKind::UProbe && kind != ProbeKind::URetProbe {
+            return Err(ProgramError::UnexpectedProgramType {});
+        }
+        Ok(Self {
+            data: ProgramData::from_bpf_prog_info(
+                name,
+                crate::MockableFd::from_fd(info.fd()?.as_fd().try_clone_to_owned()?),
+                Path::new(""),
+                info.0,
+                VerifierLogLevel::default(),
+            )?,
+            kind,
+        })
+    }
 }
 
 fn resolve_attach_path<'a, 'b, 'c, T>(
diff --git a/test/integration-test/src/tests/info.rs b/test/integration-test/src/tests/info.rs
@@ -55,6 +55,12 @@ fn test_loaded_programs() {
         programs.any(|prog| prog.id() == test_prog.id()),
         KernelVersion::new(4, 13, 0)
     );
+
+    for program in programs {
+        let mut p: SocketFilter =
+            SocketFilter::from_program_info(Some("simple_prog".to_string()), program).unwrap();
+        p.unload().unwrap();
+    }
 }
 
 #[test]
diff --git a/xtask/public-api/aya.txt b/xtask/public-api/aya.txt
