diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs index 680ba7f3..e4a63990 100644 --- a/aya/src/programs/mod.rs +++ b/aya/src/programs/mod.rs @@ -6,14 +6,7 @@ mod xdp; use libc::{close, ENOSPC}; use std::{ - cell::RefCell, - cmp, - convert::TryFrom, - ffi::CStr, - io, - os::raw::c_uint, - path::PathBuf, - rc::{Rc, Weak}, + cell::RefCell, cmp, convert::TryFrom, ffi::CStr, io, os::raw::c_uint, path::PathBuf, rc::Rc, }; use thiserror::Error; @@ -151,6 +144,12 @@ impl ProgramData { program: self.name.clone(), }) } + + pub fn link(&mut self, link: T) -> LinkRef { + let link: Rc> = Rc::new(RefCell::new(link)); + self.links.push(Rc::clone(&link)); + LinkRef::new(link) + } } const MAX_LOG_BUF_SIZE: usize = (std::u32::MAX >> 8) as usize; @@ -259,25 +258,19 @@ pub trait Link: std::fmt::Debug { } #[derive(Debug)] -pub(crate) struct LinkRef { - inner: Weak>, +pub struct LinkRef { + inner: Rc>, } -impl LinkRef { - fn new(inner: &Rc>) -> LinkRef { - LinkRef { - inner: Rc::downgrade(inner), - } +impl LinkRef { + fn new(link: Rc>) -> LinkRef { + LinkRef { inner: link } } } -impl Link for LinkRef { +impl Link for LinkRef { fn detach(&mut self) -> Result<(), ProgramError> { - if let Some(inner) = self.inner.upgrade() { - inner.borrow_mut().detach() - } else { - Err(ProgramError::AlreadyDetached) - } + self.inner.borrow_mut().detach() } } diff --git a/aya/src/programs/perf_attach.rs b/aya/src/programs/perf_attach.rs index 425fd86a..5e170391 100644 --- a/aya/src/programs/perf_attach.rs +++ b/aya/src/programs/perf_attach.rs @@ -1,5 +1,3 @@ -use std::{cell::RefCell, rc::Rc}; - use libc::close; use crate::{ @@ -32,15 +30,12 @@ impl Drop for PerfLink { } } -pub(crate) fn perf_attach(data: &mut ProgramData, fd: RawFd) -> Result { - let link = Rc::new(RefCell::new(PerfLink { perf_fd: Some(fd) })); - data.links.push(link.clone()); - +pub(crate) fn perf_attach(data: &mut ProgramData, fd: RawFd) -> Result { let prog_fd = data.fd_or_err()?; perf_event_ioctl(fd, PERF_EVENT_IOC_SET_BPF, prog_fd) .map_err(|(_, io_error)| ProgramError::PerfEventAttachError { io_error })?; perf_event_ioctl(fd, PERF_EVENT_IOC_ENABLE, 0) .map_err(|(_, io_error)| ProgramError::PerfEventAttachError { io_error })?; - Ok(LinkRef::new(&link)) + Ok(data.link(PerfLink { perf_fd: Some(fd) })) } diff --git a/aya/src/programs/probe.rs b/aya/src/programs/probe.rs index 4c186aa3..25bb9f29 100644 --- a/aya/src/programs/probe.rs +++ b/aya/src/programs/probe.rs @@ -12,7 +12,7 @@ use thiserror::Error; use crate::{ generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE, - programs::{load_program, perf_attach, Link, ProgramData, ProgramError}, + programs::{load_program, perf_attach, LinkRef, ProgramData, ProgramError}, sys::perf_event_open_probe, }; @@ -45,7 +45,7 @@ impl KProbe { fn_name: &str, offset: u64, pid: Option, - ) -> Result { + ) -> Result { attach(&mut self.data, ProbeKind::KProbe, fn_name, offset, pid) } } @@ -65,7 +65,7 @@ impl UProbe { offset: u64, target: T, pid: Option, - ) -> Result { + ) -> Result { let target = target.as_ref(); let target_str = &*target.as_os_str().to_string_lossy(); @@ -128,7 +128,7 @@ fn attach( name: &str, offset: u64, pid: Option, -) -> Result { +) -> Result { use ProbeKind::*; let perf_ty = read_sys_fs_perf_type(match kind { diff --git a/aya/src/programs/socket_filter.rs b/aya/src/programs/socket_filter.rs index 4e405de3..c9e1e151 100644 --- a/aya/src/programs/socket_filter.rs +++ b/aya/src/programs/socket_filter.rs @@ -1,5 +1,5 @@ use libc::{setsockopt, SOL_SOCKET, SO_ATTACH_BPF, SO_DETACH_BPF}; -use std::{cell::RefCell, io, mem, os::unix::prelude::RawFd, rc::Rc}; +use std::{io, mem, os::unix::prelude::RawFd}; use crate::{ generated::bpf_prog_type::BPF_PROG_TYPE_SOCKET_FILTER, @@ -16,7 +16,7 @@ impl SocketFilter { load_program(BPF_PROG_TYPE_SOCKET_FILTER, &mut self.data) } - pub fn attach(&self, socket: RawFd) -> Result { + pub fn attach(&mut self, socket: RawFd) -> Result { let prog_fd = self.data.fd_or_err()?; let ret = unsafe { @@ -33,17 +33,16 @@ impl SocketFilter { io_error: io::Error::last_os_error(), }); } - let link = Rc::new(RefCell::new(SocketFilterLink { + + Ok(self.data.link(SocketFilterLink { socket, prog_fd: Some(prog_fd), - })); - - Ok(LinkRef::new(&link)) + })) } } #[derive(Debug)] -pub(crate) struct SocketFilterLink { +struct SocketFilterLink { socket: RawFd, prog_fd: Option, } diff --git a/aya/src/programs/trace_point.rs b/aya/src/programs/trace_point.rs index 24080622..0f1d972e 100644 --- a/aya/src/programs/trace_point.rs +++ b/aya/src/programs/trace_point.rs @@ -2,7 +2,7 @@ use std::fs; use crate::{generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT, sys::perf_event_open_trace_point}; -use super::{load_program, perf_attach, Link, ProgramData, ProgramError}; +use super::{load_program, perf_attach, LinkRef, ProgramData, ProgramError}; #[derive(Debug)] pub struct TracePoint { @@ -14,7 +14,7 @@ impl TracePoint { load_program(BPF_PROG_TYPE_TRACEPOINT, &mut self.data) } - pub fn attach(&mut self, category: &str, name: &str) -> Result { + pub fn attach(&mut self, category: &str, name: &str) -> Result { let id = read_sys_fs_trace_point_id(category, name)?; let fd = perf_event_open_trace_point(id) .map_err(|(_code, io_error)| ProgramError::PerfEventOpenError { io_error })? diff --git a/aya/src/programs/xdp.rs b/aya/src/programs/xdp.rs index 8469eb6a..a125adeb 100644 --- a/aya/src/programs/xdp.rs +++ b/aya/src/programs/xdp.rs @@ -1,4 +1,4 @@ -use std::{cell::RefCell, ffi::CString, rc::Rc}; +use std::ffi::CString; use libc::if_nametoindex; @@ -25,7 +25,7 @@ impl Xdp { self.data.name.to_string() } - pub fn attach(&mut self, interface: &str) -> Result { + pub fn attach(&mut self, interface: &str) -> Result { let prog_fd = self.data.fd_or_err()?; let c_interface = CString::new(interface).unwrap(); @@ -45,10 +45,9 @@ impl Xdp { io_error, } })? as RawFd; - let link = Rc::new(RefCell::new(XdpLink::FdLink(FdLink { fd: Some(link_fd) }))); - self.data.links.push(link.clone()); - - Ok(LinkRef::new(&link)) + Ok(self + .data + .link(XdpLink::FdLink(FdLink { fd: Some(link_fd) }))) } else { unsafe { netlink_set_xdp_fd(if_index, prog_fd, None, 0) }.map_err(|io_error| { ProgramError::NetlinkXdpError { @@ -57,13 +56,10 @@ impl Xdp { } })?; - let link = Rc::new(RefCell::new(XdpLink::NlLink(NlLink { + Ok(self.data.link(XdpLink::NlLink(NlLink { if_index, prog_fd: Some(prog_fd), - }))); - self.data.links.push(link.clone()); - - Ok(LinkRef::new(&link)) + }))) } } }