From 6d5bf930810f674601cf14f522c62c8daa9c0cd6 Mon Sep 17 00:00:00 2001
From: Jordan Rome <linux@jordanrome.com>
Date: Thu, 16 Nov 2023 07:16:01 -0800
Subject: [PATCH] [uprobe] support for pid targeting for shared libs

This adds support for specifying the pid even when targeting
a uprobe/uretprobe in a library shared by multiple pids.

Example:
```
sudo bpftrace -p 1899508 -e 'uprobe:libc:getaddrinfo {
  print((comm, pid));
}
```

[Issue 2817](https://github.com/iovisor/bpftrace/issues/2817)
---
 src/attached_probe.cpp | 28 +++++++++++++++++-----------
 src/attached_probe.h   |  7 ++++---
 src/bpftrace.cpp       |  7 +++++++
 3 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/src/attached_probe.cpp b/src/attached_probe.cpp
index 97a32f755494..9613a8ad301e 100644
--- a/src/attached_probe.cpp
+++ b/src/attached_probe.cpp
@@ -191,7 +191,7 @@ AttachedProbe::AttachedProbe(Probe &probe,
       // If BPF_PROG_TYPE_RAW_TRACEPOINT is available, no need to attach prog
       // to anything -- we will simply BPF_PROG_RUN it
       if (!feature.has_raw_tp_special())
-        attach_uprobe(safe_mode);
+        attach_uprobe(0, safe_mode);
       break;
     case ProbeType::kprobe:
       attach_kprobe(safe_mode);
@@ -200,10 +200,6 @@ AttachedProbe::AttachedProbe(Probe &probe,
       check_banned_kretprobes(probe_.attach_point);
       attach_kprobe(safe_mode);
       break;
-    case ProbeType::uprobe:
-    case ProbeType::uretprobe:
-      attach_uprobe(safe_mode);
-      break;
     case ProbeType::tracepoint:
       attach_tracepoint();
       break;
@@ -239,7 +235,8 @@ AttachedProbe::AttachedProbe(Probe &probe,
                              BpfProgram &&prog,
                              int pid,
                              BPFfeature &feature,
-                             BTF &btf)
+                             BTF &btf,
+                             bool safe_mode)
     : probe_(probe), prog_(std::move(prog)), btf_(btf)
 {
   load_prog(feature);
@@ -252,6 +249,10 @@ AttachedProbe::AttachedProbe(Probe &probe,
     case ProbeType::asyncwatchpoint:
       attach_watchpoint(pid, probe.mode);
       break;
+    case ProbeType::uprobe:
+    case ProbeType::uretprobe:
+      attach_uprobe(pid, safe_mode);
+      break;
     default:
       LOG(FATAL) << "invalid attached probe type \""
                  << probetypeName(probe_.type) << "\"";
@@ -1081,7 +1082,7 @@ static void resolve_offset_uprobe_multi(const std::string &path,
   }
 }
 
-void AttachedProbe::attach_multi_uprobe(void)
+void AttachedProbe::attach_multi_uprobe(int pid)
 {
   std::vector<std::string> syms;
   std::vector<uint64_t> offsets;
@@ -1100,6 +1101,11 @@ void AttachedProbe::attach_multi_uprobe(void)
   opts.uprobe_multi.flags = probe_.type == ProbeType::uretprobe
                                 ? BPF_F_UPROBE_MULTI_RETURN
                                 : 0;
+  if (pid != 0)
+  {
+    opts.uprobe_multi.pid = pid;
+  }
+
   if (bt_verbose)
   {
     std::cout << "Attaching to " << probe_.funcs.size() << " functions"
@@ -1125,16 +1131,16 @@ void AttachedProbe::attach_multi_uprobe(void)
   }
 }
 #else
-void AttachedProbe::attach_multi_uprobe(void)
+void AttachedProbe::attach_multi_uprobe(int)
 {
 }
 #endif // HAVE_LIBBPF_UPROBE_MULTI
 
-void AttachedProbe::attach_uprobe(bool safe_mode)
+void AttachedProbe::attach_uprobe(int pid, bool safe_mode)
 {
   if (!probe_.funcs.empty())
   {
-    attach_multi_uprobe();
+    attach_multi_uprobe(pid);
     return;
   }
 
@@ -1146,7 +1152,7 @@ void AttachedProbe::attach_uprobe(bool safe_mode)
                                         eventname().c_str(),
                                         probe_.path.c_str(),
                                         offset_,
-                                        probe_.pid,
+                                        pid == 0 ? -1 : pid,
                                         0);
 
   if (perf_event_fd < 0)
diff --git a/src/attached_probe.h b/src/attached_probe.h
index 646ee43d6aa0..af2ef0e012a9 100644
--- a/src/attached_probe.h
+++ b/src/attached_probe.h
@@ -30,7 +30,8 @@ class AttachedProbe
                 BpfProgram &&prog,
                 int pid,
                 BPFfeature &feature,
-                BTF &btf);
+                BTF &btf,
+                bool safe_mode = false);
   ~AttachedProbe();
   AttachedProbe(const AttachedProbe &) = delete;
   AttachedProbe &operator=(const AttachedProbe &) = delete;
@@ -47,9 +48,9 @@ class AttachedProbe
   bool resolve_offset_uprobe(bool safe_mode);
   void load_prog(BPFfeature &feature);
   void attach_multi_kprobe(void);
-  void attach_multi_uprobe(void);
+  void attach_multi_uprobe(int pid);
   void attach_kprobe(bool safe_mode);
-  void attach_uprobe(bool safe_mode);
+  void attach_uprobe(int pid, bool safe_mode);
 
   // Note: the following usdt attachment functions will only activate a
   // semaphore if one exists.
diff --git a/src/bpftrace.cpp b/src/bpftrace.cpp
index 18179e648a66..de0224787c57 100644
--- a/src/bpftrace.cpp
+++ b/src/bpftrace.cpp
@@ -1006,6 +1006,13 @@ std::vector<std::unique_ptr<AttachedProbe>> BPFtrace::attach_probe(
 
       return ret;
     }
+    else if (probe.type == ProbeType::uprobe ||
+             probe.type == ProbeType::uretprobe)
+    {
+      ret.emplace_back(std::make_unique<AttachedProbe>(
+          probe, std::move(*program), pid, *feature_, *btf_, safe_mode_));
+      return ret;
+    }
     else if (probe.type == ProbeType::watchpoint ||
              probe.type == ProbeType::asyncwatchpoint)
     {