浏览代码

aya: implement missing bit of retprobes

Alessandro Decina 4 年之前
父节点
当前提交
f11df77f85
共有 3 个文件被更改,包括 32 次插入12 次删除
  1. 20 3
      aya/src/bpf.rs
  2. 6 0
      aya/src/obj/mod.rs
  3. 6 9
      aya/src/programs/probe.rs

+ 20 - 3
aya/src/bpf.rs

@@ -15,7 +15,10 @@ use crate::{
         btf::{Btf, BtfError},
         Object, ParseError,
     },
-    programs::{KProbe, Program, ProgramData, ProgramError, SocketFilter, TracePoint, UProbe, Xdp},
+    programs::{
+        probe::ProbeKind, KProbe, Program, ProgramData, ProgramError, SocketFilter, TracePoint,
+        UProbe, Xdp,
+    },
     sys::bpf_map_update_elem_ptr,
     util::{possible_cpus, POSSIBLE_CPUS},
 };
@@ -109,8 +112,22 @@ impl Bpf {
                     links: Vec::new(),
                 };
                 let program = match kind {
-                    crate::obj::ProgramKind::KProbe => Program::KProbe(KProbe { data }),
-                    crate::obj::ProgramKind::UProbe => Program::UProbe(UProbe { data }),
+                    crate::obj::ProgramKind::KProbe => Program::KProbe(KProbe {
+                        data,
+                        kind: ProbeKind::KProbe,
+                    }),
+                    crate::obj::ProgramKind::KRetProbe => Program::KProbe(KProbe {
+                        data,
+                        kind: ProbeKind::KRetProbe,
+                    }),
+                    crate::obj::ProgramKind::UProbe => Program::UProbe(UProbe {
+                        data,
+                        kind: ProbeKind::UProbe,
+                    }),
+                    crate::obj::ProgramKind::URetProbe => Program::UProbe(UProbe {
+                        data,
+                        kind: ProbeKind::URetProbe,
+                    }),
                     crate::obj::ProgramKind::TracePoint => Program::TracePoint(TracePoint { data }),
                     crate::obj::ProgramKind::SocketFilter => {
                         Program::SocketFilter(SocketFilter { data })

+ 6 - 0
aya/src/obj/mod.rs

@@ -58,7 +58,9 @@ pub(crate) struct Program {
 #[derive(Debug, Copy, Clone)]
 pub enum ProgramKind {
     KProbe,
+    KRetProbe,
     UProbe,
+    URetProbe,
     TracePoint,
     SocketFilter,
     Xdp,
@@ -71,7 +73,9 @@ impl FromStr for ProgramKind {
         use ProgramKind::*;
         Ok(match kind {
             "kprobe" => KProbe,
+            "kretprobe" => KRetProbe,
             "uprobe" => UProbe,
+            "uretprobe" => URetProbe,
             "xdp" => Xdp,
             "trace_point" => TracePoint,
             "socket_filter" => SocketFilter,
@@ -186,7 +190,9 @@ impl Object {
                     .insert(name.to_string(), parse_map(&section, name)?);
             }
             &[ty @ "kprobe", name]
+            | &[ty @ "kretprobe", name]
             | &[ty @ "uprobe", name]
+            | &[ty @ "uretprobe", name]
             | &[ty @ "socket_filter", name]
             | &[ty @ "xdp", name]
             | &[ty @ "trace_point", name] => {

+ 6 - 9
aya/src/programs/probe.rs

@@ -65,11 +65,13 @@ pub enum UProbeError {
 #[derive(Debug)]
 pub struct KProbe {
     pub(crate) data: ProgramData,
+    pub(crate) kind: ProbeKind,
 }
 
 #[derive(Debug)]
 pub struct UProbe {
     pub(crate) data: ProgramData,
+    pub(crate) kind: ProbeKind,
 }
 
 impl KProbe {
@@ -87,7 +89,7 @@ impl KProbe {
         offset: u64,
         pid: Option<pid_t>,
     ) -> Result<LinkRef, ProgramError> {
-        attach(&mut self.data, ProbeKind::KProbe, fn_name, offset, pid)
+        attach(&mut self.data, self.kind, fn_name, offset, pid)
     }
 }
 
@@ -147,17 +149,12 @@ impl UProbe {
             0
         };
 
-        attach(
-            &mut self.data,
-            ProbeKind::UProbe,
-            &path,
-            sym_offset + offset,
-            pid,
-        )
+        attach(&mut self.data, self.kind, &path, sym_offset + offset, pid)
     }
 }
 
-enum ProbeKind {
+#[derive(Debug, Copy, Clone)]
+pub(crate) enum ProbeKind {
     KProbe,
     KRetProbe,
     UProbe,