From f11df77f859feee2a88a69b96da4a1a22839c45a Mon Sep 17 00:00:00 2001
From: Alessandro Decina <alessandro.d@gmail.com>
Date: Wed, 17 Feb 2021 08:29:13 +0000
Subject: [PATCH] aya: implement missing bit of retprobes

---
 aya/src/bpf.rs            | 23 ++++++++++++++++++++---
 aya/src/obj/mod.rs        |  6 ++++++
 aya/src/programs/probe.rs | 15 ++++++---------
 3 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs
index 3b9fffa9..ff8c31e2 100644
--- a/aya/src/bpf.rs
+++ b/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 })
diff --git a/aya/src/obj/mod.rs b/aya/src/obj/mod.rs
index 34035e9c..807fc718 100644
--- a/aya/src/obj/mod.rs
+++ b/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] => {
diff --git a/aya/src/programs/probe.rs b/aya/src/programs/probe.rs
index 5fbcd03b..9988683e 100644
--- a/aya/src/programs/probe.rs
+++ b/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,