aya: fix detaching links on drop

https://github.com/aya-rs/aya/pull/366 broke detaching links on drop for
everything but FdLink. This restores detach on drop for all links.

aya/src/programs/cgroup_device.rs

@@ -72,9 +72,9 @@ impl CgroupDevice {
             )? as RawFd;
             self.data
                 .links
-                .insert(CgroupDeviceLink(CgroupDeviceLinkInner::Fd(FdLink::new(
-                    link_fd,
-                ))))
+                .insert(CgroupDeviceLink::new(CgroupDeviceLinkInner::Fd(
+                    FdLink::new(link_fd),
+                )))
         } else {
             bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_DEVICE).map_err(|(_, io_error)| {
                 ProgramError::SyscallError {
@@ -84,7 +84,7 @@ impl CgroupDevice {
             })?;
             self.data
                 .links
-                .insert(CgroupDeviceLink(CgroupDeviceLinkInner::ProgAttach(
+                .insert(CgroupDeviceLink::new(CgroupDeviceLinkInner::ProgAttach(
                     ProgAttachLink::new(prog_fd, cgroup_fd, BPF_CGROUP_DEVICE),
                 )))
         }

aya/src/programs/cgroup_skb.rs

@@ -105,7 +105,9 @@ impl CgroupSkb {
             )? as RawFd;
             self.data
                 .links
-                .insert(CgroupSkbLink(CgroupSkbLinkInner::Fd(FdLink::new(link_fd))))
+                .insert(CgroupSkbLink::new(CgroupSkbLinkInner::Fd(FdLink::new(
+                    link_fd,
+                ))))
         } else {
             bpf_prog_attach(prog_fd, cgroup_fd, attach_type).map_err(|(_, io_error)| {
                 ProgramError::SyscallError {
@@ -116,7 +118,7 @@ impl CgroupSkb {
 
             self.data
                 .links
-                .insert(CgroupSkbLink(CgroupSkbLinkInner::ProgAttach(
+                .insert(CgroupSkbLink::new(CgroupSkbLinkInner::ProgAttach(
                     ProgAttachLink::new(prog_fd, cgroup_fd, attach_type),
                 )))
         }

aya/src/programs/cgroup_sock.rs

@@ -81,7 +81,7 @@ impl CgroupSock {
             )? as RawFd;
             self.data
                 .links
-                .insert(CgroupSockLink(CgroupSockLinkInner::Fd(FdLink::new(
+                .insert(CgroupSockLink::new(CgroupSockLinkInner::Fd(FdLink::new(
                     link_fd,
                 ))))
         } else {
@@ -94,7 +94,7 @@ impl CgroupSock {
 
             self.data
                 .links
-                .insert(CgroupSockLink(CgroupSockLinkInner::ProgAttach(
+                .insert(CgroupSockLink::new(CgroupSockLinkInner::ProgAttach(
                     ProgAttachLink::new(prog_fd, cgroup_fd, attach_type),
                 )))
         }

aya/src/programs/cgroup_sock_addr.rs

@@ -82,7 +82,7 @@ impl CgroupSockAddr {
             )? as RawFd;
             self.data
                 .links
-                .insert(CgroupSockAddrLink(CgroupSockAddrLinkInner::Fd(
+                .insert(CgroupSockAddrLink::new(CgroupSockAddrLinkInner::Fd(
                     FdLink::new(link_fd),
                 )))
         } else {
@@ -93,11 +93,13 @@ impl CgroupSockAddr {
                 }
             })?;
 
-            self.data
-                .links
-                .insert(CgroupSockAddrLink(CgroupSockAddrLinkInner::ProgAttach(
-                    ProgAttachLink::new(prog_fd, cgroup_fd, attach_type),
-                )))
+            self.data.links.insert(CgroupSockAddrLink::new(
+                CgroupSockAddrLinkInner::ProgAttach(ProgAttachLink::new(
+                    prog_fd,
+                    cgroup_fd,
+                    attach_type,
+                )),
+            ))
         }
     }
 

aya/src/programs/cgroup_sockopt.rs

@@ -79,9 +79,9 @@ impl CgroupSockopt {
             )? as RawFd;
             self.data
                 .links
-                .insert(CgroupSockoptLink(CgroupSockoptLinkInner::Fd(FdLink::new(
-                    link_fd,
-                ))))
+                .insert(CgroupSockoptLink::new(CgroupSockoptLinkInner::Fd(
+                    FdLink::new(link_fd),
+                )))
         } else {
             bpf_prog_attach(prog_fd, cgroup_fd, attach_type).map_err(|(_, io_error)| {
                 ProgramError::SyscallError {
@@ -92,7 +92,7 @@ impl CgroupSockopt {
 
             self.data
                 .links
-                .insert(CgroupSockoptLink(CgroupSockoptLinkInner::ProgAttach(
+                .insert(CgroupSockoptLink::new(CgroupSockoptLinkInner::ProgAttach(
                     ProgAttachLink::new(prog_fd, cgroup_fd, attach_type),
                 )))
         }

aya/src/programs/cgroup_sysctl.rs

@@ -74,9 +74,9 @@ impl CgroupSysctl {
             )? as RawFd;
             self.data
                 .links
-                .insert(CgroupSysctlLink(CgroupSysctlLinkInner::Fd(FdLink::new(
-                    link_fd,
-                ))))
+                .insert(CgroupSysctlLink::new(CgroupSysctlLinkInner::Fd(
+                    FdLink::new(link_fd),
+                )))
         } else {
             bpf_prog_attach(prog_fd, cgroup_fd, BPF_CGROUP_SYSCTL).map_err(|(_, io_error)| {
                 ProgramError::SyscallError {
@@ -87,7 +87,7 @@ impl CgroupSysctl {
 
             self.data
                 .links
-                .insert(CgroupSysctlLink(CgroupSysctlLinkInner::ProgAttach(
+                .insert(CgroupSysctlLink::new(CgroupSysctlLinkInner::ProgAttach(
                     ProgAttachLink::new(prog_fd, cgroup_fd, BPF_CGROUP_SYSCTL),
                 )))
         }

aya/src/programs/extension.rs

@@ -95,7 +95,9 @@ impl Extension {
                 call: "bpf_link_create".to_owned(),
                 io_error,
             })? as RawFd;
-        self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
+        self.data
+            .links
+            .insert(ExtensionLink::new(FdLink::new(link_fd)))
     }
 
     /// Attaches the extension to another program.
@@ -123,7 +125,9 @@ impl Extension {
                 call: "bpf_link_create".to_owned(),
                 io_error,
             })? as RawFd;
-        self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
+        self.data
+            .links
+            .insert(ExtensionLink::new(FdLink::new(link_fd)))
     }
 
     /// Detaches the extension.

aya/src/programs/links.rs

@@ -79,17 +79,17 @@ impl<T: Link> Drop for LinkMap<T> {
 
 /// The identifier of an `FdLink`.
 #[derive(Debug, Hash, Eq, PartialEq)]
-pub struct FdLinkId(pub(crate) Option<RawFd>);
+pub struct FdLinkId(pub(crate) RawFd);
 
 /// A file descriptor link.
 #[derive(Debug)]
 pub struct FdLink {
-    pub(crate) fd: Option<RawFd>,
+    pub(crate) fd: RawFd,
 }
 
 impl FdLink {
     pub(crate) fn new(fd: RawFd) -> FdLink {
-        FdLink { fd: Some(fd) }
+        FdLink { fd }
     }
 
     /// Pins the link to a BPF file system.
@@ -121,21 +121,18 @@ impl FdLink {
     /// let pinned_link = fd_link.pin("/sys/fs/bpf/example")?;
     /// # Ok::<(), Error>(())
     /// ```
-    pub fn pin<P: AsRef<Path>>(mut self, path: P) -> Result<PinnedLink, PinError> {
-        let fd = self.fd.take().ok_or_else(|| PinError::NoFd {
-            name: "link".to_string(),
-        })?;
+    pub fn pin<P: AsRef<Path>>(self, path: P) -> Result<PinnedLink, PinError> {
         let path_string =
             CString::new(path.as_ref().to_string_lossy().into_owned()).map_err(|e| {
                 PinError::InvalidPinPath {
                     error: e.to_string(),
                 }
             })?;
-        bpf_pin_object(fd, &path_string).map_err(|(_, io_error)| PinError::SyscallError {
+        bpf_pin_object(self.fd, &path_string).map_err(|(_, io_error)| PinError::SyscallError {
             name: "BPF_OBJ_PIN".to_string(),
             io_error,
         })?;
-        Ok(PinnedLink::new(PathBuf::from(path.as_ref()), fd))
+        Ok(PinnedLink::new(PathBuf::from(path.as_ref()), self))
     }
 }
 
@@ -147,27 +144,27 @@ impl Link for FdLink {
     }
 
     fn detach(self) -> Result<(), ProgramError> {
-        // detach is a noop since it consumes self. once self is consumed,
-        // drop will be triggered and the link will be detached.
+        // detach is a noop since it consumes self. once self is consumed, drop will be triggered
+        // and the link will be detached.
+        //
+        // Other links don't need to do this since they use define_link_wrapper!, but FdLink is a
+        // bit special in that it defines a custom ::new() so it can't use the macro.
         Ok(())
     }
 }
 
-impl Drop for FdLink {
-    fn drop(&mut self) {
-        if let Some(fd) = self.fd.take() {
-            // Safety: libc
-            unsafe { close(fd) };
-        }
-    }
-}
-
 impl From<PinnedLink> for FdLink {
     fn from(p: PinnedLink) -> Self {
         p.inner
     }
 }
 
+impl Drop for FdLink {
+    fn drop(&mut self) {
+        unsafe { close(self.fd) };
+    }
+}
+
 /// A pinned file descriptor link.
 ///
 /// This link has been pinned to the BPF filesystem. On drop, the file descriptor that backs
@@ -180,11 +177,8 @@ pub struct PinnedLink {
 }
 
 impl PinnedLink {
-    fn new(path: PathBuf, fd: RawFd) -> Self {
-        PinnedLink {
-            inner: FdLink::new(fd),
-            path,
-        }
+    fn new(path: PathBuf, link: FdLink) -> Self {
+        PinnedLink { inner: link, path }
     }
 
     /// Creates a [`PinnedLink`] from a valid path on bpffs.
@@ -196,7 +190,10 @@ impl PinnedLink {
                 code,
                 io_error,
             })? as RawFd;
-        Ok(PinnedLink::new(path.as_ref().to_path_buf(), fd))
+        Ok(PinnedLink::new(
+            path.as_ref().to_path_buf(),
+            FdLink::new(fd),
+        ))
     }
 
     /// Removes the pinned link from the filesystem and returns an [`FdLink`].
@@ -254,29 +251,56 @@ macro_rules! define_link_wrapper {
 
         #[$doc1]
         #[derive(Debug)]
-        pub struct $wrapper($base);
+        pub struct $wrapper(Option<$base>);
+
+        #[allow(dead_code)]
+        // allow dead code since currently XDP is the only consumer of inner and
+        // into_inner
+        impl $wrapper {
+            fn new(base: $base) -> $wrapper {
+                $wrapper(Some(base))
+            }
+
+            fn inner(&self) -> &$base {
+                self.0.as_ref().unwrap()
+            }
+
+            fn into_inner(mut self) -> $base {
+                self.0.take().unwrap()
+            }
+        }
+
+        impl Drop for $wrapper {
+            fn drop(&mut self) {
+                use crate::programs::links::Link;
+
+                if let Some(base) = self.0.take() {
+                    let _ = base.detach();
+                }
+            }
+        }
 
         impl crate::programs::Link for $wrapper {
             type Id = $wrapper_id;
 
             fn id(&self) -> Self::Id {
-                $wrapper_id(self.0.id())
+                $wrapper_id(self.0.as_ref().unwrap().id())
             }
 
-            fn detach(self) -> Result<(), ProgramError> {
-                self.0.detach()
+            fn detach(mut self) -> Result<(), ProgramError> {
+                self.0.take().unwrap().detach()
             }
         }
 
         impl From<$base> for $wrapper {
             fn from(b: $base) -> $wrapper {
-                $wrapper(b)
+                $wrapper(Some(b))
             }
         }
 
         impl From<$wrapper> for $base {
-            fn from(w: $wrapper) -> $base {
-                w.0
+            fn from(mut w: $wrapper) -> $base {
+                w.0.take().unwrap()
             }
         }
     };

aya/src/programs/sk_lookup.rs

@@ -70,7 +70,9 @@ impl SkLookup {
                 io_error,
             },
         )? as RawFd;
-        self.data.links.insert(SkLookupLink(FdLink::new(link_fd)))
+        self.data
+            .links
+            .insert(SkLookupLink::new(FdLink::new(link_fd)))
     }
 
     /// Takes ownership of the link referenced by the provided link_id.

aya/src/programs/sk_msg.rs

@@ -88,7 +88,7 @@ impl SkMsg {
                 io_error,
             }
         })?;
-        self.data.links.insert(SkMsgLink(ProgAttachLink::new(
+        self.data.links.insert(SkMsgLink::new(ProgAttachLink::new(
             prog_fd,
             map_fd,
             BPF_SK_MSG_VERDICT,

aya/src/programs/sk_skb.rs

@@ -84,9 +84,11 @@ impl SkSkb {
                 io_error,
             }
         })?;
-        self.data
-            .links
-            .insert(SkSkbLink(ProgAttachLink::new(prog_fd, map_fd, attach_type)))
+        self.data.links.insert(SkSkbLink::new(ProgAttachLink::new(
+            prog_fd,
+            map_fd,
+            attach_type,
+        )))
     }
 
     /// Detaches the program.

aya/src/programs/sock_ops.rs

@@ -68,7 +68,7 @@ impl SockOps {
                 io_error,
             }
         })?;
-        self.data.links.insert(SockOpsLink(ProgAttachLink::new(
+        self.data.links.insert(SockOpsLink::new(ProgAttachLink::new(
             prog_fd,
             cgroup_fd,
             BPF_CGROUP_SOCK_OPS,

aya/src/programs/tc.rs

@@ -165,7 +165,7 @@ impl SchedClassifier {
         }
         .map_err(|io_error| TcError::NetlinkError { io_error })?;
 
-        self.data.links.insert(SchedClassifierLink(TcLink {
+        self.data.links.insert(SchedClassifierLink::new(TcLink {
             if_index: if_index as i32,
             attach_type,
             priority,

aya/src/programs/xdp.rs

@@ -117,16 +117,18 @@ impl Xdp {
             )? as RawFd;
             self.data
                 .links
-                .insert(XdpLink(XdpLinkInner::FdLink(FdLink::new(link_fd))))
+                .insert(XdpLink::new(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 })?;
 
-            self.data.links.insert(XdpLink(XdpLinkInner::NlLink(NlLink {
-                if_index,
-                prog_fd,
-                flags,
-            })))
+            self.data
+                .links
+                .insert(XdpLink::new(XdpLinkInner::NlLink(NlLink {
+                    if_index,
+                    prog_fd,
+                    flags,
+                })))
         }
     }
 
@@ -150,9 +152,9 @@ impl Xdp {
     /// Ownership of the link will transfer to this program.
     pub fn attach_to_link(&mut self, link: XdpLink) -> Result<XdpLinkId, ProgramError> {
         let prog_fd = self.data.fd_or_err()?;
-        match &link.0 {
+        match link.inner() {
             XdpLinkInner::FdLink(fd_link) => {
-                let link_fd = fd_link.fd.unwrap();
+                let link_fd = fd_link.fd;
                 bpf_link_update(link_fd, prog_fd, None, 0).map_err(|(_, io_error)| {
                     ProgramError::SyscallError {
                         call: "bpf_link_update".to_string(),
@@ -163,7 +165,7 @@ impl Xdp {
                 mem::forget(link);
                 self.data
                     .links
-                    .insert(XdpLink(XdpLinkInner::FdLink(FdLink::new(link_fd))))
+                    .insert(XdpLink::new(XdpLinkInner::FdLink(FdLink::new(link_fd))))
             }
             XdpLinkInner::NlLink(nl_link) => {
                 let if_index = nl_link.if_index;
@@ -176,11 +178,13 @@ impl Xdp {
                 }
                 // dispose of link and avoid detach on drop
                 mem::forget(link);
-                self.data.links.insert(XdpLink(XdpLinkInner::NlLink(NlLink {
-                    if_index,
-                    prog_fd,
-                    flags,
-                })))
+                self.data
+                    .links
+                    .insert(XdpLink::new(XdpLinkInner::NlLink(NlLink {
+                        if_index,
+                        prog_fd,
+                        flags,
+                    })))
             }
         }
     }
@@ -246,7 +250,7 @@ impl TryFrom<XdpLink> for FdLink {
     type Error = LinkError;
 
     fn try_from(value: XdpLink) -> Result<Self, Self::Error> {
-        if let XdpLinkInner::FdLink(fd) = value.0 {
+        if let XdpLinkInner::FdLink(fd) = value.into_inner() {
             Ok(fd)
         } else {
             Err(LinkError::InvalidLink)
@@ -259,15 +263,14 @@ impl TryFrom<FdLink> for XdpLink {
 
     fn try_from(fd_link: FdLink) -> Result<Self, Self::Error> {
         // unwrap of fd_link.fd will not panic since it's only None when being dropped.
-        let info = bpf_link_get_info_by_fd(fd_link.fd.unwrap()).map_err(|io_error| {
-            LinkError::SyscallError {
+        let info =
+            bpf_link_get_info_by_fd(fd_link.fd).map_err(|io_error| LinkError::SyscallError {
                 call: "BPF_OBJ_GET_INFO_BY_FD".to_string(),
                 code: 0,
                 io_error,
-            }
-        })?;
+            })?;
         if info.type_ == (bpf_link_type::BPF_LINK_TYPE_XDP as u32) {
-            return Ok(XdpLink(XdpLinkInner::FdLink(fd_link)));
+            return Ok(XdpLink::new(XdpLinkInner::FdLink(fd_link)));
         }
         Err(LinkError::InvalidLink)
     }

test/integration-ebpf/src/test.rs

@@ -1,9 +1,13 @@
 #![no_std]
 #![no_main]
 
-use aya_bpf::{bindings::xdp_action, macros::xdp, programs::XdpContext};
+use aya_bpf::{
+    bindings::xdp_action,
+    macros::{kprobe, xdp},
+    programs::{ProbeContext, XdpContext},
+};
 
-#[xdp(name = "test_unload")]
+#[xdp(name = "test_unload_xdp")]
 pub fn pass(ctx: XdpContext) -> u32 {
     match unsafe { try_pass(ctx) } {
         Ok(ret) => ret,
@@ -15,6 +19,12 @@ unsafe fn try_pass(_ctx: XdpContext) -> Result<u32, u32> {
     Ok(xdp_action::XDP_PASS)
 }
 
+#[kprobe]
+// truncated name to match bpftool output
+pub fn test_unload_kpr(_ctx: ProbeContext) -> u32 {
+    0
+}
+
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }

test/integration-test/src/tests/load.rs

@@ -5,7 +5,7 @@ use aya::{
     maps::Array,
     programs::{
         links::{FdLink, PinnedLink},
-        TracePoint, Xdp, XdpFlags,
+        KProbe, TracePoint, Xdp, XdpFlags,
     },
     Bpf,
 };
@@ -84,28 +84,63 @@ macro_rules! assert_loaded {
 }
 
 #[integration_test]
-fn unload() {
+fn unload_xdp() {
     let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/test");
     let mut bpf = Bpf::load(bytes).unwrap();
-    let prog: &mut Xdp = bpf.program_mut("test_unload").unwrap().try_into().unwrap();
+    let prog: &mut Xdp = bpf
+        .program_mut("test_unload_xdp")
+        .unwrap()
+        .try_into()
+        .unwrap();
     prog.load().unwrap();
+    assert_loaded!("test_unload_xdp", true);
     let link = prog.attach("lo", XdpFlags::default()).unwrap();
     {
-        let _link_owned = prog.take_link(link);
+        let _link_owned = prog.take_link(link).unwrap();
         prog.unload().unwrap();
-        assert_loaded!("test_unload", true);
+        assert_loaded!("test_unload_xdp", true);
     };
 
-    assert_loaded!("test_unload", false);
+    assert_loaded!("test_unload_xdp", false);
     prog.load().unwrap();
 
-    assert_loaded!("test_unload", true);
+    assert_loaded!("test_unload_xdp", true);
     prog.attach("lo", XdpFlags::default()).unwrap();
 
-    assert_loaded!("test_unload", true);
+    assert_loaded!("test_unload_xdp", true);
+    prog.unload().unwrap();
+
+    assert_loaded!("test_unload_xdp", false);
+}
+
+#[integration_test]
+fn unload_kprobe() {
+    let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/test");
+    let mut bpf = Bpf::load(bytes).unwrap();
+    let prog: &mut KProbe = bpf
+        .program_mut("test_unload_kpr")
+        .unwrap()
+        .try_into()
+        .unwrap();
+    prog.load().unwrap();
+    assert_loaded!("test_unload_kpr", true);
+    let link = prog.attach("try_to_wake_up", 0).unwrap();
+    {
+        let _link_owned = prog.take_link(link).unwrap();
+        prog.unload().unwrap();
+        assert_loaded!("test_unload_kpr", true);
+    };
+
+    assert_loaded!("test_unload_kpr", false);
+    prog.load().unwrap();
+
+    assert_loaded!("test_unload_kpr", true);
+    prog.attach("try_to_wake_up", 0).unwrap();
+
+    assert_loaded!("test_unload_kpr", true);
     prog.unload().unwrap();
 
-    assert_loaded!("test_unload", false);
+    assert_loaded!("test_unload_kpr", false);
 }
 
 #[integration_test]
@@ -117,26 +152,30 @@ fn pin_link() {
 
     let bytes = include_bytes_aligned!("../../../../target/bpfel-unknown-none/debug/test");
     let mut bpf = Bpf::load(bytes).unwrap();
-    let prog: &mut Xdp = bpf.program_mut("test_unload").unwrap().try_into().unwrap();
+    let prog: &mut Xdp = bpf
+        .program_mut("test_unload_xdp")
+        .unwrap()
+        .try_into()
+        .unwrap();
     prog.load().unwrap();
     let link_id = prog.attach("lo", XdpFlags::default()).unwrap();
     let link = prog.take_link(link_id).unwrap();
-    assert_loaded!("test_unload", true);
+    assert_loaded!("test_unload_xdp", true);
 
     let fd_link: FdLink = link.try_into().unwrap();
     let pinned = fd_link.pin("/sys/fs/bpf/aya-xdp-test-lo").unwrap();
 
     // because of the pin, the program is still attached
     prog.unload().unwrap();
-    assert_loaded!("test_unload", true);
+    assert_loaded!("test_unload_xdp", true);
 
     // delete the pin, but the program is still attached
     let new_link = pinned.unpin().unwrap();
-    assert_loaded!("test_unload", true);
+    assert_loaded!("test_unload_xdp", true);
 
     // finally when new_link is dropped we're detached
     drop(new_link);
-    assert_loaded!("test_unload", false);
+    assert_loaded!("test_unload_xdp", false);
 }
 
 #[integration_test]