aya: Implement forget_link

Fixes #51 Signed-off-by: Dave Tucker <dave@dtucker.co.uk>
aya-rs · May 6, 2022 · 45ab0a1 · 45ab0a1
1 parent 4afc5ea
commit 45ab0a1
Show file tree

Hide file tree

Showing 24 changed files with 319 additions and 8 deletions.
diff --git a/aya/src/programs/cgroup_skb.rs b/aya/src/programs/cgroup_skb.rs
@@ -116,6 +116,14 @@ 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<CgroupSkbLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
+
     /// Detaches the program.
     ///
     /// See [CgroupSkb::attach].
@@ -155,6 +163,7 @@ impl Link for CgroupSkbLinkInner {
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [CgroupSkb::forget_link]
     CgroupSkbLink,
     /// The type returned by [CgroupSkb::attach]. Can be passed to [CgroupSkb::detach].
     CgroupSkbLinkId,

diff --git a/aya/src/programs/extension.rs b/aya/src/programs/extension.rs
@@ -146,9 +146,18 @@ 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<ExtensionLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [Extension::forget_link]
     ExtensionLink,
     /// The type returned by [Extension::attach]. Can be passed to [Extension::detach].
     ExtensionLinkId,

diff --git a/aya/src/programs/fentry.rs b/aya/src/programs/fentry.rs
@@ -75,9 +75,18 @@ 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<FEntryLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [FEntry::forget_link]
     FEntryLink,
     /// The type returned by [FEntry::attach]. Can be passed to [FEntry::detach].
     FEntryLinkId,

diff --git a/aya/src/programs/fexit.rs b/aya/src/programs/fexit.rs
@@ -75,9 +75,18 @@ 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<FExitLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [FExit::forget_link]
     FExitLink,
     /// The type returned by [FExit::attach]. Can be passed to [FExit::detach].
     FExitLinkId,

diff --git a/aya/src/programs/kprobe.rs b/aya/src/programs/kprobe.rs
@@ -76,9 +76,18 @@ 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<KProbeLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [KProbe::forget_link]
     KProbeLink,
     /// The type returned by [KProbe::attach]. Can be passed to [KProbe::detach].
     KProbeLinkId,

diff --git a/aya/src/programs/links.rs b/aya/src/programs/links.rs
@@ -43,6 +43,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 +62,7 @@ pub(crate) struct FdLinkId(pub(crate) RawFd);
 
 #[derive(Debug)]
 pub(crate) struct FdLink {
-    fd: RawFd,
+    pub(crate) fd: RawFd,
 }
 
 impl FdLink {
@@ -119,13 +123,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;
@@ -257,4 +262,32 @@ mod tests {
         assert!(*l1_detached.borrow() == 1);
         assert!(*l2_detached.borrow() == 1);
     }
+
+    #[test]
+    fn test_forget() {
+        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);
+    }
 }
diff --git a/aya/src/programs/lirc_mode2.rs b/aya/src/programs/lirc_mode2.rs
@@ -81,6 +81,14 @@ 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<LircLink, ProgramError> {
+        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)?;

diff --git a/aya/src/programs/lsm.rs b/aya/src/programs/lsm.rs
@@ -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<LsmLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [Lsm::forget_link]
     LsmLink,
     /// The type returned by [Lsm::attach]. Can be passed to [Lsm::detach].
     LsmLinkId,

diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs
@@ -355,6 +355,10 @@ impl<T: Link> ProgramData<T> {
         })?;
         Ok(())
     }
