aya: Implement forget_link

Fixes #51

Signed-off-by: Dave Tucker <dave@dtucker.co.uk>

aya/src/programs/cgroup_skb.rs

@@ -9,7 +9,8 @@ use crate::{
         bpf_prog_type::BPF_PROG_TYPE_CGROUP_SKB,
     },
     programs::{
-        define_link_wrapper, load_program, FdLink, Link, ProgAttachLink, ProgramData, ProgramError,
+        define_link_wrapper, load_program, FdLink, Link, OwnedLink, ProgAttachLink, ProgramData,
+        ProgramError,
     },
     sys::{bpf_link_create, bpf_prog_attach, kernel_version},
 };
@@ -116,6 +117,17 @@ impl CgroupSkb {
         }
     }
 
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: CgroupSkbLinkId,
+    ) -> Result<OwnedLink<CgroupSkbLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
+
     /// Detaches the program.
     ///
     /// See [CgroupSkb::attach].
@@ -155,6 +167,7 @@ impl Link for CgroupSkbLinkInner {
 }
 
 define_link_wrapper!(
+    /// The link used by [CgroupSkb] programs.
     CgroupSkbLink,
     /// The type returned by [CgroupSkb::attach]. Can be passed to [CgroupSkb::detach].
     CgroupSkbLinkId,

aya/src/programs/extension.rs

@@ -6,7 +6,9 @@ use object::Endianness;
 use crate::{
     generated::{bpf_attach_type::BPF_CGROUP_INET_INGRESS, bpf_prog_type::BPF_PROG_TYPE_EXT},
     obj::btf::BtfKind,
-    programs::{define_link_wrapper, load_program, FdLink, FdLinkId, ProgramData, ProgramError},
+    programs::{
+        define_link_wrapper, load_program, FdLink, FdLinkId, OwnedLink, ProgramData, ProgramError,
+    },
     sys::{self, bpf_link_create},
     Btf,
 };
@@ -146,9 +148,21 @@ impl Extension {
     pub fn detach(&mut self, link_id: ExtensionLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: ExtensionLinkId,
+    ) -> Result<OwnedLink<ExtensionLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
 }
 
 define_link_wrapper!(
+    /// The link used by [Extension] programs.
     ExtensionLink,
     /// The type returned by [Extension::attach]. Can be passed to [Extension::detach].
     ExtensionLinkId,

aya/src/programs/fentry.rs

@@ -5,7 +5,7 @@ use crate::{
     obj::btf::{Btf, BtfKind},
     programs::{
         define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
-        ProgramData, ProgramError,
+        OwnedLink, ProgramData, ProgramError,
     },
 };
 
@@ -75,9 +75,21 @@ impl FEntry {
     pub fn detach(&mut self, link_id: FEntryLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: FEntryLinkId,
+    ) -> Result<OwnedLink<FEntryLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
 }
 
 define_link_wrapper!(
+    /// The link used by [FEntry] programs.
     FEntryLink,
     /// The type returned by [FEntry::attach]. Can be passed to [FEntry::detach].
     FEntryLinkId,

aya/src/programs/fexit.rs

@@ -5,7 +5,7 @@ use crate::{
     obj::btf::{Btf, BtfKind},
     programs::{
         define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
-        ProgramData, ProgramError,
+        OwnedLink, ProgramData, ProgramError,
     },
 };
 
@@ -75,9 +75,21 @@ impl FExit {
     pub fn detach(&mut self, link_id: FExitLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: FExitLinkId,
+    ) -> Result<OwnedLink<FExitLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
 }
 
 define_link_wrapper!(
+    /// The link used by [FExit] programs.
     FExitLink,
     /// The type returned by [FExit::attach]. Can be passed to [FExit::detach].
     FExitLinkId,

aya/src/programs/kprobe.rs

@@ -8,7 +8,7 @@ use crate::{
         define_link_wrapper, load_program,
         perf_attach::{PerfLink, PerfLinkId},
         probe::{attach, ProbeKind},
-        ProgramData, ProgramError,
+        OwnedLink, ProgramData, ProgramError,
     },
 };
 
@@ -76,9 +76,21 @@ impl KProbe {
     pub fn detach(&mut self, link_id: KProbeLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: KProbeLinkId,
+    ) -> Result<OwnedLink<KProbeLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
 }
 
 define_link_wrapper!(
+    /// The link used by [KProbe] programs.
     KProbeLink,
     /// The type returned by [KProbe::attach]. Can be passed to [KProbe::detach].
     KProbeLinkId,

aya/src/programs/links.rs

@@ -1,19 +1,53 @@
 use libc::{close, dup};
+
 use std::{
+    borrow::Borrow,
     collections::{hash_map::Entry, HashMap},
+    ops::Deref,
     os::unix::prelude::RawFd,
 };
 
 use crate::{generated::bpf_attach_type, programs::ProgramError, sys::bpf_prog_detach};
 
-pub(crate) trait Link: std::fmt::Debug + 'static {
+/// A Link
+pub trait Link: std::fmt::Debug + 'static {
+    /// Unique Id
     type Id: std::fmt::Debug + std::hash::Hash + Eq + PartialEq;
 
+    /// Returns the link id
     fn id(&self) -> Self::Id;
 
+    /// Detaches the Link
     fn detach(self) -> Result<(), ProgramError>;
 }
 
+/// An owned link that automatically detaches the inner link when dropped.
+pub struct OwnedLink<T: Link> {
+    inner: Option<T>,
+}
+
+impl<T: Link> OwnedLink<T> {
+    pub(crate) fn new(inner: T) -> Self {
+        Self { inner: Some(inner) }
+    }
+}
+
+impl<T: Link> Deref for OwnedLink<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        self.inner.borrow().as_ref().unwrap()
+    }
+}
+
+impl<T: Link> Drop for OwnedLink<T> {
+    fn drop(&mut self) {
+        if let Some(link) = self.inner.take() {
+            link.detach().unwrap();
+        }
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct LinkMap<T: Link> {
     links: HashMap<T::Id, T>,
@@ -43,6 +77,10 @@ impl<T: Link> LinkMap<T> {
             .ok_or(ProgramError::NotAttached)?
             .detach()
     }
+
+    pub(crate) fn forget(&mut self, link_id: T::Id) -> Result<T, ProgramError> {
+        self.links.remove(&link_id).ok_or(ProgramError::NotAttached)
+    }
 }
 
 impl<T: Link> Drop for LinkMap<T> {
@@ -58,7 +96,7 @@ pub(crate) struct FdLinkId(pub(crate) RawFd);
 
 #[derive(Debug)]
 pub(crate) struct FdLink {
-    fd: RawFd,
+    pub(crate) fd: RawFd,
 }
 
 impl FdLink {
@@ -119,13 +157,14 @@ impl Link for ProgAttachLink {
 }
 
 macro_rules! define_link_wrapper {
-    ($wrapper:ident, #[$doc:meta] $wrapper_id:ident, $base:ident, $base_id:ident) => {
-        #[$doc]
+    (#[$doc1:meta] $wrapper:ident, #[$doc2:meta] $wrapper_id:ident, $base:ident, $base_id:ident) => {
+        #[$doc2]
         #[derive(Debug, Hash, Eq, PartialEq)]
         pub struct $wrapper_id($base_id);
 
+        #[$doc1]
         #[derive(Debug)]
-        pub(crate) struct $wrapper($base);
+        pub struct $wrapper($base);
 
         impl crate::programs::Link for $wrapper {
             type Id = $wrapper_id;
@@ -153,7 +192,7 @@ pub(crate) use define_link_wrapper;
 mod tests {
     use std::{cell::RefCell, rc::Rc};
 
-    use crate::programs::ProgramError;
+    use crate::programs::{OwnedLink, ProgramError};
 
     use super::{Link, LinkMap};
 
@@ -257,4 +296,58 @@ mod tests {
         assert!(*l1_detached.borrow() == 1);
         assert!(*l2_detached.borrow() == 1);
     }
+
+    #[test]
+    fn test_owned_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 owned_l1 = {
+            let mut links = LinkMap::new();
+            let id1 = links.insert(l1).unwrap();
+            links.insert(l2).unwrap();
+            // manually forget one link
+            let owned_l1 = links.forget(id1);
+            assert!(*l1_detached.borrow() == 0);
+            assert!(*l2_detached.borrow() == 0);
+            owned_l1.unwrap()
+        };
+
+        // l2 is detached on `Drop`, but l1 is still alive
+        assert!(*l1_detached.borrow() == 0);
+        assert!(*l2_detached.borrow() == 1);
+
+        // manually detach l1
+        assert!(owned_l1.detach().is_ok());
+        assert!(*l1_detached.borrow() == 1);
+        assert!(*l2_detached.borrow() == 1);
+    }
+
+    #[test]
+    fn test_owned_drop() {
+        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 forget one link and wrap in OwnedLink
+            let _ = OwnedLink {
+                inner: Some(links.forget(id1).unwrap()),
+            };
+
+            // OwnedLink was dropped in the statement above
+            assert!(*l1_detached.borrow() == 1);
+            assert!(*l2_detached.borrow() == 0);
+        };
+
+        assert!(*l1_detached.borrow() == 1);
+        assert!(*l2_detached.borrow() == 1);
+    }
 }

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, ProgramData, ProgramError, ProgramInfo},
+    programs::{load_program, query, Link, OwnedLink, ProgramData, ProgramError, ProgramInfo},
     sys::{bpf_obj_get_info_by_fd, bpf_prog_attach, bpf_prog_detach, bpf_prog_get_fd_by_id},
 };
 
@@ -81,6 +81,17 @@ impl LircMode2 {
         self.data.links.remove(link_id)
     }
 
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(
+        &mut self,
+        link_id: LircLinkId,
+    ) -> Result<OwnedLink<LircLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
+
     /// Queries the lirc device for attached programs.
     pub fn query<T: AsRawFd>(target_fd: T) -> Result<Vec<LircLink>, ProgramError> {
         let prog_ids = query(target_fd.as_raw_fd(), BPF_LIRC_MODE2, 0, &mut None)?;
@@ -108,6 +119,7 @@ impl LircMode2 {
 pub struct LircLinkId(RawFd, RawFd);
 
 #[derive(Debug)]
+/// An LircMode2 Link
 pub struct LircLink {
     prog_fd: RawFd,
     target_fd: RawFd,
@@ -121,6 +133,7 @@ impl LircLink {
         }
     }
 
+    /// Get ProgramInfo from this link
     pub fn info(&self) -> Result<ProgramInfo, ProgramError> {
         match bpf_obj_get_info_by_fd(self.prog_fd) {
             Ok(info) => Ok(ProgramInfo(info)),

aya/src/programs/lsm.rs

@@ -4,7 +4,7 @@ use crate::{
     obj::btf::{Btf, BtfKind},
     programs::{
         define_link_wrapper, load_program, utils::attach_raw_tracepoint, FdLink, FdLinkId,
-        ProgramData, ProgramError,
+        OwnedLink, ProgramData, ProgramError,
     },
 };
 
@@ -80,9 +80,18 @@ impl Lsm {
     pub fn detach(&mut self, link_id: LsmLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Takes ownership of the link referenced by the provided link_id.
+    ///
+    /// The link will be detached on `Drop` and the caller is now responsible
+    /// for managing its lifetime.
+    pub fn forget_link(&mut self, link_id: LsmLinkId) -> Result<OwnedLink<LsmLink>, ProgramError> {
+        Ok(OwnedLink::new(self.data.forget_link(link_id)?))
+    }
 }
 
 define_link_wrapper!(
+    /// The link used by [Lsm] programs.
     LsmLink,
     /// The type returned by [Lsm::attach]. Can be passed to [Lsm::detach].
     LsmLinkId,