From 3a5256f101d43a8b9167df01dd3fc0dfdb0e6f00 Mon Sep 17 00:00:00 2001
From: Sasha Goldshtein <goldshtn@gmail.com>
Date: Mon, 20 Feb 2017 15:42:57 +0000
Subject: [PATCH] trace, argdist: Treat small USDT arguments correctly

trace and argdist currently only work correctly for USDT arguments
whose size is exactly 8 bytes. Smaller types, such as chars, shorts,
ints (signed or unsigned) are not treated correctly. The reason is
that the produced program would invoke the `bpf_usdt_readarg` helper
with the address of a u64 local variable, and then cast that variable
to the user-specified type derived from the format string. However,
the `bpf_usdt_readarg` rewriting then passes `sizeof(u64)` to the
generated `bpf_..._readarg` macro, which then fails to read anything
because the provided size doesn't match the argument size it knows
about.

The fix is fairly easy: instead of declaring a u64 unconditionally
and reading into that variable with `bpf_usdt_readarg`, declare a
variable that has the correct type according to what we know about
the USDT probe.
---
 tools/argdist.py | 7 +++++--
 tools/trace.py   | 7 +++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tools/argdist.py b/tools/argdist.py
index 09df3c368147..9902086dd339 100755
--- a/tools/argdist.py
+++ b/tools/argdist.py
@@ -293,9 +293,12 @@ def _generate_hash_field(self, i):
         def _generate_usdt_arg_assignment(self, i):
                 expr = self.exprs[i]
                 if self.probe_type == "u" and expr[0:3] == "arg":
-                        return ("        u64 %s = 0;\n" +
+                        arg_index = int(expr[3])
+                        arg_ctype = self.usdt_ctx.get_probe_arg_ctype(
+                                self.function, arg_index - 1)
+                        return ("        %s %s = 0;\n" +
                                 "        bpf_usdt_readarg(%s, ctx, &%s);\n") \
-                                % (expr, expr[3], expr)
+                                % (arg_ctype, expr, expr[3], expr)
                 else:
                         return ""
 
diff --git a/tools/trace.py b/tools/trace.py
index 3a957e9b304b..f2a874189798 100755
--- a/tools/trace.py
+++ b/tools/trace.py
@@ -321,9 +321,12 @@ def _generate_field_assign(self, idx):
                 expr = self.values[idx].strip()
                 text = ""
                 if self.probe_type == "u" and expr[0:3] == "arg":
-                        text = ("        u64 %s = 0;\n" +
+                        arg_index = int(expr[3])
+                        arg_ctype = self.usdt.get_probe_arg_ctype(
+                                self.usdt_name, arg_index - 1)
+                        text = ("        %s %s = 0;\n" +
                                 "        bpf_usdt_readarg(%s, ctx, &%s);\n") \
-                                % (expr, expr[3], expr)
+                                % (arg_ctype, expr, expr[3], expr)
 
                 if field_type == "s":
                         return text + """