+
+    pub(crate) fn forget_link(&mut self, link_id: T::Id) -> Result<T, ProgramError> {
+        self.links.forget(link_id)
+    }
 }
 
 fn load_program<T: Link>(

diff --git a/aya/src/programs/perf_attach.rs b/aya/src/programs/perf_attach.rs
@@ -11,7 +11,7 @@ use crate::{
 pub struct PerfLinkId(RawFd);
 
 #[derive(Debug)]
-pub(crate) struct PerfLink {
+pub struct PerfLink {
     perf_fd: RawFd,
     probe_kind: Option<ProbeKind>,
     event_alias: Option<String>,

diff --git a/aya/src/programs/perf_event.rs b/aya/src/programs/perf_event.rs
@@ -177,4 +177,12 @@ impl PerfEvent {
     pub fn detach(&mut self, link_id: PerfLinkId) -> 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: PerfLinkId) -> Result<PerfLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
diff --git a/aya/src/programs/raw_trace_point.rs b/aya/src/programs/raw_trace_point.rs
@@ -59,9 +59,21 @@ impl RawTracePoint {
     pub fn detach(&mut self, link_id: RawTracePointLinkId) -> 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: RawTracePointLinkId,
+    ) -> Result<RawTracePointLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [RawTracePoint::forget_link]
     RawTracePointLink,
     /// The type returned by [RawTracePoint::attach]. Can be passed to [RawTracePoint::detach].
     RawTracePointLinkId,

diff --git a/aya/src/programs/sk_msg.rs b/aya/src/programs/sk_msg.rs
@@ -94,9 +94,18 @@ impl SkMsg {
     pub fn detach(&mut self, link_id: SkMsgLinkId) -> 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: SkMsgLinkId) -> Result<SkMsgLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [SkMsg::forget_link]
     SkMsgLink,
     /// The type returned by [SkMsg::attach]. Can be passed to [SkMsg::detach].
     SkMsgLinkId,

diff --git a/aya/src/programs/sk_skb.rs b/aya/src/programs/sk_skb.rs
@@ -89,9 +89,18 @@ impl SkSkb {
     pub fn detach(&mut self, link_id: SkSkbLinkId) -> 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: SkSkbLinkId) -> Result<SkSkbLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [SkSkb::forget_link]
     SkSkbLink,
     /// The type returned by [SkSkb::attach]. Can be passed to [SkSkb::detach].
     SkSkbLinkId,

diff --git a/aya/src/programs/sock_ops.rs b/aya/src/programs/sock_ops.rs
@@ -81,9 +81,18 @@ impl SockOps {
     pub fn detach(&mut self, link_id: SockOpsLinkId) -> 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: SockOpsLinkId) -> Result<SockOpsLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [SockOps::forget_link]
     SockOpsLink,
     /// The type returned by [SockOps::attach]. Can be passed to [SockOps::detach].
     SockOpsLinkId,

diff --git a/aya/src/programs/socket_filter.rs b/aya/src/programs/socket_filter.rs
@@ -101,14 +101,25 @@ impl SocketFilter {
     pub fn detach(&mut self, link_id: SocketFilterLinkId) -> 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: SocketFilterLinkId,
+    ) -> Result<SocketFilterLink, ProgramError> {
+        self.data.forget_link(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)]
-pub(crate) struct SocketFilterLink {
+pub struct SocketFilterLink {
     socket: RawFd,
     prog_fd: RawFd,
 }

diff --git a/aya/src/programs/tc.rs b/aya/src/programs/tc.rs
@@ -141,6 +141,17 @@ impl SchedClassifier {
     pub fn detach(&mut self, link_id: SchedClassifierLinkId) -> 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: SchedClassifierLinkId,
+    ) -> Result<SchedClassifierLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]
@@ -168,6 +179,7 @@ impl Link for TcLink {
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [SchedClassifier::forget_link]
     SchedClassifierLink,
     /// The type returned by [SchedClassifier::attach]. Can be passed to [SchedClassifier::detach].
     SchedClassifierLinkId,

diff --git a/aya/src/programs/tp_btf.rs b/aya/src/programs/tp_btf.rs
@@ -78,9 +78,21 @@ impl BtfTracePoint {
     pub fn detach(&mut self, link_id: BtfTracePointLinkId) -> 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: BtfTracePointLinkId,
+    ) -> Result<BtfTracePointLink, ProgramError> {
+        self.data.forget_link(link_id)
+    }
 }
 
 define_link_wrapper!(
+    /// An owned link, obtained using [BtfTracePoint::forget_link]
     BtfTracePointLink,
     /// The type returned by [BtfTracePoint::attach]. Can be passed to [BtfTracePoint::detach].
     BtfTracePointLinkId,