Sfoglia il codice sorgente

bpf: make program() and program_mut() return inner program types

program() and program_mut() are now generic and can return the inner
program type, which is what you want 99.999% of the time, eg:

    let prog = bpf.program_mut::<&mut Xdp>("foo")?.unwrap();
    prog.load()?;
    prog.attach("eth0")?;

The methods will fail if the requested program is not of the expected
type (eg if you try to retrieve a kprobe as an xdp program).
Alessandro Decina 4 anni fa
parent
commit
a8c212377f
5 ha cambiato i file con 80 aggiunte e 110 eliminazioni
  1. 13 74
      src/bpf.rs
  2. 60 27
      src/programs/mod.rs
  3. 1 3
      src/programs/probe.rs
  4. 5 5
      src/programs/socket_filter.rs
  5. 1 1
      src/programs/trace_point.rs

+ 13 - 74
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<Option<T>, <T as TryFrom<&'a Program>>::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<Option<T>, <T as TryFrom<&'a mut Program>>::Error> {
+        self.programs
+            .get_mut(name)
+            .map(|p| T::try_from(p))
+            .transpose()
     }
 }
 

+ 60 - 27
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<RawFd> {
-        self.data().fd
-    }
-}
-
-macro_rules! impl_program_fd {
-    ($($struct_name:ident),+ $(,)?) => {
-        $(
-            impl ProgramFd for $struct_name {
-                fn fd(&self) -> Option<RawFd> {
-                    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<RawFd> {
+        self.data().fd
+    }
+}
+
+macro_rules! impl_program_fd {
+    ($($struct_name:ident),+ $(,)?) => {
+        $(
+            impl ProgramFd for $struct_name {
+                fn fd(&self) -> Option<RawFd> {
+                    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);

+ 1 - 3
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, io::Error> = LdSoCache::load("/etc/ld.so.cache");
 }

+ 5 - 5
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 {

+ 1 - 1
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,