From 3fa7ba1e8a819896510388c91c839ac53557be48 Mon Sep 17 00:00:00 2001
From: Sasha Goldshtein <goldshtn@gmail.com>
Date: Sat, 14 Jan 2017 11:17:40 +0000
Subject: [PATCH] argdist, trace: Support non-C identifier names

When argdist or trace face a function that has characters
in its name that are not valid in C identifier, they now
replace these characters with an underscore (`_`) when
generating function names and structure names to include
in the BPF program. As a result, it is now possible to
trace functions that have these identifiers in their names,
such as Golang functions like `fmt.Println`.
---
 tools/argdist.py | 20 +++++++++++++-------
 tools/trace.py   |  3 +--
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/tools/argdist.py b/tools/argdist.py
index 2e3aad5d976c..3bcf43a47740 100755
--- a/tools/argdist.py
+++ b/tools/argdist.py
@@ -159,7 +159,7 @@ def _validate_specifier(self):
                 if parts[0] not in ["r", "p", "t", "u"]:
                         self._bail("probe type must be 'p', 'r', 't', or 'u'" +
                                    " but got '%s'" % parts[0])
-                if re.match(r"\w+\(.*\)", parts[2]) is None:
+                if re.match(r"\S+\(.*\)", parts[2]) is None:
                         self._bail(("function signature '%s' has an invalid " +
                                     "format") % parts[2])
 
@@ -173,6 +173,9 @@ def _parse_exprs(self, exprs):
                         self._bail("no exprs specified")
                 self.exprs = exprs.split(',')
 
+        def _make_valid_identifier(self, ident):
+                return re.sub(r'[^A-Za-z0-9_]', '_', ident)
+
         def __init__(self, tool, type, specifier):
                 self.usdt_ctx = None
                 self.streq_functions = ""
@@ -196,8 +199,9 @@ def __init__(self, tool, type, specifier):
                         self.tp_event = self.function
                 elif self.probe_type == "u":
                         self.library = parts[1]
-                        self.probe_func_name = "%s_probe%d" % \
-                                (self.function, Probe.next_probe_index)
+                        self.probe_func_name = self._make_valid_identifier(
+                                "%s_probe%d" % \
+                                (self.function, Probe.next_probe_index))
                         self._enable_usdt_probe()
                 else:
                         self.library = parts[1]
@@ -233,10 +237,12 @@ def check(expr):
                 self.entry_probe_required = self.probe_type == "r" and \
                         (any(map(check, self.exprs)) or check(self.filter))
 
-                self.probe_func_name = "%s_probe%d" % \
-                        (self.function, Probe.next_probe_index)
-                self.probe_hash_name = "%s_hash%d" % \
-                        (self.function, Probe.next_probe_index)
+                self.probe_func_name = self._make_valid_identifier(
+                        "%s_probe%d" % \
+                        (self.function, Probe.next_probe_index))
+                self.probe_hash_name = self._make_valid_identifier(
+                        "%s_hash%d" % \
+                        (self.function, Probe.next_probe_index))
                 Probe.next_probe_index += 1
 
         def _enable_usdt_probe(self):
diff --git a/tools/trace.py b/tools/trace.py
index 80212b4bc759..8ac92010f159 100755
--- a/tools/trace.py
+++ b/tools/trace.py
@@ -76,8 +76,7 @@ def __init__(self, probe, string_size, kernel_stack, user_stack):
                 self.probe_num = Probe.probe_count
                 self.probe_name = "probe_%s_%d" % \
                                 (self._display_function(), self.probe_num)
-                if self.probe_name.find(".") > 0:   # for golang
-                    self.probe_name = self.probe_name.replace(".", "_DOT_")
+                self.probe_name = re.sub(r'[^A-Za-z0-9_]', '_', self.probe_name)
 
         def __str__(self):
                 return "%s:%s:%s FLT=%s ACT=%s/%s" % (self.probe_type,