From 5d95132bb3e4d0c77534dfa5ffda09f005671cdd Mon Sep 17 00:00:00 2001
From: Konrad Sztyber <konrad.sztyber@intel.com>
Date: Fri, 8 Oct 2021 13:09:25 +0200
Subject: [PATCH] scripts/bpftrace: use SO names in bpftrace scripts

If an USDT probe is defined within a shared object, bpftrace expects the
path to that shared object instead of an executable.  This means that we
need to replace __EXE__ in the bpftrace scripts differently, depending
on whether the application was linked statically or dynamically.

This is now done by the `scripts/bpf/gen.py` script.  It lists all
available probes, along with their locations, of a process.  Then, it
matches them to those described in a bpftrace script replacing the
__EXE__ markers with the listed location (either a path to the
executable or a shared library).  If a bpftrace script uses a probe that
isn't listed, its __EXE__ will be replaced by a path to the executable.

Signed-off-by: Konrad Sztyber <konrad.sztyber@intel.com>
Change-Id: I7c323d5f7d948ea57cf8d4d3132e4d59a2de594f
Reviewed-on: https://review.spdk.io/gerrit/c/spdk/spdk/+/9807
Community-CI: Broadcom CI <spdk-ci.pdl@broadcom.com>
Tested-by: SPDK CI Jenkins <sys_sgci@intel.com>
Reviewed-by: Tomasz Zawadzki <tomasz.zawadzki@intel.com>
Reviewed-by: Jim Harris <james.r.harris@intel.com>
---
 scripts/bpf/gen.py  | 53 +++++++++++++++++++++++++++++++++++++++++++++
 scripts/bpftrace.sh |  2 +-
 2 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100755 scripts/bpf/gen.py

diff --git a/scripts/bpf/gen.py b/scripts/bpf/gen.py
new file mode 100755
index 000000000..f204cffa9
--- /dev/null
+++ b/scripts/bpf/gen.py
@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+
+from argparse import ArgumentParser
+import os
+import re
+import subprocess
+import sys
+
+
+class TraceProcess:
+    def __init__(self, pid):
+        self._path = os.readlink(f'/proc/{pid}/exe')
+        self._pid = pid
+        self._probes = self._init_probes()
+
+    def _init_probes(self):
+        lines = subprocess.check_output(['bpftrace', '-l', '-p', str(self._pid)], text=True)
+        probes = {}
+        for line in lines.split('\n'):
+            parts = line.split(':')
+            if len(parts) < 3:
+                continue
+            ptype, path, function = parts[0], parts[1], parts[-1]
+            probes[(ptype, function)] = path
+        return probes
+
+    def fixup(self, script):
+        pregs = [re.compile(r'({}):__EXE__:(\w+)'.format(ptype)) for ptype in ['usdt', 'uprobe']]
+        with open(script, 'r') as file:
+            lines = file.readlines()
+        result = ''
+        for line in lines:
+            for regex in pregs:
+                match = regex.match(line)
+                if match is not None:
+                    ptype, function = match.groups()
+                    path = self._probes.get((ptype, function), self._path)
+                    line = line.replace('__EXE__', path)
+                    break
+            result += line.replace('__EXE__', self._path).replace('__PID__', str(self._pid))
+        return result
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser(description='bpftrace script generator replacing special ' +
+                            'variables in the scripts with appropriate values')
+    parser.add_argument('-p', '--pid', type=int, required=True, help='PID of a traced process')
+    parser.add_argument('scripts', metavar='SCRIPTS', type=str, nargs='+',
+                        help='bpftrace scripts to process')
+    args = parser.parse_args(sys.argv[1:])
+    proc = TraceProcess(args.pid)
+    for script in args.scripts:
+        print(proc.fixup(script))
diff --git a/scripts/bpftrace.sh b/scripts/bpftrace.sh
index 63e9b96ad..12050d0e7 100755
--- a/scripts/bpftrace.sh
+++ b/scripts/bpftrace.sh
@@ -10,7 +10,7 @@ if [ $# -lt 2 ]; then
 fi
 SCRIPTS_DIR=$(readlink -f $(dirname $0))
 BIN_PATH=$(readlink -f /proc/$1/exe)
-BPF_SCRIPT=$(sed "s#__EXE__#${BIN_PATH}#g" "${@:2}" | sed "s#__PID__#${1}#g")
+BPF_SCRIPT=$($SCRIPTS_DIR/bpf/gen.py -p $1 "${@:2}")
 BPF_SCRIPT+=$($SCRIPTS_DIR/bpf/gen_enums.sh)
 if [ -n "$ECHO_SCRIPT" ]; then
 	echo "$BPF_SCRIPT"