Browse Source

programs: add internal API to create links

Alessandro Decina 4 years ago
parent
commit
f88ca1f1f1

+ 14 - 21
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<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)
+    }
 }
 
 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<T: Link> {
-    inner: Weak<RefCell<T>>,
+pub struct LinkRef {
+    inner: Rc<RefCell<dyn Link>>,
 }
 
-impl<T: Link> LinkRef<T> {
-    fn new(inner: &Rc<RefCell<T>>) -> LinkRef<T> {
-        LinkRef {
-            inner: Rc::downgrade(inner),
-        }
+impl LinkRef {
+    fn new(link: Rc<RefCell<dyn Link>>) -> LinkRef {
+        LinkRef { inner: link }
     }
 }
 
-impl<T: Link> Link for LinkRef<T> {
+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()
     }
 }
 

+ 2 - 7
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<impl Link, ProgramError> {
-    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<LinkRef, ProgramError> {
     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) }))
 }

+ 4 - 4
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<pid_t>,
-    ) -> Result<impl Link, ProgramError> {
+    ) -> Result<LinkRef, ProgramError> {
         attach(&mut self.data, ProbeKind::KProbe, fn_name, offset, pid)
     }
 }
@@ -65,7 +65,7 @@ impl UProbe {
         offset: u64,
         target: T,
         pid: Option<pid_t>,
-    ) -> Result<impl Link, ProgramError> {
+    ) -> Result<LinkRef, ProgramError> {
         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<pid_t>,
-) -> Result<impl Link, ProgramError> {
+) -> Result<LinkRef, ProgramError> {
     use ProbeKind::*;
 
     let perf_ty = read_sys_fs_perf_type(match kind {

+ 6 - 7
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<impl Link, ProgramError> {
+    pub fn attach(&mut self, socket: RawFd) -> Result<LinkRef, ProgramError> {
         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<RawFd>,
 }

+ 2 - 2
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<impl Link, ProgramError> {
+    pub fn attach(&mut self, category: &str, name: &str) -> Result<LinkRef, ProgramError> {
         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 })?

+ 7 - 11
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<impl Link, ProgramError> {
+    pub fn attach(&mut self, interface: &str) -> Result<LinkRef, ProgramError> {
         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))
+            })))
         }
     }
 }