Bladeren bron

Merge pull request #249 from alessandrod/new-links

aya: rework links
Dave Tucker 3 jaren geleden
bovenliggende
commit
b039ac5

+ 1 - 1
aya/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "aya"
-version = "0.10.7"
+version = "0.11.0-dev.0"
 description = "An eBPF library with a focus on developer experience and operability."
 keywords = ["ebpf", "bpf", "linux", "kernel"]
 license = "MIT OR Apache-2.0"

+ 55 - 38
aya/src/bpf.rs

@@ -405,57 +405,58 @@ impl<'a> BpfLoader<'a> {
                 } else {
                     None
                 };
-                let data = ProgramData {
-                    name: prog_name,
-                    obj,
-                    fd: None,
-                    links: Vec::new(),
-                    expected_attach_type: None,
-                    attach_btf_obj_fd: None,
-                    attach_btf_id: None,
-                    attach_prog_fd: None,
-                    btf_fd,
-                };
+                let section = obj.section.clone();
+
                 let program = if self.extensions.contains(name.as_str()) {
-                    Program::Extension(Extension { data })
+                    Program::Extension(Extension {
+                        data: ProgramData::new(prog_name, obj, btf_fd),
+                    })
                 } else {
-                    match &data.obj.section {
+                    match &section {
                         ProgramSection::KProbe { .. } => Program::KProbe(KProbe {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             kind: ProbeKind::KProbe,
                         }),
                         ProgramSection::KRetProbe { .. } => Program::KProbe(KProbe {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             kind: ProbeKind::KRetProbe,
                         }),
                         ProgramSection::UProbe { .. } => Program::UProbe(UProbe {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             kind: ProbeKind::UProbe,
                         }),
                         ProgramSection::URetProbe { .. } => Program::UProbe(UProbe {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             kind: ProbeKind::URetProbe,
                         }),
-                        ProgramSection::TracePoint { .. } => {
-                            Program::TracePoint(TracePoint { data })
-                        }
+                        ProgramSection::TracePoint { .. } => Program::TracePoint(TracePoint {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                         ProgramSection::SocketFilter { .. } => {
-                            Program::SocketFilter(SocketFilter { data })
+                            Program::SocketFilter(SocketFilter {
+                                data: ProgramData::new(prog_name, obj, btf_fd),
+                            })
                         }
-                        ProgramSection::Xdp { .. } => Program::Xdp(Xdp { data }),
-                        ProgramSection::SkMsg { .. } => Program::SkMsg(SkMsg { data }),
+                        ProgramSection::Xdp { .. } => Program::Xdp(Xdp {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
+                        ProgramSection::SkMsg { .. } => Program::SkMsg(SkMsg {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                         ProgramSection::SkSkbStreamParser { .. } => Program::SkSkb(SkSkb {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             kind: SkSkbKind::StreamParser,
                         }),
                         ProgramSection::SkSkbStreamVerdict { .. } => Program::SkSkb(SkSkb {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             kind: SkSkbKind::StreamVerdict,
                         }),
-                        ProgramSection::SockOps { .. } => Program::SockOps(SockOps { data }),
+                        ProgramSection::SockOps { .. } => Program::SockOps(SockOps {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                         ProgramSection::SchedClassifier { .. } => {
                             Program::SchedClassifier(SchedClassifier {
-                                data,
+                                data: ProgramData::new(prog_name, obj, btf_fd),
                                 name: unsafe {
                                     CString::from_vec_unchecked(Vec::from(name.clone()))
                                         .into_boxed_c_str()
@@ -463,29 +464,45 @@ impl<'a> BpfLoader<'a> {
                             })
                         }
                         ProgramSection::CgroupSkb { .. } => Program::CgroupSkb(CgroupSkb {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             expected_attach_type: None,
                         }),
                         ProgramSection::CgroupSkbIngress { .. } => Program::CgroupSkb(CgroupSkb {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             expected_attach_type: Some(CgroupSkbAttachType::Ingress),
                         }),
                         ProgramSection::CgroupSkbEgress { .. } => Program::CgroupSkb(CgroupSkb {
-                            data,
+                            data: ProgramData::new(prog_name, obj, btf_fd),
                             expected_attach_type: Some(CgroupSkbAttachType::Egress),
                         }),
-                        ProgramSection::LircMode2 { .. } => Program::LircMode2(LircMode2 { data }),
-                        ProgramSection::PerfEvent { .. } => Program::PerfEvent(PerfEvent { data }),
+                        ProgramSection::LircMode2 { .. } => Program::LircMode2(LircMode2 {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
+                        ProgramSection::PerfEvent { .. } => Program::PerfEvent(PerfEvent {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                         ProgramSection::RawTracePoint { .. } => {
-                            Program::RawTracePoint(RawTracePoint { data })
+                            Program::RawTracePoint(RawTracePoint {
+                                data: ProgramData::new(prog_name, obj, btf_fd),
+                            })
                         }
-                        ProgramSection::Lsm { .. } => Program::Lsm(Lsm { data }),
+                        ProgramSection::Lsm { .. } => Program::Lsm(Lsm {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                         ProgramSection::BtfTracePoint { .. } => {
-                            Program::BtfTracePoint(BtfTracePoint { data })
+                            Program::BtfTracePoint(BtfTracePoint {
+                                data: ProgramData::new(prog_name, obj, btf_fd),
+                            })
                         }
-                        ProgramSection::FEntry { .. } => Program::FEntry(FEntry { data }),
-                        ProgramSection::FExit { .. } => Program::FExit(FExit { data }),
-                        ProgramSection::Extension { .. } => Program::Extension(Extension { data }),
+                        ProgramSection::FEntry { .. } => Program::FEntry(FEntry {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
+                        ProgramSection::FExit { .. } => Program::FExit(FExit {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
+                        ProgramSection::Extension { .. } => Program::Extension(Extension {
+                            data: ProgramData::new(prog_name, obj, btf_fd),
+                        }),
                     }
                 };
                 (name, program)

+ 64 - 12
aya/src/programs/cgroup_skb.rs

@@ -1,16 +1,19 @@
-use std::os::unix::prelude::{AsRawFd, RawFd};
+use std::{
+    hash::Hash,
+    os::unix::prelude::{AsRawFd, RawFd},
+};
 
 use crate::{
     generated::{
         bpf_attach_type::{BPF_CGROUP_INET_EGRESS, BPF_CGROUP_INET_INGRESS},
         bpf_prog_type::BPF_PROG_TYPE_CGROUP_SKB,
     },
-    programs::{load_program, LinkRef, ProgAttachLink, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, FdLink, Link, ProgAttachLink, ProgramData, ProgramError,
+    },
     sys::{bpf_link_create, bpf_prog_attach, kernel_version},
 };
 
-use super::FdLink;
-
 /// A program used to inspect or filter network activity for a given cgroup.
 ///
 /// [`CgroupSkb`] programs can be used to inspect or filter network activity
@@ -51,14 +54,12 @@ use super::FdLink;
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_CGROUP_SKB")]
 pub struct CgroupSkb {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<CgroupSkbLink>,
     pub(crate) expected_attach_type: Option<CgroupSkbAttachType>,
 }
 
 impl CgroupSkb {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_CGROUP_SKB, &mut self.data)
     }
@@ -74,11 +75,13 @@ impl CgroupSkb {
     }
 
     /// Attaches the program to the given cgroup.
+    ///
+    /// The returned value can be used to detach, see [CgroupSkb::detach].
     pub fn attach<T: AsRawFd>(
         &mut self,
         cgroup: T,
         attach_type: CgroupSkbAttachType,
-    ) -> Result<LinkRef, ProgramError> {
+    ) -> Result<CgroupSkbLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let cgroup_fd = cgroup.as_raw_fd();
 
@@ -94,7 +97,9 @@ impl CgroupSkb {
                     io_error,
                 },
             )? as RawFd;
-            Ok(self.data.link(FdLink { fd: Some(link_fd) }))
+            self.data
+                .links
+                .insert(CgroupSkbLink(CgroupSkbLinkInner::Fd(FdLink::new(link_fd))))
         } else {
             bpf_prog_attach(prog_fd, cgroup_fd, attach_type).map_err(|(_, io_error)| {
                 ProgramError::SyscallError {
@@ -103,13 +108,60 @@ impl CgroupSkb {
                 }
             })?;
 
-            Ok(self
-                .data
-                .link(ProgAttachLink::new(prog_fd, cgroup_fd, attach_type)))
+            self.data
+                .links
+                .insert(CgroupSkbLink(CgroupSkbLinkInner::ProgAttach(
+                    ProgAttachLink::new(prog_fd, cgroup_fd, attach_type),
+                )))
+        }
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [CgroupSkb::attach].
+    pub fn detach(&mut self, link_id: CgroupSkbLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
+}
+
+#[derive(Debug, Hash, Eq, PartialEq)]
+enum CgroupSkbLinkIdInner {
+    Fd(<FdLink as Link>::Id),
+    ProgAttach(<ProgAttachLink as Link>::Id),
+}
+
+#[derive(Debug)]
+enum CgroupSkbLinkInner {
+    Fd(FdLink),
+    ProgAttach(ProgAttachLink),
+}
+
+impl Link for CgroupSkbLinkInner {
+    type Id = CgroupSkbLinkIdInner;
+
+    fn id(&self) -> Self::Id {
+        match self {
+            CgroupSkbLinkInner::Fd(fd) => CgroupSkbLinkIdInner::Fd(fd.id()),
+            CgroupSkbLinkInner::ProgAttach(p) => CgroupSkbLinkIdInner::ProgAttach(p.id()),
+        }
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        match self {
+            CgroupSkbLinkInner::Fd(fd) => fd.detach(),
+            CgroupSkbLinkInner::ProgAttach(p) => p.detach(),
         }
     }
 }
 
+define_link_wrapper!(
+    CgroupSkbLink,
+    /// The type returned by [CgroupSkb::attach]. Can be passed to [CgroupSkb::detach].
+    CgroupSkbLinkId,
+    CgroupSkbLinkInner,
+    CgroupSkbLinkIdInner
+);
+
 /// Defines where to attach a [`CgroupSkb`] program.
 #[derive(Copy, Clone, Debug)]
 pub enum CgroupSkbAttachType {

+ 24 - 8
aya/src/programs/extension.rs

@@ -6,7 +6,7 @@ use object::Endianness;
 use crate::{
     generated::{bpf_attach_type::BPF_CGROUP_INET_INGRESS, bpf_prog_type::BPF_PROG_TYPE_EXT},
     obj::btf::BtfKind,
-    programs::{load_program, FdLink, LinkRef, ProgramData, ProgramError},
+    programs::{define_link_wrapper, load_program, FdLink, FdLinkId, ProgramData, ProgramError},
     sys::{self, bpf_link_create},
     Btf,
 };
@@ -48,7 +48,7 @@ pub enum ExtensionError {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_EXT")]
 pub struct Extension {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<ExtensionLink>,
 }
 
 impl Extension {
@@ -63,8 +63,6 @@ impl Extension {
     /// The extension code will be loaded but inactive until it's attached.
     /// There are no restrictions on what functions may be replaced, so you could replace
     /// the main entry point of your program with an extension.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load<T: AsRawFd>(&mut self, program: T, func_name: &str) -> Result<(), ProgramError> {
         let target_prog_fd = program.as_raw_fd();
 
@@ -122,11 +120,13 @@ impl Extension {
         load_program(BPF_PROG_TYPE_EXT, &mut self.data)
     }
 
-    /// Attaches the extension
+    /// Attaches the extension.
     ///
     /// Attaches the extension effectively replacing the original target function.
-    /// Detaching the returned link restores the original function.
-    pub fn attach(&mut self) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach the extension and restore the
+    /// original function, see [Extension::detach].
+    pub fn attach(&mut self) -> Result<ExtensionLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let target_fd = self.data.attach_prog_fd.ok_or(ProgramError::NotLoaded)?;
         let btf_id = self.data.attach_btf_id.ok_or(ProgramError::NotLoaded)?;
@@ -136,6 +136,22 @@ impl Extension {
                 call: "bpf_link_create".to_owned(),
                 io_error,
             })? as RawFd;
-        Ok(self.data.link(FdLink { fd: Some(link_fd) }))
+        self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
+    }
+
+    /// Detaches the extension.
+    ///
+    /// Detaching restores the original code overridden by the extension program.
+    /// See [Extension::attach].
+    pub fn detach(&mut self, link_id: ExtensionLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 }
+
+define_link_wrapper!(
+    ExtensionLink,
+    /// The type returned by [Extension::attach]. Can be passed to [Extension::detach].
+    ExtensionLinkId,
+    FdLink,
+    FdLinkId
+);

+ 25 - 6
aya/src/programs/fentry.rs

@@ -1,8 +1,12 @@
 //! fentry programs.
+
 use crate::{
     generated::{bpf_attach_type::BPF_TRACE_FENTRY, bpf_prog_type::BPF_PROG_TYPE_TRACING},
     obj::btf::{Btf, BtfKind},
-    programs::{load_program, utils::attach_raw_tracepoint, LinkRef, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
+        ProgramData, ProgramError,
+    },
 };
 
 /// A program that can be attached to the entry point of (almost) any kernel
@@ -43,14 +47,12 @@ use crate::{
 #[doc(alias = "BPF_TRACE_FENTRY")]
 #[doc(alias = "BPF_PROG_TYPE_TRACING")]
 pub struct FEntry {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<FEntryLink>,
 }
 
 impl FEntry {
     /// Loads the program inside the kernel.
     ///
-    /// See also [`Program::load`](crate::programs::Program::load).
-    ///
     /// Loads the program so it's executed when the kernel function `fn_name`
     /// is entered. The `btf` argument must contain the BTF info for the
     /// running kernel.
@@ -60,8 +62,25 @@ impl FEntry {
         load_program(BPF_PROG_TYPE_TRACING, &mut self.data)
     }
 
-    /// Attaches the program
-    pub fn attach(&mut self) -> Result<LinkRef, ProgramError> {
+    /// Attaches the program.
+    ///
+    /// The returned value can be used to detach, see [FEntry::detach].
+    pub fn attach(&mut self) -> Result<FEntryLinkId, ProgramError> {
         attach_raw_tracepoint(&mut self.data, None)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [FEntry::attach].
+    pub fn detach(&mut self, link_id: FEntryLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
+
+define_link_wrapper!(
+    FEntryLink,
+    /// The type returned by [FEntry::attach]. Can be passed to [FEntry::detach].
+    FEntryLinkId,
+    FdLink,
+    FdLinkId
+);

+ 25 - 6
aya/src/programs/fexit.rs

@@ -1,8 +1,12 @@
 //! fexit programs.
+
 use crate::{
     generated::{bpf_attach_type::BPF_TRACE_FEXIT, bpf_prog_type::BPF_PROG_TYPE_TRACING},
     obj::btf::{Btf, BtfKind},
-    programs::{load_program, utils::attach_raw_tracepoint, LinkRef, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
+        ProgramData, ProgramError,
+    },
 };
 
 /// A program that can be attached to the exit point of (almost) anny kernel
@@ -43,14 +47,12 @@ use crate::{
 #[doc(alias = "BPF_TRACE_FEXIT")]
 #[doc(alias = "BPF_PROG_TYPE_TRACING")]
 pub struct FExit {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<FExitLink>,
 }
 
 impl FExit {
     /// Loads the program inside the kernel.
     ///
-    /// See also [`Program::load`](crate::programs::Program::load).
-    ///
     /// Loads the program so it's executed when the kernel function `fn_name`
     /// is exited. The `btf` argument must contain the BTF info for the running
     /// kernel.
@@ -60,8 +62,25 @@ impl FExit {
         load_program(BPF_PROG_TYPE_TRACING, &mut self.data)
     }
 
-    /// Attaches the program
-    pub fn attach(&mut self) -> Result<LinkRef, ProgramError> {
+    /// Attaches the program.
+    ///
+    /// The returned value can be used to detach, see [FExit::detach].
+    pub fn attach(&mut self) -> Result<FExitLinkId, ProgramError> {
         attach_raw_tracepoint(&mut self.data, None)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [FExit::attach].
+    pub fn detach(&mut self, link_id: FExitLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
+
+define_link_wrapper!(
+    FExitLink,
+    /// The type returned by [FExit::attach]. Can be passed to [FExit::detach].
+    FExitLinkId,
+    FdLink,
+    FdLinkId
+);

+ 22 - 6
aya/src/programs/kprobe.rs

@@ -5,9 +5,10 @@ use thiserror::Error;
 use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
     programs::{
-        load_program,
+        define_link_wrapper, load_program,
+        perf_attach::{PerfLink, PerfLinkId},
         probe::{attach, ProbeKind},
-        LinkRef, ProgramData, ProgramError,
+        ProgramData, ProgramError,
     },
 };
 
@@ -38,14 +39,12 @@ use crate::{
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_KPROBE")]
 pub struct KProbe {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<KProbeLink>,
     pub(crate) kind: ProbeKind,
 }
 
 impl KProbe {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_KPROBE, &mut self.data)
     }
@@ -65,11 +64,28 @@ impl KProbe {
     /// If the program is a `kprobe`, it is attached to the *start* address of the target function.
     /// Conversely if the program is a `kretprobe`, it is attached to the return address of the
     /// target function.
-    pub fn attach(&mut self, fn_name: &str, offset: u64) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach from the given function, see [KProbe::detach].
+    pub fn attach(&mut self, fn_name: &str, offset: u64) -> Result<KProbeLinkId, ProgramError> {
         attach(&mut self.data, self.kind, fn_name, offset, None)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [KProbe::attach].
+    pub fn detach(&mut self, link_id: KProbeLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
 
+define_link_wrapper!(
+    KProbeLink,
+    /// The type returned by [KProbe::attach]. Can be passed to [KProbe::detach].
+    KProbeLinkId,
+    PerfLink,
+    PerfLinkId
+);
+
 /// The type returned when attaching a [`KProbe`] fails.
 #[derive(Debug, Error)]
 pub enum KProbeError {

+ 260 - 0
aya/src/programs/links.rs

@@ -0,0 +1,260 @@
+use libc::{close, dup};
+use std::{
+    collections::{hash_map::Entry, HashMap},
+    os::unix::prelude::RawFd,
+};
+
+use crate::{generated::bpf_attach_type, programs::ProgramError, sys::bpf_prog_detach};
+
+pub(crate) trait Link: std::fmt::Debug + 'static {
+    type Id: std::fmt::Debug + std::hash::Hash + Eq + PartialEq;
+
+    fn id(&self) -> Self::Id;
+
+    fn detach(self) -> Result<(), ProgramError>;
+}
+
+#[derive(Debug)]
+pub(crate) struct LinkMap<T: Link> {
+    links: HashMap<T::Id, T>,
+}
+
+impl<T: Link> LinkMap<T> {
+    pub(crate) fn new() -> LinkMap<T> {
+        LinkMap {
+            links: HashMap::new(),
+        }
+    }
+
+    pub(crate) fn insert(&mut self, link: T) -> Result<T::Id, ProgramError> {
+        let id = link.id();
+
+        match self.links.entry(link.id()) {
+            Entry::Occupied(_) => return Err(ProgramError::AlreadyAttached),
+            Entry::Vacant(e) => e.insert(link),
+        };
+
+        Ok(id)
+    }
+
+    pub(crate) fn remove(&mut self, link_id: T::Id) -> Result<(), ProgramError> {
+        self.links
+            .remove(&link_id)
+            .ok_or(ProgramError::NotAttached)?
+            .detach()
+    }
+}
+
+impl<T: Link> Drop for LinkMap<T> {
+    fn drop(&mut self) {
+        for (_, link) in self.links.drain() {
+            let _ = link.detach();
+        }
+    }
+}
+
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub(crate) struct FdLinkId(pub(crate) RawFd);
+
+#[derive(Debug)]
+pub(crate) struct FdLink {
+    fd: RawFd,
+}
+
+impl FdLink {
+    pub(crate) fn new(fd: RawFd) -> FdLink {
+        FdLink { fd }
+    }
+}
+
+impl Link for FdLink {
+    type Id = FdLinkId;
+
+    fn id(&self) -> Self::Id {
+        FdLinkId(self.fd)
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        unsafe { close(self.fd) };
+        Ok(())
+    }
+}
+
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub(crate) struct ProgAttachLinkId(RawFd, RawFd, bpf_attach_type);
+
+#[derive(Debug)]
+pub(crate) struct ProgAttachLink {
+    prog_fd: RawFd,
+    target_fd: RawFd,
+    attach_type: bpf_attach_type,
+}
+
+impl ProgAttachLink {
+    pub(crate) fn new(
+        prog_fd: RawFd,
+        target_fd: RawFd,
+        attach_type: bpf_attach_type,
+    ) -> ProgAttachLink {
+        ProgAttachLink {
+            prog_fd,
+            target_fd: unsafe { dup(target_fd) },
+            attach_type,
+        }
+    }
+}
+
+impl Link for ProgAttachLink {
+    type Id = ProgAttachLinkId;
+
+    fn id(&self) -> Self::Id {
+        ProgAttachLinkId(self.prog_fd, self.target_fd, self.attach_type)
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        let _ = bpf_prog_detach(self.prog_fd, self.target_fd, self.attach_type);
+        unsafe { close(self.target_fd) };
+        Ok(())
+    }
+}
+
+macro_rules! define_link_wrapper {
+    ($wrapper:ident, #[$doc:meta] $wrapper_id:ident, $base:ident, $base_id:ident) => {
+        #[$doc]
+        #[derive(Debug, Hash, Eq, PartialEq)]
+        pub struct $wrapper_id($base_id);
+
+        #[derive(Debug)]
+        pub(crate) struct $wrapper($base);
+
+        impl crate::programs::Link for $wrapper {
+            type Id = $wrapper_id;
+
+            fn id(&self) -> Self::Id {
+                $wrapper_id(self.0.id())
+            }
+
+            fn detach(self) -> Result<(), ProgramError> {
+                self.0.detach()
+            }
+        }
+
+        impl From<$base> for $wrapper {
+            fn from(b: $base) -> $wrapper {
+                $wrapper(b)
+            }
+        }
+    };
+}
+
+pub(crate) use define_link_wrapper;
+
+#[cfg(test)]
+mod tests {
+    use std::{cell::RefCell, rc::Rc};
+
+    use crate::programs::ProgramError;
+
+    use super::{Link, LinkMap};
+
+    #[derive(Debug, Hash, Eq, PartialEq)]
+    struct TestLinkId(u8, u8);
+
+    #[derive(Debug)]
+    struct TestLink {
+        id: (u8, u8),
+        detached: Rc<RefCell<u8>>,
+    }
+
+    impl TestLink {
+        fn new(a: u8, b: u8) -> TestLink {
+            TestLink {
+                id: (a, b),
+                detached: Rc::new(RefCell::new(0)),
+            }
+        }
+    }
+
+    impl Link for TestLink {
+        type Id = TestLinkId;
+
+        fn id(&self) -> Self::Id {
+            TestLinkId(self.id.0, self.id.1)
+        }
+
+        fn detach(self) -> Result<(), ProgramError> {
+            *self.detached.borrow_mut() += 1;
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn test_link_map() {
+        let mut links = LinkMap::new();
+        let l1 = TestLink::new(1, 2);
+        let l1_detached = Rc::clone(&l1.detached);
+        let l2 = TestLink::new(1, 3);
+        let l2_detached = Rc::clone(&l2.detached);
+
+        let id1 = links.insert(l1).unwrap();
+        let id2 = links.insert(l2).unwrap();
+
+        assert!(*l1_detached.borrow() == 0);
+        assert!(*l2_detached.borrow() == 0);
+
+        assert!(links.remove(id1).is_ok());
+        assert!(*l1_detached.borrow() == 1);
+        assert!(*l2_detached.borrow() == 0);
+
+        assert!(links.remove(id2).is_ok());
+        assert!(*l1_detached.borrow() == 1);
+        assert!(*l2_detached.borrow() == 1);
+    }
+
+    #[test]
+    fn test_already_attached() {
+        let mut links = LinkMap::new();
+
+        links.insert(TestLink::new(1, 2)).unwrap();
+        assert!(matches!(
+            links.insert(TestLink::new(1, 2)),
+            Err(ProgramError::AlreadyAttached)
+        ));
+    }
+
+    #[test]
+    fn test_not_attached() {
+        let mut links = LinkMap::new();
+
+        let l1 = TestLink::new(1, 2);
+        let l1_id1 = l1.id();
+        let l1_id2 = l1.id();
+        links.insert(TestLink::new(1, 2)).unwrap();
+        links.remove(l1_id1).unwrap();
+        assert!(matches!(
+            links.remove(l1_id2),
+            Err(ProgramError::NotAttached)
+        ));
+    }
+
+    #[test]
+    fn test_drop_detach() {
+        let l1 = TestLink::new(1, 2);
+        let l1_detached = Rc::clone(&l1.detached);
+        let l2 = TestLink::new(1, 3);
+        let l2_detached = Rc::clone(&l2.detached);
+
+        {
+            let mut links = LinkMap::new();
+            let id1 = links.insert(l1).unwrap();
+            links.insert(l2).unwrap();
+            // manually remove one link
+            assert!(links.remove(id1).is_ok());
+            assert!(*l1_detached.borrow() == 1);
+            assert!(*l2_detached.borrow() == 0);
+        }
+        // remove the other on drop
+        assert!(*l1_detached.borrow() == 1);
+        assert!(*l2_detached.borrow() == 1);
+    }
+}

+ 36 - 39
aya/src/programs/lirc_mode2.rs

@@ -2,7 +2,7 @@ use std::os::unix::prelude::{AsRawFd, RawFd};
 
 use crate::{
     generated::{bpf_attach_type::BPF_LIRC_MODE2, bpf_prog_type::BPF_PROG_TYPE_LIRC_MODE2},
-    programs::{load_program, query, Link, LinkRef, ProgramData, ProgramError, ProgramInfo},
+    programs::{load_program, query, Link, ProgramData, ProgramError, ProgramInfo},
     sys::{bpf_obj_get_info_by_fd, bpf_prog_attach, bpf_prog_detach, bpf_prog_get_fd_by_id},
 };
 
@@ -48,19 +48,19 @@ use libc::{close, dup};
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_LIRC_MODE2")]
 pub struct LircMode2 {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<LircLink>,
 }
 
 impl LircMode2 {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_LIRC_MODE2, &mut self.data)
     }
 
     /// Attaches the program to the given lirc device.
-    pub fn attach<T: AsRawFd>(&mut self, lircdev: T) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [LircMode2::detach].
+    pub fn attach<T: AsRawFd>(&mut self, lircdev: T) -> Result<LircLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let lircdev_fd = lircdev.as_raw_fd();
 
@@ -71,7 +71,14 @@ impl LircMode2 {
             }
         })?;
 
-        Ok(self.data.link(LircLink::new(prog_fd, lircdev_fd)))
+        self.data.links.insert(LircLink::new(prog_fd, lircdev_fd))
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [LircMode2::attach].
+    pub fn detach(&mut self, link_id: LircLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 
     /// Queries the lirc device for attached programs.
@@ -91,60 +98,50 @@ impl LircMode2 {
 
         Ok(prog_fds
             .into_iter()
-            .map(|prog_fd| LircLink {
-                prog_fd: Some(prog_fd),
-                target_fd: Some(unsafe { dup(target_fd.as_raw_fd()) }),
-            })
+            .map(|prog_fd| LircLink::new(prog_fd, target_fd.as_raw_fd()))
             .collect())
     }
 }
 
+/// The type returned by [LircMode2::attach]. Can be passed to [LircMode2::detach].
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct LircLinkId(RawFd, RawFd);
+
 #[derive(Debug)]
 pub struct LircLink {
-    prog_fd: Option<RawFd>,
-    target_fd: Option<RawFd>,
+    prog_fd: RawFd,
+    target_fd: RawFd,
 }
 
 impl LircLink {
     pub(crate) fn new(prog_fd: RawFd, target_fd: RawFd) -> LircLink {
         LircLink {
-            prog_fd: Some(prog_fd),
-            target_fd: Some(unsafe { dup(target_fd) }),
+            prog_fd,
+            target_fd: unsafe { dup(target_fd) },
         }
     }
 
     pub fn info(&self) -> Result<ProgramInfo, ProgramError> {
-        if let Some(fd) = self.prog_fd {
-            match bpf_obj_get_info_by_fd(fd) {
-                Ok(info) => Ok(ProgramInfo(info)),
-                Err(io_error) => Err(ProgramError::SyscallError {
-                    call: "bpf_obj_get_info_by_fd".to_owned(),
-                    io_error,
-                }),
-            }
-        } else {
-            Err(ProgramError::AlreadyDetached)
+        match bpf_obj_get_info_by_fd(self.prog_fd) {
+            Ok(info) => Ok(ProgramInfo(info)),
+            Err(io_error) => Err(ProgramError::SyscallError {
+                call: "bpf_obj_get_info_by_fd".to_owned(),
+                io_error,
+            }),
         }
     }
 }
 
 impl Link for LircLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if let Some(prog_fd) = self.prog_fd.take() {
-            let target_fd = self.target_fd.take().unwrap();
-            let _ = bpf_prog_detach(prog_fd, target_fd, BPF_LIRC_MODE2);
-            unsafe { close(target_fd) };
-            Ok(())
-        } else {
-            Err(ProgramError::AlreadyDetached)
-        }
+    type Id = LircLinkId;
+
+    fn id(&self) -> Self::Id {
+        LircLinkId(self.prog_fd, self.target_fd)
     }
-}
 
-impl Drop for LircLink {
-    fn drop(&mut self) {
-        if let Some(target_fd) = self.target_fd.take() {
-            unsafe { close(target_fd) };
-        }
+    fn detach(self) -> Result<(), ProgramError> {
+        let _ = bpf_prog_detach(self.prog_fd, self.target_fd, BPF_LIRC_MODE2);
+        unsafe { close(self.target_fd) };
+        Ok(())
     }
 }

+ 23 - 5
aya/src/programs/lsm.rs

@@ -2,7 +2,10 @@
 use crate::{
     generated::{bpf_attach_type::BPF_LSM_MAC, bpf_prog_type::BPF_PROG_TYPE_LSM},
     obj::btf::{Btf, BtfKind},
-    programs::{load_program, utils::attach_raw_tracepoint, LinkRef, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
+        ProgramData, ProgramError,
+    },
 };
 
 /// A program that attaches to Linux LSM hooks. Used to implement security policy and
@@ -46,14 +49,12 @@ use crate::{
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_LSM")]
 pub struct Lsm {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<LsmLink>,
 }
 
 impl Lsm {
     /// Loads the program inside the kernel.
     ///
-    /// See also [`Program::load`](crate::programs::Program::load).
-    ///
     /// # Arguments
     ///
     /// * `lsm_hook_name` - full name of the LSM hook that the program should
@@ -67,7 +68,24 @@ impl Lsm {
     }
 
     /// Attaches the program.
-    pub fn attach(&mut self) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [Lsm::detach].
+    pub fn attach(&mut self) -> Result<LsmLinkId, ProgramError> {
         attach_raw_tracepoint(&mut self.data, None)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [Lsm::attach].
+    pub fn detach(&mut self, link_id: LsmLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
+
+define_link_wrapper!(
+    LsmLink,
+    /// The type returned by [Lsm::attach]. Can be passed to [Lsm::detach].
+    LsmLinkId,
+    FdLink,
+    FdLinkId
+);

+ 89 - 183
aya/src/programs/mod.rs

@@ -41,6 +41,7 @@ mod extension;
 mod fentry;
 mod fexit;
 mod kprobe;
+mod links;
 mod lirc_mode2;
 mod lsm;
 mod perf_attach;
@@ -58,37 +59,36 @@ mod uprobe;
 mod utils;
 mod xdp;
 
-use libc::{close, dup, ENOSPC};
+use libc::ENOSPC;
 use std::{
-    cell::RefCell,
     convert::TryFrom,
     ffi::CString,
     io,
     os::unix::io::{AsRawFd, RawFd},
     path::Path,
-    rc::Rc,
 };
 use thiserror::Error;
 
-pub use cgroup_skb::{CgroupSkb, CgroupSkbAttachType};
-pub use extension::{Extension, ExtensionError};
-pub use fentry::FEntry;
-pub use fexit::FExit;
-pub use kprobe::{KProbe, KProbeError};
-pub use lirc_mode2::LircMode2;
-pub use lsm::Lsm;
+pub use cgroup_skb::{CgroupSkb, CgroupSkbAttachType, CgroupSkbLinkId};
+pub use extension::{Extension, ExtensionError, ExtensionLinkId};
+pub use fentry::{FEntry, FEntryLinkId};
+pub use fexit::{FExit, FExitLinkId};
+pub use kprobe::{KProbe, KProbeError, KProbeLinkId};
+use links::*;
+pub use lirc_mode2::{LircLinkId, LircMode2};
+pub use lsm::{Lsm, LsmLinkId};
 use perf_attach::*;
 pub use perf_event::{PerfEvent, PerfEventScope, PerfTypeId, SamplePolicy};
 pub use probe::ProbeKind;
-pub use raw_trace_point::RawTracePoint;
-pub use sk_msg::SkMsg;
-pub use sk_skb::{SkSkb, SkSkbKind};
-pub use sock_ops::SockOps;
-pub use socket_filter::{SocketFilter, SocketFilterError};
-pub use tc::{SchedClassifier, TcAttachType, TcError};
-pub use tp_btf::BtfTracePoint;
-pub use trace_point::{TracePoint, TracePointError};
-pub use uprobe::{UProbe, UProbeError};
+pub use raw_trace_point::{RawTracePoint, RawTracePointLinkId};
+pub use sk_msg::{SkMsg, SkMsgLinkId};
+pub use sk_skb::{SkSkb, SkSkbKind, SkSkbLinkId};
+pub use sock_ops::{SockOps, SockOpsLinkId};
+pub use socket_filter::{SocketFilter, SocketFilterError, SocketFilterLinkId};
+pub use tc::{SchedClassifier, SchedClassifierLinkId, TcAttachType, TcError};
+pub use tp_btf::{BtfTracePoint, BtfTracePointLinkId};
+pub use trace_point::{TracePoint, TracePointError, TracePointLinkId};
+pub use uprobe::{UProbe, UProbeError, UProbeLinkId};
 pub use xdp::{Xdp, XdpError, XdpFlags};
 
 use crate::{
@@ -96,7 +96,7 @@ use crate::{
     maps::MapError,
     obj::{self, btf::BtfError, Function, KernelVersion},
     sys::{
-        bpf_get_object, bpf_load_program, bpf_obj_get_info_by_fd, bpf_pin_object, bpf_prog_detach,
+        bpf_get_object, bpf_load_program, bpf_obj_get_info_by_fd, bpf_pin_object,
         bpf_prog_get_fd_by_id, bpf_prog_query, retry_with_verifier_logs, BpfLoadProgramAttrs,
     },
     util::VerifierLog,
@@ -113,9 +113,9 @@ pub enum ProgramError {
     #[error("the program is not loaded")]
     NotLoaded,
 
-    /// The program is already detached.
-    #[error("the program was already detached")]
-    AlreadyDetached,
+    /// The program is already attached.
+    #[error("the program was already attached")]
+    AlreadyAttached,
 
     /// The program is not attached.
     #[error("the program is not attached")]
@@ -251,20 +251,6 @@ pub enum Program {
 }
 
 impl Program {
-    /// Loads the program in the kernel.
-    ///
-    /// # Errors
-    ///
-    /// If the load operation fails, the method returns
-    /// [`ProgramError::LoadError`] and the error's `verifier_log` field
-    /// contains the output from the kernel verifier.
-    ///
-    /// If the program is already loaded, [`ProgramError::AlreadyLoaded`] is
-    /// returned.
-    pub fn load(&mut self) -> Result<(), ProgramError> {
-        load_program(self.prog_type(), self.data_mut())
-    }
-
     /// Returns the low level program type.
     pub fn prog_type(&self) -> bpf_prog_type {
         use crate::generated::bpf_prog_type::*;
@@ -292,62 +278,35 @@ impl Program {
 
     /// Pin the program to the provided path
     pub fn pin<P: AsRef<Path>>(&mut self, path: P) -> Result<(), ProgramError> {
-        self.data_mut().pin(path)
-    }
-
-    fn data(&self) -> &ProgramData {
-        match self {
-            Program::KProbe(p) => &p.data,
-            Program::UProbe(p) => &p.data,
-            Program::TracePoint(p) => &p.data,
-            Program::SocketFilter(p) => &p.data,
-            Program::Xdp(p) => &p.data,
-            Program::SkMsg(p) => &p.data,
-            Program::SkSkb(p) => &p.data,
-            Program::SockOps(p) => &p.data,
-            Program::SchedClassifier(p) => &p.data,
-            Program::CgroupSkb(p) => &p.data,
-            Program::LircMode2(p) => &p.data,
-            Program::PerfEvent(p) => &p.data,
-            Program::RawTracePoint(p) => &p.data,
-            Program::Lsm(p) => &p.data,
-            Program::BtfTracePoint(p) => &p.data,
-            Program::FEntry(p) => &p.data,
-            Program::FExit(p) => &p.data,
-            Program::Extension(p) => &p.data,
-        }
-    }
-
-    fn data_mut(&mut self) -> &mut ProgramData {
         match self {
-            Program::KProbe(p) => &mut p.data,
-            Program::UProbe(p) => &mut p.data,
-            Program::TracePoint(p) => &mut p.data,
-            Program::SocketFilter(p) => &mut p.data,
-            Program::Xdp(p) => &mut p.data,
-            Program::SkMsg(p) => &mut p.data,
-            Program::SkSkb(p) => &mut p.data,
-            Program::SockOps(p) => &mut p.data,
-            Program::SchedClassifier(p) => &mut p.data,
-            Program::CgroupSkb(p) => &mut p.data,
-            Program::LircMode2(p) => &mut p.data,
-            Program::PerfEvent(p) => &mut p.data,
-            Program::RawTracePoint(p) => &mut p.data,
-            Program::Lsm(p) => &mut p.data,
-            Program::BtfTracePoint(p) => &mut p.data,
-            Program::FEntry(p) => &mut p.data,
-            Program::FExit(p) => &mut p.data,
-            Program::Extension(p) => &mut p.data,
+            Program::KProbe(p) => p.data.pin(path),
+            Program::UProbe(p) => p.data.pin(path),
+            Program::TracePoint(p) => p.data.pin(path),
+            Program::SocketFilter(p) => p.data.pin(path),
+            Program::Xdp(p) => p.data.pin(path),
+            Program::SkMsg(p) => p.data.pin(path),
+            Program::SkSkb(p) => p.data.pin(path),
+            Program::SockOps(p) => p.data.pin(path),
+            Program::SchedClassifier(p) => p.data.pin(path),
+            Program::CgroupSkb(p) => p.data.pin(path),
+            Program::LircMode2(p) => p.data.pin(path),
+            Program::PerfEvent(p) => p.data.pin(path),
+            Program::RawTracePoint(p) => p.data.pin(path),
+            Program::Lsm(p) => p.data.pin(path),
+            Program::BtfTracePoint(p) => p.data.pin(path),
+            Program::FEntry(p) => p.data.pin(path),
+            Program::FExit(p) => p.data.pin(path),
+            Program::Extension(p) => p.data.pin(path),
         }
     }
 }
 
-#[derive(Debug, Clone)]
-pub(crate) struct ProgramData {
+#[derive(Debug)]
+pub(crate) struct ProgramData<T: Link> {
     pub(crate) name: Option<String>,
     pub(crate) obj: obj::Program,
     pub(crate) fd: Option<RawFd>,
-    pub(crate) links: Vec<Rc<RefCell<dyn Link>>>,
+    pub(crate) links: LinkMap<T>,
     pub(crate) expected_attach_type: Option<bpf_attach_type>,
     pub(crate) attach_btf_obj_fd: Option<u32>,
     pub(crate) attach_btf_id: Option<u32>,
@@ -355,15 +314,29 @@ pub(crate) struct ProgramData {
     pub(crate) btf_fd: Option<RawFd>,
 }
 
-impl ProgramData {
-    fn fd_or_err(&self) -> Result<RawFd, ProgramError> {
-        self.fd.ok_or(ProgramError::NotLoaded)
+impl<T: Link> ProgramData<T> {
+    pub(crate) fn new(
+        name: Option<String>,
+        obj: obj::Program,
+        btf_fd: Option<RawFd>,
+    ) -> ProgramData<T> {
+        ProgramData {
+            name,
+            obj,
+            fd: None,
+            links: LinkMap::new(),
+            expected_attach_type: None,
+            attach_btf_obj_fd: None,
+            attach_btf_id: None,
+            attach_prog_fd: None,
+            btf_fd,
+        }
     }
+}
 
-    pub fn link<T: Link + 'static>(&mut self, link: T) -> LinkRef {
-        let link: Rc<RefCell<dyn Link>> = Rc::new(RefCell::new(link));
-        self.links.push(Rc::clone(&link));
-        LinkRef::new(link)
+impl<T: Link> ProgramData<T> {
+    fn fd_or_err(&self) -> Result<RawFd, ProgramError> {
+        self.fd.ok_or(ProgramError::NotLoaded)
     }
 
     pub fn pin<P: AsRef<Path>>(&mut self, path: P) -> Result<(), ProgramError> {
@@ -384,7 +357,10 @@ impl ProgramData {
     }
 }
 
-fn load_program(prog_type: bpf_prog_type, data: &mut ProgramData) -> Result<(), ProgramError> {
+fn load_program<T: Link>(
+    prog_type: bpf_prog_type,
+    data: &mut ProgramData<T>,
+) -> Result<(), ProgramError> {
     let ProgramData { obj, fd, .. } = data;
     if fd.is_some() {
         return Err(ProgramError::AlreadyLoaded);
@@ -500,98 +476,28 @@ pub(crate) fn query<T: AsRawFd>(
     }
 }
 
-/// Detach an attached program
-pub trait Link: std::fmt::Debug {
-    /// detaches an attached program
-    fn detach(&mut self) -> Result<(), ProgramError>;
-}
-
-/// The return type of `program.attach(...)`.
-///
-/// [`LinkRef`] implements the [`Link`] trait and can be used to detach a
-/// program.
-#[derive(Debug)]
-pub struct LinkRef {
-    inner: Rc<RefCell<dyn Link>>,
-}
-
-impl LinkRef {
-    fn new(link: Rc<RefCell<dyn Link>>) -> LinkRef {
-        LinkRef { inner: link }
-    }
-}
-
-impl Link for LinkRef {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        self.inner.borrow_mut().detach()
-    }
-}
-
-#[derive(Debug)]
-pub(crate) struct FdLink {
-    fd: Option<RawFd>,
-}
-
-impl Link for FdLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if let Some(fd) = self.fd.take() {
-            unsafe { close(fd) };
-            Ok(())
-        } else {
-            Err(ProgramError::AlreadyDetached)
-        }
-    }
-}
-
-impl Drop for FdLink {
-    fn drop(&mut self) {
-        let _ = self.detach();
-    }
-}
-
-#[derive(Debug)]
-struct ProgAttachLink {
-    prog_fd: Option<RawFd>,
-    target_fd: Option<RawFd>,
-    attach_type: bpf_attach_type,
-}
-
-impl ProgAttachLink {
-    pub(crate) fn new(
-        prog_fd: RawFd,
-        target_fd: RawFd,
-        attach_type: bpf_attach_type,
-    ) -> ProgAttachLink {
-        ProgAttachLink {
-            prog_fd: Some(prog_fd),
-            target_fd: Some(unsafe { dup(target_fd) }),
-            attach_type,
-        }
-    }
-}
-
-impl Link for ProgAttachLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if let Some(prog_fd) = self.prog_fd.take() {
-            let target_fd = self.target_fd.take().unwrap();
-            let _ = bpf_prog_detach(prog_fd, target_fd, self.attach_type);
-            unsafe { close(target_fd) };
-            Ok(())
-        } else {
-            Err(ProgramError::AlreadyDetached)
-        }
-    }
-}
-
-impl Drop for ProgAttachLink {
-    fn drop(&mut self) {
-        let _ = self.detach();
-    }
-}
-
 impl ProgramFd for Program {
     fn fd(&self) -> Option<RawFd> {
-        self.data().fd
+        match self {
+            Program::KProbe(p) => p.data.fd,
+            Program::UProbe(p) => p.data.fd,
+            Program::TracePoint(p) => p.data.fd,
+            Program::SocketFilter(p) => p.data.fd,
+            Program::Xdp(p) => p.data.fd,
+            Program::SkMsg(p) => p.data.fd,
+            Program::SkSkb(p) => p.data.fd,
+            Program::SockOps(p) => p.data.fd,
+            Program::SchedClassifier(p) => p.data.fd,
+            Program::CgroupSkb(p) => p.data.fd,
+            Program::LircMode2(p) => p.data.fd,
+            Program::PerfEvent(p) => p.data.fd,
+            Program::RawTracePoint(p) => p.data.fd,
+            Program::Lsm(p) => p.data.fd,
+            Program::BtfTracePoint(p) => p.data.fd,
+            Program::FEntry(p) => p.data.fd,
+            Program::FExit(p) => p.data.fd,
+            Program::Extension(p) => p.data.fd,
+        }
     }
 }
 

+ 36 - 33
aya/src/programs/perf_attach.rs

@@ -2,64 +2,64 @@ use libc::close;
 use std::os::unix::io::RawFd;
 
 use crate::{
-    programs::{probe::detach_debug_fs, ProbeKind},
+    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,
 };
 
-use super::{Link, LinkRef, ProgramData, ProgramError};
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct PerfLinkId(RawFd);
 
 #[derive(Debug)]
-struct PerfLink {
-    perf_fd: Option<RawFd>,
+pub(crate) struct PerfLink {
+    perf_fd: RawFd,
     probe_kind: Option<ProbeKind>,
     event_alias: Option<String>,
 }
 
 impl Link for PerfLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if let Some(fd) = self.perf_fd.take() {
-            let _ = perf_event_ioctl(fd, PERF_EVENT_IOC_DISABLE, 0);
-            unsafe { close(fd) };
+    type Id = PerfLinkId;
 
-            if let Some(probe_kind) = self.probe_kind.take() {
-                if let Some(event_alias) = self.event_alias.take() {
-                    let _ = detach_debug_fs(probe_kind, &event_alias);
-                }
-            }
+    fn id(&self) -> Self::Id {
+        PerfLinkId(self.perf_fd)
+    }
 
-            Ok(())
-        } else {
-            Err(ProgramError::AlreadyDetached)
+    fn detach(mut self) -> Result<(), ProgramError> {
+        let _ = perf_event_ioctl(self.perf_fd, PERF_EVENT_IOC_DISABLE, 0);
+        unsafe { close(self.perf_fd) };
+
+        if let Some(probe_kind) = self.probe_kind.take() {
+            if let Some(event_alias) = self.event_alias.take() {
+                let _ = detach_debug_fs(probe_kind, &event_alias);
+            }
         }
-    }
-}
 
-impl Drop for PerfLink {
-    fn drop(&mut self) {
-        let _ = self.detach();
+        Ok(())
     }
 }
 
-pub(crate) fn perf_attach(data: &mut ProgramData, fd: RawFd) -> Result<LinkRef, ProgramError> {
+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_debugfs(
-    data: &mut ProgramData,
+pub(crate) fn perf_attach_debugfs<T: Link + From<PerfLink>>(
+    data: &mut ProgramData<T>,
     fd: RawFd,
     probe_kind: ProbeKind,
     event_alias: String,
-) -> Result<LinkRef, ProgramError> {
+) -> Result<T::Id, ProgramError> {
     perf_attach_either(data, fd, Some(probe_kind), Some(event_alias))
 }
 
-fn perf_attach_either(
-    data: &mut ProgramData,
+fn perf_attach_either<T: Link + From<PerfLink>>(
+    data: &mut ProgramData<T>,
     fd: RawFd,
     probe_kind: Option<ProbeKind>,
     event_alias: Option<String>,
-) -> Result<LinkRef, ProgramError> {
+) -> Result<T::Id, ProgramError> {
     let prog_fd = data.fd_or_err()?;
     perf_event_ioctl(fd, PERF_EVENT_IOC_SET_BPF, prog_fd).map_err(|(_, io_error)| {
         ProgramError::SyscallError {
@@ -74,9 +74,12 @@ fn perf_attach_either(
         }
     })?;
 
-    Ok(data.link(PerfLink {
-        perf_fd: Some(fd),
-        probe_kind,
-        event_alias,
-    }))
+    data.links.insert(
+        PerfLink {
+            perf_fd: fd,
+            probe_kind,
+            event_alias,
+        }
+        .into(),
+    )
 }

+ 26 - 12
aya/src/programs/perf_event.rs

@@ -1,16 +1,23 @@
 //! Perf event programs.
-use crate::{generated::bpf_prog_type::BPF_PROG_TYPE_PERF_EVENT, sys::perf_event_open};
-
-use crate::generated::perf_type_id::{
-    PERF_TYPE_BREAKPOINT, PERF_TYPE_HARDWARE, PERF_TYPE_HW_CACHE, PERF_TYPE_RAW,
-    PERF_TYPE_SOFTWARE, PERF_TYPE_TRACEPOINT,
-};
-
 pub use crate::generated::{
     perf_hw_cache_id, perf_hw_cache_op_id, perf_hw_cache_op_result_id, perf_hw_id, perf_sw_ids,
 };
 
-use super::{load_program, perf_attach, LinkRef, ProgramData, ProgramError};
+use crate::{
+    generated::{
+        bpf_prog_type::BPF_PROG_TYPE_PERF_EVENT,
+        perf_type_id::{
+            PERF_TYPE_BREAKPOINT, PERF_TYPE_HARDWARE, PERF_TYPE_HW_CACHE, PERF_TYPE_RAW,
+            PERF_TYPE_SOFTWARE, PERF_TYPE_TRACEPOINT,
+        },
+    },
+    programs::{
+        load_program, perf_attach,
+        perf_attach::{PerfLink, PerfLinkId},
+        ProgramData, ProgramError,
+    },
+    sys::perf_event_open,
+};
 
 /// The type of perf event
 #[repr(u32)]
@@ -112,13 +119,11 @@ pub enum PerfEventScope {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_PERF_EVENT")]
 pub struct PerfEvent {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<PerfLink>,
 }
 
 impl PerfEvent {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_PERF_EVENT, &mut self.data)
     }
@@ -128,13 +133,15 @@ impl PerfEvent {
     /// The possible values and encoding of the `config` argument depends on the
     /// `perf_type`. See `perf_sw_ids`, `perf_hw_id`, `perf_hw_cache_id`,
     /// `perf_hw_cache_op_id` and `perf_hw_cache_op_result_id`.
+    ///
+    /// The returned value can be used to detach, see [PerfEvent::detach].
     pub fn attach(
         &mut self,
         perf_type: PerfTypeId,
         config: u64,
         scope: PerfEventScope,
         sample_policy: SamplePolicy,
-    ) -> Result<LinkRef, ProgramError> {
+    ) -> Result<PerfLinkId, ProgramError> {
         let (sample_period, sample_frequency) = match sample_policy {
             SamplePolicy::Period(period) => (period, None),
             SamplePolicy::Frequency(frequency) => (0, Some(frequency)),
@@ -163,4 +170,11 @@ impl PerfEvent {
 
         perf_attach(&mut self.data, fd)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [PerfEvent::attach].
+    pub fn detach(&mut self, link_id: PerfLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }

+ 5 - 5
aya/src/programs/probe.rs

@@ -7,8 +7,8 @@ use std::{
 
 use crate::{
     programs::{
-        kprobe::KProbeError, perf_attach, perf_attach_debugfs,
-        trace_point::read_sys_fs_trace_point_id, uprobe::UProbeError, LinkRef, ProgramData,
+        kprobe::KProbeError, perf_attach, perf_attach::PerfLink, perf_attach_debugfs,
+        trace_point::read_sys_fs_trace_point_id, uprobe::UProbeError, Link, ProgramData,
         ProgramError,
     },
     sys::{kernel_version, perf_event_open_probe, perf_event_open_trace_point},
@@ -36,13 +36,13 @@ impl ProbeKind {
     }
 }
 
-pub(crate) fn attach(
-    program_data: &mut ProgramData,
+pub(crate) fn attach<T: Link + From<PerfLink>>(
+    program_data: &mut ProgramData<T>,
     kind: ProbeKind,
     fn_name: &str,
     offset: u64,
     pid: Option<pid_t>,
-) -> Result<LinkRef, ProgramError> {
+) -> Result<T::Id, ProgramError> {
     // https://github.com/torvalds/linux/commit/e12f03d7031a977356e3d7b75a68c2185ff8d155
     // Use debugfs to create probe
     let k_ver = kernel_version().unwrap();

+ 23 - 5
aya/src/programs/raw_trace_point.rs

@@ -3,7 +3,10 @@ use std::ffi::CString;
 
 use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_RAW_TRACEPOINT,
-    programs::{load_program, utils::attach_raw_tracepoint, LinkRef, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
+        ProgramData, ProgramError,
+    },
 };
 
 /// A program that can be attached at a pre-defined kernel trace point, but also
@@ -33,20 +36,35 @@ use crate::{
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_RAW_TRACEPOINT")]
 pub struct RawTracePoint {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<RawTracePointLink>,
 }
 
 impl RawTracePoint {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_RAW_TRACEPOINT, &mut self.data)
     }
 
     /// Attaches the program to the given tracepoint.
-    pub fn attach(&mut self, tp_name: &str) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [RawTracePoint::detach].
+    pub fn attach(&mut self, tp_name: &str) -> Result<RawTracePointLinkId, ProgramError> {
         let tp_name_c = CString::new(tp_name).unwrap();
         attach_raw_tracepoint(&mut self.data, Some(&tp_name_c))
     }
+
+    /// Detaches from a tracepoint.
+    ///
+    /// See [RawTracePoint::attach].
+    pub fn detach(&mut self, link_id: RawTracePointLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
+
+define_link_wrapper!(
+    RawTracePointLink,
+    /// The type returned by [RawTracePoint::attach]. Can be passed to [RawTracePoint::detach].
+    RawTracePointLinkId,
+    FdLink,
+    FdLinkId
+);

+ 28 - 8
aya/src/programs/sk_msg.rs

@@ -1,7 +1,10 @@
 use crate::{
     generated::{bpf_attach_type::BPF_SK_MSG_VERDICT, bpf_prog_type::BPF_PROG_TYPE_SK_MSG},
     maps::sock::SocketMap,
-    programs::{load_program, LinkRef, ProgAttachLink, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, ProgAttachLink, ProgAttachLinkId, ProgramData,
+        ProgramError,
+    },
     sys::bpf_prog_attach,
 };
 
@@ -56,19 +59,19 @@ use crate::{
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_SK_MSG")]
 pub struct SkMsg {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<SkMsgLink>,
 }
 
 impl SkMsg {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_SK_MSG, &mut self.data)
     }
 
     /// Attaches the program to the given sockmap.
-    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [SkMsg::detach].
+    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<SkMsgLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let map_fd = map.fd_or_err()?;
 
@@ -78,8 +81,25 @@ impl SkMsg {
                 io_error,
             }
         })?;
-        Ok(self
-            .data
-            .link(ProgAttachLink::new(prog_fd, map_fd, BPF_SK_MSG_VERDICT)))
+        self.data.links.insert(SkMsgLink(ProgAttachLink::new(
+            prog_fd,
+            map_fd,
+            BPF_SK_MSG_VERDICT,
+        )))
+    }
+
+    /// Detaches the program from a sockmap.
+    ///
+    /// See [SkMsg::attach].
+    pub fn detach(&mut self, link_id: SkMsgLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 }
+
+define_link_wrapper!(
+    SkMsgLink,
+    /// The type returned by [SkMsg::attach]. Can be passed to [SkMsg::detach].
+    SkMsgLinkId,
+    ProgAttachLink,
+    ProgAttachLinkId
+);

+ 26 - 8
aya/src/programs/sk_skb.rs

@@ -4,7 +4,10 @@ use crate::{
         bpf_prog_type::BPF_PROG_TYPE_SK_SKB,
     },
     maps::sock::SocketMap,
-    programs::{load_program, LinkRef, ProgAttachLink, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, ProgAttachLink, ProgAttachLinkId, ProgramData,
+        ProgramError,
+    },
     sys::bpf_prog_attach,
 };
 
@@ -48,20 +51,20 @@ pub enum SkSkbKind {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_SK_SKB")]
 pub struct SkSkb {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<SkSkbLink>,
     pub(crate) kind: SkSkbKind,
 }
 
 impl SkSkb {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_SK_SKB, &mut self.data)
     }
 
     /// Attaches the program to the given socket map.
-    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [SkSkb::detach].
+    pub fn attach(&mut self, map: &dyn SocketMap) -> Result<SkSkbLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let map_fd = map.fd_or_err()?;
 
@@ -75,8 +78,23 @@ impl SkSkb {
                 io_error,
             }
         })?;
-        Ok(self
-            .data
-            .link(ProgAttachLink::new(prog_fd, map_fd, attach_type)))
+        self.data
+            .links
+            .insert(SkSkbLink(ProgAttachLink::new(prog_fd, map_fd, attach_type)))
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [SkSkb::attach].
+    pub fn detach(&mut self, link_id: SkSkbLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 }
+
+define_link_wrapper!(
+    SkSkbLink,
+    /// The type returned by [SkSkb::attach]. Can be passed to [SkSkb::detach].
+    SkSkbLinkId,
+    ProgAttachLink,
+    ProgAttachLinkId
+);

+ 28 - 8
aya/src/programs/sock_ops.rs

@@ -2,7 +2,10 @@ use std::os::unix::io::AsRawFd;
 
 use crate::{
     generated::{bpf_attach_type::BPF_CGROUP_SOCK_OPS, bpf_prog_type::BPF_PROG_TYPE_SOCK_OPS},
-    programs::{load_program, LinkRef, ProgAttachLink, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, ProgAttachLink, ProgAttachLinkId, ProgramData,
+        ProgramError,
+    },
     sys::bpf_prog_attach,
 };
 
@@ -43,19 +46,19 @@ use crate::{
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_SOCK_OPS")]
 pub struct SockOps {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<SockOpsLink>,
 }
 
 impl SockOps {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_SOCK_OPS, &mut self.data)
     }
 
     /// Attaches the program to the given cgroup.
-    pub fn attach<T: AsRawFd>(&mut self, cgroup: T) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [SockOps::detach].
+    pub fn attach<T: AsRawFd>(&mut self, cgroup: T) -> Result<SockOpsLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let cgroup_fd = cgroup.as_raw_fd();
 
@@ -65,8 +68,25 @@ impl SockOps {
                 io_error,
             }
         })?;
-        Ok(self
-            .data
-            .link(ProgAttachLink::new(prog_fd, cgroup_fd, BPF_CGROUP_SOCK_OPS)))
+        self.data.links.insert(SockOpsLink(ProgAttachLink::new(
+            prog_fd,
+            cgroup_fd,
+            BPF_CGROUP_SOCK_OPS,
+        )))
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [SockOps::attach].
+    pub fn detach(&mut self, link_id: SockOpsLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 }
+
+define_link_wrapper!(
+    SockOpsLink,
+    /// The type returned by [SockOps::attach]. Can be passed to [SockOps::detach].
+    SockOpsLinkId,
+    ProgAttachLink,
+    ProgAttachLinkId
+);

+ 34 - 30
aya/src/programs/socket_filter.rs

@@ -7,7 +7,7 @@ use thiserror::Error;
 
 use crate::{
     generated::{bpf_prog_type::BPF_PROG_TYPE_SOCKET_FILTER, SO_ATTACH_BPF, SO_DETACH_BPF},
-    programs::{load_program, Link, LinkRef, ProgramData, ProgramError},
+    programs::{load_program, Link, ProgramData, ProgramError},
 };
 
 /// The type returned when attaching a [`SocketFilter`] fails.
@@ -60,19 +60,19 @@ pub enum SocketFilterError {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_SOCKET_FILTER")]
 pub struct SocketFilter {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<SocketFilterLink>,
 }
 
 impl SocketFilter {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_SOCKET_FILTER, &mut self.data)
     }
 
     /// Attaches the filter on the given socket.
-    pub fn attach<T: AsRawFd>(&mut self, socket: T) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach from the socket, see [SocketFilter::detach].
+    pub fn attach<T: AsRawFd>(&mut self, socket: T) -> Result<SocketFilterLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let socket = socket.as_raw_fd();
 
@@ -92,40 +92,44 @@ impl SocketFilter {
             .into());
         }
 
-        Ok(self.data.link(SocketFilterLink {
-            socket,
-            prog_fd: Some(prog_fd),
-        }))
+        self.data.links.insert(SocketFilterLink { socket, prog_fd })
+    }
+
+    /// Detaches the program.
+    ///
+    /// See [SocketFilter::attach].
+    pub fn detach(&mut self, link_id: SocketFilterLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 }
 
+/// The type returned by [SocketFilter::attach]. Can be passed to [SocketFilter::detach].
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub struct SocketFilterLinkId(RawFd, RawFd);
+
 #[derive(Debug)]
-struct SocketFilterLink {
+pub(crate) struct SocketFilterLink {
     socket: RawFd,
-    prog_fd: Option<RawFd>,
+    prog_fd: RawFd,
 }
 
 impl Link for SocketFilterLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if let Some(fd) = self.prog_fd.take() {
-            unsafe {
-                setsockopt(
-                    self.socket,
-                    SOL_SOCKET,
-                    SO_DETACH_BPF as i32,
-                    &fd as *const _ as *const _,
-                    mem::size_of::<RawFd>() as u32,
-                );
-            }
-            Ok(())
-        } else {
-            Err(ProgramError::AlreadyDetached)
-        }
+    type Id = SocketFilterLinkId;
+
+    fn id(&self) -> Self::Id {
+        SocketFilterLinkId(self.socket, self.prog_fd)
     }
-}
 
-impl Drop for SocketFilterLink {
-    fn drop(&mut self) {
-        let _ = self.detach();
+    fn detach(self) -> Result<(), ProgramError> {
+        unsafe {
+            setsockopt(
+                self.socket,
+                SOL_SOCKET,
+                SO_DETACH_BPF as i32,
+                &self.prog_fd as *const _ as *const _,
+                mem::size_of::<RawFd>() as u32,
+            );
+        }
+        Ok(())
     }
 }

+ 40 - 29
aya/src/programs/tc.rs

@@ -4,14 +4,13 @@ use thiserror::Error;
 use std::{
     ffi::{CStr, CString},
     io,
-    os::unix::io::RawFd,
 };
 
 use crate::{
     generated::{
         bpf_prog_type::BPF_PROG_TYPE_SCHED_CLS, TC_H_CLSACT, TC_H_MIN_EGRESS, TC_H_MIN_INGRESS,
     },
-    programs::{load_program, Link, LinkRef, ProgramData, ProgramError},
+    programs::{define_link_wrapper, load_program, Link, ProgramData, ProgramError},
     sys::{
         netlink_find_filter_with_name, netlink_qdisc_add_clsact, netlink_qdisc_attach,
         netlink_qdisc_detach,
@@ -20,7 +19,7 @@ use crate::{
 };
 
 /// Traffic control attach type.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
 pub enum TcAttachType {
     /// Attach to ingress.
     Ingress,
@@ -72,7 +71,7 @@ pub enum TcAttachType {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_SCHED_CLS")]
 pub struct SchedClassifier {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<SchedClassifierLink>,
     pub(crate) name: Box<CStr>,
 }
 
@@ -91,14 +90,6 @@ pub enum TcError {
     AlreadyAttached,
 }
 
-#[derive(Debug)]
-struct TcLink {
-    if_index: i32,
-    attach_type: TcAttachType,
-    prog_fd: Option<RawFd>,
-    priority: u32,
-}
-
 impl TcAttachType {
     pub(crate) fn parent(&self) -> u32 {
         match self {
@@ -111,14 +102,14 @@ impl TcAttachType {
 
 impl SchedClassifier {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_SCHED_CLS, &mut self.data)
     }
 
     /// Attaches the program to the given `interface`.
     ///
+    /// The returned value can be used to detach, see [SchedClassifier::detach].
+    ///
     /// # Errors
     ///
     /// [`TcError::NetlinkError`] is returned if attaching fails. A common cause
@@ -129,7 +120,7 @@ impl SchedClassifier {
         &mut self,
         interface: &str,
         attach_type: TcAttachType,
-    ) -> Result<LinkRef, ProgramError> {
+    ) -> Result<SchedClassifierLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let if_index = ifindex_from_ifname(interface)
             .map_err(|io_error| TcError::NetlinkError { io_error })?;
@@ -137,33 +128,53 @@ impl SchedClassifier {
             unsafe { netlink_qdisc_attach(if_index as i32, &attach_type, prog_fd, &self.name) }
                 .map_err(|io_error| TcError::NetlinkError { io_error })?;
 
-        Ok(self.data.link(TcLink {
+        self.data.links.insert(SchedClassifierLink(TcLink {
             if_index: if_index as i32,
             attach_type,
-            prog_fd: Some(prog_fd),
             priority,
         }))
     }
-}
 
-impl Drop for TcLink {
-    fn drop(&mut self) {
-        let _ = self.detach();
+    /// Detaches the program.
+    ///
+    /// See [SchedClassifier::attach].
+    pub fn detach(&mut self, link_id: SchedClassifierLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
     }
 }
 
+#[derive(Debug, Hash, Eq, PartialEq)]
+pub(crate) struct TcLinkId(i32, TcAttachType, u32);
+
+#[derive(Debug)]
+struct TcLink {
+    if_index: i32,
+    attach_type: TcAttachType,
+    priority: u32,
+}
+
 impl Link for TcLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if self.prog_fd.take().is_some() {
-            unsafe { netlink_qdisc_detach(self.if_index, &self.attach_type, self.priority) }
-                .map_err(|io_error| TcError::NetlinkError { io_error })?;
-            Ok(())
-        } else {
-            Err(ProgramError::AlreadyDetached)
-        }
+    type Id = TcLinkId;
+
+    fn id(&self) -> Self::Id {
+        TcLinkId(self.if_index, self.attach_type, self.priority)
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        unsafe { netlink_qdisc_detach(self.if_index, &self.attach_type, self.priority) }
+            .map_err(|io_error| TcError::NetlinkError { io_error })?;
+        Ok(())
     }
 }
 
+define_link_wrapper!(
+    SchedClassifierLink,
+    /// The type returned by [SchedClassifier::attach]. Can be passed to [SchedClassifier::detach].
+    SchedClassifierLinkId,
+    TcLink,
+    TcLinkId
+);
+
 /// Add the `clasct` qdisc to the given interface.
 ///
 /// The `clsact` qdisc must be added to an interface before [`SchedClassifier`]

+ 23 - 5
aya/src/programs/tp_btf.rs

@@ -2,7 +2,10 @@
 use crate::{
     generated::{bpf_attach_type::BPF_TRACE_RAW_TP, bpf_prog_type::BPF_PROG_TYPE_TRACING},
     obj::btf::{Btf, BtfKind},
-    programs::{load_program, utils::attach_raw_tracepoint, LinkRef, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
+        ProgramData, ProgramError,
+    },
 };
 
 /// Marks a function as a [BTF-enabled raw tracepoint][1] eBPF program that can be attached at
@@ -44,14 +47,12 @@ use crate::{
 #[doc(alias = "BPF_TRACE_RAW_TP")]
 #[doc(alias = "BPF_PROG_TYPE_TRACING")]
 pub struct BtfTracePoint {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<BtfTracePointLink>,
 }
 
 impl BtfTracePoint {
     /// Loads the program inside the kernel.
     ///
-    /// See also [`Program::load`](crate::programs::Program::load).
-    ///
     /// # Arguments
     ///
     /// * `tracepoint` - full name of the tracepoint that we should attach to
@@ -65,7 +66,24 @@ impl BtfTracePoint {
     }
 
     /// Attaches the program.
-    pub fn attach(&mut self) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [BtfTracePoint::detach].
+    pub fn attach(&mut self) -> Result<BtfTracePointLinkId, ProgramError> {
         attach_raw_tracepoint(&mut self.data, None)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [BtfTracePoint::attach].
+    pub fn detach(&mut self, link_id: BtfTracePointLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
+
+define_link_wrapper!(
+    BtfTracePointLink,
+    /// The type returned by [BtfTracePoint::attach]. Can be passed to [BtfTracePoint::detach].
+    BtfTracePointLinkId,
+    FdLink,
+    FdLinkId
+);

+ 28 - 7
aya/src/programs/trace_point.rs

@@ -1,9 +1,15 @@
 use std::{fs, io};
 use thiserror::Error;
 
-use crate::{generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT, sys::perf_event_open_trace_point};
-
-use super::{load_program, perf_attach, LinkRef, ProgramData, ProgramError};
+use crate::{
+    generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT,
+    programs::{
+        define_link_wrapper, load_program,
+        perf_attach::{perf_attach, PerfLink, PerfLinkId},
+        ProgramData, ProgramError,
+    },
+    sys::perf_event_open_trace_point,
+};
 
 /// The type returned when attaching a [`TracePoint`] fails.
 #[derive(Debug, Error)]
@@ -55,13 +61,11 @@ pub enum TracePointError {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_TRACEPOINT")]
 pub struct TracePoint {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<TracePointLink>,
 }
 
 impl TracePoint {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_TRACEPOINT, &mut self.data)
     }
@@ -70,7 +74,9 @@ impl TracePoint {
     ///
     /// For a list of the available event categories and names, see
     /// `/sys/kernel/debug/tracing/events`.
-    pub fn attach(&mut self, category: &str, name: &str) -> Result<LinkRef, ProgramError> {
+    ///
+    /// The returned value can be used to detach, see [TracePoint::detach].
+    pub fn attach(&mut self, category: &str, name: &str) -> Result<TracePointLinkId, ProgramError> {
         let id = read_sys_fs_trace_point_id(category, name)?;
         let fd = perf_event_open_trace_point(id, None).map_err(|(_code, io_error)| {
             ProgramError::SyscallError {
@@ -81,8 +87,23 @@ impl TracePoint {
 
         perf_attach(&mut self.data, fd)
     }
+
+    /// Detaches from a trace point.
+    ///
+    /// See [TracePoint::attach].
+    pub fn detach(&mut self, link_id: TracePointLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
 
+define_link_wrapper!(
+    TracePointLink,
+    /// The type returned by [TracePoint::attach]. Can be passed to [TracePoint::detach].
+    TracePointLinkId,
+    PerfLink,
+    PerfLinkId
+);
+
 pub(crate) fn read_sys_fs_trace_point_id(
     category: &str,
     name: &str,

+ 22 - 7
aya/src/programs/uprobe.rs

@@ -16,9 +16,10 @@ use thiserror::Error;
 use crate::{
     generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE,
     programs::{
-        load_program,
+        define_link_wrapper, load_program,
+        perf_attach::{PerfLink, PerfLinkId},
         probe::{attach, ProbeKind},
-        LinkRef, ProgramData, ProgramError,
+        ProgramData, ProgramError,
     },
 };
 
@@ -40,14 +41,12 @@ const LD_SO_CACHE_HEADER: &str = "glibc-ld.so.cache1.1";
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_KPROBE")]
 pub struct UProbe {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<UProbeLink>,
     pub(crate) kind: ProbeKind,
 }
 
 impl UProbe {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         load_program(BPF_PROG_TYPE_KPROBE, &mut self.data)
     }
@@ -69,16 +68,17 @@ impl UProbe {
     /// a library name (eg: `"libc"`).
     ///
     /// If the program is an `uprobe`, it is attached to the *start* address of the target
-    /// function.  Instead if the program is a `kretprobe`, it is attached to the return address of
+    /// function.  Instead if the program is a `uretprobe`, it is attached to the return address of
     /// the target function.
     ///
+    /// The returned value can be used to detach, see [UProbe::detach].
     pub fn attach<T: AsRef<Path>>(
         &mut self,
         fn_name: Option<&str>,
         offset: u64,
         target: T,
         pid: Option<pid_t>,
-    ) -> Result<LinkRef, ProgramError> {
+    ) -> Result<UProbeLinkId, ProgramError> {
         let target = target.as_ref();
         let target_str = &*target.as_os_str().to_string_lossy();
 
@@ -121,8 +121,23 @@ impl UProbe {
 
         attach(&mut self.data, self.kind, &path, sym_offset + offset, pid)
     }
+
+    /// Detaches the program.
+    ///
+    /// See [UProbe::attach].
+    pub fn detach(&mut self, link_id: UProbeLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
 
+define_link_wrapper!(
+    UProbeLink,
+    /// The type returned by [UProbe::attach]. Can be passed to [UProbe::detach].
+    UProbeLinkId,
+    PerfLink,
+    PerfLinkId
+);
+
 /// The type returned when attaching an [`UProbe`] fails.
 #[derive(Debug, Error)]
 pub enum UProbeError {

+ 5 - 5
aya/src/programs/utils.rs

@@ -2,15 +2,15 @@
 use std::{ffi::CStr, os::unix::io::RawFd};
 
 use crate::{
-    programs::{FdLink, LinkRef, ProgramData, ProgramError},
+    programs::{FdLink, Link, ProgramData, ProgramError},
     sys::bpf_raw_tracepoint_open,
 };
 
 /// Attaches the program to a raw tracepoint.
-pub(crate) fn attach_raw_tracepoint(
-    program_data: &mut ProgramData,
+pub(crate) fn attach_raw_tracepoint<T: Link + From<FdLink>>(
+    program_data: &mut ProgramData<T>,
     tp_name: Option<&CStr>,
-) -> Result<LinkRef, ProgramError> {
+) -> Result<T::Id, ProgramError> {
     let prog_fd = program_data.fd_or_err()?;
 
     let pfd = bpf_raw_tracepoint_open(tp_name, prog_fd).map_err(|(_code, io_error)| {
@@ -20,5 +20,5 @@ pub(crate) fn attach_raw_tracepoint(
         }
     })? as RawFd;
 
-    Ok(program_data.link(FdLink { fd: Some(pfd) }))
+    program_data.links.insert(FdLink::new(pfd).into())
 }

+ 60 - 34
aya/src/programs/xdp.rs

@@ -1,6 +1,6 @@
 use bitflags;
 use libc::if_nametoindex;
-use std::{ffi::CString, io, os::unix::io::RawFd};
+use std::{ffi::CString, hash::Hash, io, os::unix::io::RawFd};
 use thiserror::Error;
 
 use crate::{
@@ -10,7 +10,7 @@ use crate::{
         XDP_FLAGS_DRV_MODE, XDP_FLAGS_HW_MODE, XDP_FLAGS_REPLACE, XDP_FLAGS_SKB_MODE,
         XDP_FLAGS_UPDATE_IF_NOEXIST,
     },
-    programs::{load_program, FdLink, Link, LinkRef, ProgramData, ProgramError},
+    programs::{define_link_wrapper, load_program, FdLink, Link, ProgramData, ProgramError},
     sys::{bpf_link_create, kernel_version, netlink_set_xdp_fd},
 };
 
@@ -68,13 +68,11 @@ bitflags! {
 #[derive(Debug)]
 #[doc(alias = "BPF_PROG_TYPE_XDP")]
 pub struct Xdp {
-    pub(crate) data: ProgramData,
+    pub(crate) data: ProgramData<XdpLink>,
 }
 
 impl Xdp {
     /// Loads the program inside the kernel.
-    ///
-    /// See also [`Program::load`](crate::programs::Program::load).
     pub fn load(&mut self) -> Result<(), ProgramError> {
         self.data.expected_attach_type = Some(bpf_attach_type::BPF_XDP);
         load_program(BPF_PROG_TYPE_XDP, &mut self.data)
@@ -82,6 +80,8 @@ impl Xdp {
 
     /// Attaches the program to the given `interface`.
     ///
+    /// The returned value can be used to detach, see [Xdp::detach].
+    ///
     /// # Errors
     ///
     /// If the given `interface` does not exist
@@ -91,7 +91,7 @@ impl Xdp {
     /// kernels `>= 5.9.0`, and instead
     /// [`XdpError::NetlinkError`] is returned for older
     /// kernels.
-    pub fn attach(&mut self, interface: &str, flags: XdpFlags) -> Result<LinkRef, ProgramError> {
+    pub fn attach(&mut self, interface: &str, flags: XdpFlags) -> Result<XdpLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
         let c_interface = CString::new(interface).unwrap();
         let if_index = unsafe { if_nametoindex(c_interface.as_ptr()) } as RawFd;
@@ -109,63 +109,89 @@ impl Xdp {
                     io_error,
                 },
             )? as RawFd;
-            Ok(self
-                .data
-                .link(XdpLink::FdLink(FdLink { fd: Some(link_fd) })))
+            self.data
+                .links
+                .insert(XdpLink(XdpLinkInner::FdLink(FdLink::new(link_fd))))
         } else {
             unsafe { netlink_set_xdp_fd(if_index, prog_fd, None, flags.bits) }
                 .map_err(|io_error| XdpError::NetlinkError { io_error })?;
 
-            Ok(self.data.link(XdpLink::NlLink(NlLink {
+            self.data.links.insert(XdpLink(XdpLinkInner::NlLink(NlLink {
                 if_index,
-                prog_fd: Some(prog_fd),
+                prog_fd,
                 flags,
             })))
         }
     }
+
+    /// Detaches the program.
+    ///
+    /// See [Xdp::attach].
+    pub fn detach(&mut self, link_id: XdpLinkId) -> Result<(), ProgramError> {
+        self.data.links.remove(link_id)
+    }
 }
 
 #[derive(Debug)]
-struct NlLink {
+pub(crate) struct NlLink {
     if_index: i32,
-    prog_fd: Option<RawFd>,
+    prog_fd: RawFd,
     flags: XdpFlags,
 }
 
 impl Link for NlLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
-        if let Some(fd) = self.prog_fd.take() {
-            let k_ver = kernel_version().unwrap();
-            let flags = if k_ver >= (5, 7, 0) {
-                self.flags.bits | XDP_FLAGS_REPLACE
-            } else {
-                self.flags.bits
-            };
-            let _ = unsafe { netlink_set_xdp_fd(self.if_index, -1, Some(fd), flags) };
-            Ok(())
+    type Id = (i32, RawFd);
+
+    fn id(&self) -> Self::Id {
+        (self.if_index, self.prog_fd)
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
+        let k_ver = kernel_version().unwrap();
+        let flags = if k_ver >= (5, 7, 0) {
+            self.flags.bits | XDP_FLAGS_REPLACE
         } else {
-            Err(ProgramError::AlreadyDetached)
-        }
+            self.flags.bits
+        };
+        let _ = unsafe { netlink_set_xdp_fd(self.if_index, -1, Some(self.prog_fd), flags) };
+        Ok(())
     }
 }
 
-impl Drop for NlLink {
-    fn drop(&mut self) {
-        let _ = self.detach();
-    }
+#[derive(Debug, Hash, Eq, PartialEq)]
+enum XdpLinkIdInner {
+    FdLinkId(<FdLink as Link>::Id),
+    NlLinkId(<NlLink as Link>::Id),
 }
 
 #[derive(Debug)]
-enum XdpLink {
+enum XdpLinkInner {
     FdLink(FdLink),
     NlLink(NlLink),
 }
 
-impl Link for XdpLink {
-    fn detach(&mut self) -> Result<(), ProgramError> {
+impl Link for XdpLinkInner {
+    type Id = XdpLinkIdInner;
+
+    fn id(&self) -> Self::Id {
+        match self {
+            XdpLinkInner::FdLink(link) => XdpLinkIdInner::FdLinkId(link.id()),
+            XdpLinkInner::NlLink(link) => XdpLinkIdInner::NlLinkId(link.id()),
+        }
+    }
+
+    fn detach(self) -> Result<(), ProgramError> {
         match self {
-            XdpLink::FdLink(link) => link.detach(),
-            XdpLink::NlLink(link) => link.detach(),
+            XdpLinkInner::FdLink(link) => link.detach(),
+            XdpLinkInner::NlLink(link) => link.detach(),
         }
     }
 }
+
+define_link_wrapper!(
+    XdpLink,
+    /// The type returned by [Xdp::attach]. Can be passed to [Xdp::detach].
+    XdpLinkId,
+    XdpLinkInner,
+    XdpLinkIdInner
+);

+ 1 - 3
aya/src/sys/bpf.rs

@@ -21,13 +21,11 @@ use crate::{
     },
     maps::PerCpuValues,
     obj::btf::{FuncSecInfo, LineSecInfo},
-    sys::{kernel_version, SysResult},
+    sys::{kernel_version, syscall, SysResult, Syscall},
     util::VerifierLog,
     Pod, BPF_OBJ_NAME_LEN,
 };
 
-use super::{syscall, Syscall};
-
 pub(crate) fn bpf_create_map(name: &CStr, def: &bpf_map_def) -> SysResult {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };