diff --git a/src/bpf.rs b/src/bpf.rs index 00439a1f..d23d2735 100644 --- a/src/bpf.rs +++ b/src/bpf.rs @@ -121,82 +121,21 @@ impl Bpf { .transpose() } - pub fn program(&self, name: &str) -> Option<&Program> { - self.programs.get(name) - } - - pub fn program_mut(&mut self, name: &str) -> Option<&mut Program> { - self.programs.get_mut(name) - } - - pub fn kprobe(&self, name: &str) -> Option<&KProbe> { - match self.programs.get(name) { - Some(Program::KProbe(kprobe)) => Some(kprobe), - _ => None, - } - } - - pub fn kprobe_mut(&mut self, name: &str) -> Option<&mut KProbe> { - match self.programs.get_mut(name) { - Some(Program::KProbe(kprobe)) => Some(kprobe), - _ => None, - } - } - - pub fn uprobe(&self, name: &str) -> Option<&UProbe> { - match self.programs.get(name) { - Some(Program::UProbe(uprobe)) => Some(uprobe), - _ => None, - } - } - - pub fn uprobe_mut(&mut self, name: &str) -> Option<&mut UProbe> { - match self.programs.get_mut(name) { - Some(Program::UProbe(uprobe)) => Some(uprobe), - _ => None, - } - } - - pub fn trace_point(&self, name: &str) -> Option<&TracePoint> { - match self.programs.get(name) { - Some(Program::TracePoint(trace_point)) => Some(trace_point), - _ => None, - } - } - - pub fn trace_point_mut(&mut self, name: &str) -> Option<&mut TracePoint> { - match self.programs.get_mut(name) { - Some(Program::TracePoint(trace_point)) => Some(trace_point), - _ => None, - } - } - - pub fn socket_filter(&self, name: &str) -> Option<&SocketFilter> { - match self.programs.get(name) { - Some(Program::SocketFilter(socket_filter)) => Some(socket_filter), - _ => None, - } - } - - pub fn socket_filter_mut(&mut self, name: &str) -> Option<&mut SocketFilter> { - match self.programs.get_mut(name) { - Some(Program::SocketFilter(socket_filter)) => Some(socket_filter), - _ => None, - } - } - - pub fn xdp(&self, name: &str) -> Option<&Xdp> { - match self.programs.get(name) { - Some(Program::Xdp(xdp)) => Some(xdp), - _ => None, - } + pub fn program<'a, 'slf: 'a, T: TryFrom<&'a Program>>( + &'slf self, + name: &str, + ) -> Result, >::Error> { + self.programs.get(name).map(|p| T::try_from(p)).transpose() } - pub fn xdp_mut(&mut self, name: &str) -> Option<&mut Xdp> { - match self.programs.get_mut(name) { - Some(Program::Xdp(xdp)) => Some(xdp), - _ => None, - } + pub fn program_mut<'a, 'slf: 'a, T: TryFrom<&'a mut Program>>( + &'slf mut self, + name: &str, + ) -> Result, >::Error> { + self.programs + .get_mut(name) + .map(|p| T::try_from(p)) + .transpose() } } diff --git a/src/programs/mod.rs b/src/programs/mod.rs index c87b52f4..9b8c1601 100644 --- a/src/programs/mod.rs +++ b/src/programs/mod.rs @@ -5,24 +5,24 @@ mod trace_point; mod xdp; use libc::{close, ENOSPC}; -use perf_attach::*; -pub use probe::*; -pub use socket_filter::*; -pub use trace_point::*; -pub use xdp::*; - use std::{ cell::RefCell, cmp, + convert::TryFrom, ffi::CStr, io, os::raw::c_uint, path::PathBuf, rc::{Rc, Weak}, }; - use thiserror::Error; +use perf_attach::*; +pub use probe::*; +pub use socket_filter::*; +pub use trace_point::*; +pub use xdp::*; + use crate::{obj, syscalls::bpf_load_program, RawFd}; #[derive(Debug, Error)] pub enum ProgramError { @@ -73,6 +73,9 @@ pub enum ProgramError { #[error("setsockopt SO_ATTACH_BPF failed: {io_error}")] SocketFilterError { io_error: io::Error }, + #[error("unexpected program type")] + UnexpectedProgramType, + #[error("{message}")] Other { message: String }, } @@ -127,26 +130,6 @@ impl Program { } } -impl ProgramFd for Program { - fn fd(&self) -> Option { - self.data().fd - } -} - -macro_rules! impl_program_fd { - ($($struct_name:ident),+ $(,)?) => { - $( - impl ProgramFd for $struct_name { - fn fd(&self) -> Option { - self.data.fd - } - } - )+ - } -} - -impl_program_fd!(KProbe, UProbe, TracePoint, SocketFilter, Xdp); - #[derive(Debug)] pub(crate) struct ProgramData { pub(crate) name: String, @@ -312,3 +295,53 @@ impl Drop for FdLink { let _ = self.detach(); } } + +impl ProgramFd for Program { + fn fd(&self) -> Option { + self.data().fd + } +} + +macro_rules! impl_program_fd { + ($($struct_name:ident),+ $(,)?) => { + $( + impl ProgramFd for $struct_name { + fn fd(&self) -> Option { + self.data.fd + } + } + )+ + } +} + +impl_program_fd!(KProbe, UProbe, TracePoint, SocketFilter, Xdp); + +macro_rules! impl_try_from_program { + ($($ty:ident),+ $(,)?) => { + $( + impl<'a> TryFrom<&'a Program> for &'a $ty { + type Error = ProgramError; + + fn try_from(program: &'a Program) -> Result<&'a $ty, ProgramError> { + match program { + Program::$ty(p) => Ok(p), + _ => Err(ProgramError::UnexpectedProgramType), + } + } + } + + impl<'a> TryFrom<&'a mut Program> for &'a mut $ty { + type Error = ProgramError; + + fn try_from(program: &'a mut Program) -> Result<&'a mut $ty, ProgramError> { + match program { + Program::$ty(p) => Ok(p), + _ => Err(ProgramError::UnexpectedProgramType), + } + } + } + )+ + } +} + +impl_try_from_program!(KProbe, UProbe, TracePoint, SocketFilter, Xdp); diff --git a/src/programs/probe.rs b/src/programs/probe.rs index 1780386d..a554f455 100644 --- a/src/programs/probe.rs +++ b/src/programs/probe.rs @@ -12,12 +12,10 @@ use thiserror::Error; use crate::{ generated::bpf_prog_type::BPF_PROG_TYPE_KPROBE, - programs::{load_program, ProgramData, ProgramError}, + programs::{load_program, perf_attach, Link, ProgramData, ProgramError}, syscalls::perf_event_open_probe, }; -use super::{perf_attach, Link}; - lazy_static! { static ref LD_SO_CACHE: Result = LdSoCache::load("/etc/ld.so.cache"); } diff --git a/src/programs/socket_filter.rs b/src/programs/socket_filter.rs index bb655de6..f50979dc 100644 --- a/src/programs/socket_filter.rs +++ b/src/programs/socket_filter.rs @@ -1,10 +1,10 @@ -use std::{io, mem, os::unix::prelude::RawFd}; - use libc::{setsockopt, SOL_SOCKET, SO_ATTACH_BPF}; +use std::{io, mem, os::unix::prelude::RawFd}; -use crate::generated::bpf_prog_type::BPF_PROG_TYPE_SOCKET_FILTER; - -use super::{load_program, ProgramData, ProgramError}; +use crate::{ + generated::bpf_prog_type::BPF_PROG_TYPE_SOCKET_FILTER, + programs::{load_program, ProgramData, ProgramError}, +}; #[derive(Debug)] pub struct SocketFilter { diff --git a/src/programs/trace_point.rs b/src/programs/trace_point.rs index 474edd2a..d7f2fe60 100644 --- a/src/programs/trace_point.rs +++ b/src/programs/trace_point.rs @@ -1,4 +1,4 @@ -use std::{convert::TryFrom, fs}; +use std::fs; use crate::{ generated::bpf_prog_type::BPF_PROG_TYPE_TRACEPOINT, syscalls::perf_event_open_trace_point,