usdt: Implement `bpf_usdt_readarg` as frontend action

examples/lua/usdt_ruby.lua

@@ -18,10 +18,14 @@ limitations under the License.
 local program = [[
 #include <uapi/linux/ptrace.h>
 int trace_method(struct pt_regs *ctx) {
-    char fn_name[128] = {};
-    bpf_usdt_readarg_p(method__entry_2, ctx, &fn_name, sizeof(fn_name));
-    bpf_trace_printk("%s(...)\n", &fn_name);
-    return 0;
+  uint64_t addr;
+  bpf_usdt_readarg(2, ctx, &addr);
+
+  char fn_name[128] = {};
+  bpf_probe_read(&fn_name, sizeof(fn_name), (void *)addr);
+
+  bpf_trace_printk("%s(...)\n", fn_name);
+  return 0;
 };
 ]]
 

src/cc/export/helpers.h

@@ -408,7 +408,9 @@ int incr_cksum_l3(void *off, u64 oldval, u64 newval) asm("llvm.bpf.extra");
 int incr_cksum_l4(void *off, u64 oldval, u64 newval, u64 flags) asm("llvm.bpf.extra");
 int bpf_num_cpus() asm("llvm.bpf.extra");
 
-#define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
+struct pt_regs;
+int bpf_usdt_readarg(int argc, struct pt_regs *ctx, void *arg) asm("llvm.bpf.extra");
+int bpf_usdt_readarg_p(int argc, struct pt_regs *ctx, void *buf, u64 len) asm("llvm.bpf.extra");
 
 #ifdef __powerpc__
 #define PT_REGS_PARM1(ctx)	((ctx)->gpr[3])
@@ -434,10 +436,7 @@ int bpf_num_cpus() asm("llvm.bpf.extra");
 #error "bcc does not support this platform yet"
 #endif
 
-#define bpf_usdt_readarg(probearg, ctx) _bpf_readarg_##probearg(ctx)
-#define bpf_usdt_readarg_p(probearg, ctx, buf, len) {\
-  u64 __addr = bpf_usdt_readarg(probearg, ctx); \
-  bpf_probe_read(buf, len, (void *)__addr); }
+#define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
 
 #endif
 )********"

src/cc/frontends/clang/b_frontend_action.cc

@@ -254,6 +254,7 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
   // put each non-static non-inline function decl in its own section, to be
   // extracted by the MemoryManager
   if (D->isExternallyVisible() && D->hasBody()) {
+    current_fn_ = D->getName();
     string attr = string("__attribute__((section(\"") + BPF_FN_PREFIX + D->getName().str() + "\")))\n";
     rewriter_.InsertText(D->getLocStart(), attr);
     if (D->param_size() > MAX_CALLING_CONV_REGS + 1) {
@@ -470,7 +471,21 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
             numcpu = 1;
           text = to_string(numcpu);
           rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
-        }
+        } else if (Decl->getName() == "bpf_usdt_readarg_p") {
+	 text = "({ u64 __addr = 0x0; ";
+	 text += "_bpf_readarg_" + current_fn_ + "_" + args[0] + "(" +
+	  args[1] + ", &__addr, sizeof(__addr));";
+	 text += "bpf_probe_read(" + args[2] + ", " + args[3] +
+	  ", (void *)__addr);";
+	 text += "})";
+	 rewriter_.ReplaceText(
+	   SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
+        } else if (Decl->getName() == "bpf_usdt_readarg") {
+	 text = "_bpf_readarg_" + current_fn_ + "_" + args[0] + "(" +
+	  args[1] + ", " + args[2] + ", sizeof(*(" + args[2] + ")))";
+	 rewriter_.ReplaceText(
+	   SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
+	}
       }
     }
   }

src/cc/frontends/clang/b_frontend_action.h

@@ -83,6 +83,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   std::vector<TableDesc> &tables_;  /// store the open FDs
   std::vector<clang::ParmVarDecl *> fn_args_;
   std::set<clang::Expr *> visited_;
+  std::string current_fn_;
 };
 
 // Do a depth-first search to rewrite all pointers that need to be probed

src/cc/usdt.cc

@@ -130,40 +130,6 @@ bool Probe::disable(int pid) {
   return result;
 }
 
-bool Probe::usdt_thunks(std::ostream &stream, const std::string &prefix) {
-  assert(!locations_.empty());
-  for (size_t i = 0; i < locations_.size(); ++i) {
-    tfm::format(
-        stream,
-        "int %s_thunk_%d(struct pt_regs *ctx) { return %s(ctx, %d); }\n",
-        prefix, i, prefix, i);
-  }
-  return true;
-}
-
-bool Probe::usdt_cases(std::ostream &stream, const optional<int> &pid) {
-  assert(!locations_.empty());
-  const size_t arg_count = locations_[0].arguments_.size();
-
-  for (size_t arg_n = 0; arg_n < arg_count; ++arg_n) {
-    tfm::format(stream, "%s arg%d = 0;\n", largest_arg_type(arg_n), arg_n + 1);
-  }
-
-  for (size_t loc_n = 0; loc_n < locations_.size(); ++loc_n) {
-    Location &location = locations_[loc_n];
-    tfm::format(stream, "if (__loc_id == %d) {\n", loc_n);
-
-    for (size_t arg_n = 0; arg_n < location.arguments_.size(); ++arg_n) {
-      Argument &arg = location.arguments_[arg_n];
-      if (!arg.assign_to_local(stream, tfm::format("arg%d", arg_n + 1),
-                               bin_path_, pid))
-        return false;
-    }
-    stream << "}\n";
-  }
-  return true;
-}
-
 std::string Probe::largest_arg_type(size_t arg_n) {
   Argument *largest = nullptr;
   for (Location &location : locations_) {
@@ -177,7 +143,8 @@ std::string Probe::largest_arg_type(size_t arg_n) {
   return largest->ctype();
 }
 
-bool Probe::usdt_getarg(std::ostream &stream, const optional<int> &pid) {
+bool Probe::usdt_getarg(std::ostream &stream,
+    const std::string &fn_name, const optional<int> &pid) {
   const size_t arg_count = locations_[0].arguments_.size();
 
   if (arg_count == 0)
@@ -185,18 +152,21 @@ bool Probe::usdt_getarg(std::ostream &stream, const optional<int> &pid) {
 
   for (size_t arg_n = 0; arg_n < arg_count; ++arg_n) {
     std::string ctype = largest_arg_type(arg_n);
+    std::string cptr = tfm::format("*((%s *)dest)", ctype);
+
     tfm::format(stream,
-                "static inline %s _bpf_readarg_%s_%d(struct pt_regs *ctx) {\n"
-                "  %s result = 0x0;\n",
-                ctype, name_, arg_n + 1, ctype);
+        "static inline int _bpf_readarg_%s_%d("
+        "struct pt_regs *ctx, void *dest, size_t len) {\n"
+        "  if (len != sizeof(%s)) return -1;\n",
+        fn_name, arg_n + 1, ctype);
 
     if (locations_.size() == 1) {
       Location &location = locations_.front();
       stream << "  ";
-      if (!location.arguments_[arg_n].assign_to_local(stream, "result",
+      if (!location.arguments_[arg_n].assign_to_local(stream, cptr,
                                                       bin_path_, pid))
         return false;
-      stream << "\n";
+      stream << "\n  return 0;\n}\n";
     } else {
       stream << "  switch(ctx->ip) {\n";
       for (Location &location : locations_) {
@@ -206,15 +176,15 @@ bool Probe::usdt_getarg(std::ostream &stream, const optional<int> &pid) {
           return false;
 
         tfm::format(stream, "  case 0x%xULL: ", global_address);
-        if (!location.arguments_[arg_n].assign_to_local(stream, "result",
+        if (!location.arguments_[arg_n].assign_to_local(stream, cptr,
                                                         bin_path_, pid))
           return false;
 
-        stream << " break;\n";
+        stream << " return 0;\n";
       }
       stream << "  }\n";
+      stream << "  return -1;\n}\n";
     }
-    stream << "  return result;\n}\n";
   }
   return true;
 }
@@ -277,7 +247,7 @@ Probe *Context::get(const std::string &probe_name) const {
 bool Context::generate_usdt_args(std::ostream &stream) {
   stream << "#include <uapi/linux/ptrace.h>\n";
   for (auto &p : uprobes_) {
-    if (!p.first->usdt_getarg(stream, pid_))
+    if (!p.first->usdt_getarg(stream, p.second, pid_))
       return false;
   }
   return true;

src/cc/usdt.h

@@ -147,10 +147,7 @@ public:
   size_t num_arguments() const { return locations_.front().arguments_.size(); }
 
   uint64_t address(size_t n = 0) const { return locations_[n].address_; }
-
-  bool usdt_thunks(std::ostream &stream, const std::string &prefix);
-  bool usdt_cases(std::ostream &stream, const optional<int> &pid = nullopt);
-  bool usdt_getarg(std::ostream &stream, const optional<int> &pid = nullopt);
+  bool usdt_getarg(std::ostream &stream, const std::string &fn_name, const optional<int> &pid = nullopt);
 
   bool need_enable() const { return semaphore_ != 0x0; }
   bool enable(int pid);

tests/cc/test_usdt_probes.cc

@@ -52,13 +52,6 @@ TEST_CASE("test finding a probe in our own process", "[usdt]") {
     REQUIRE(probe->need_enable() == false);
 
     REQUIRE(a_probed_function() != 0);
-
-    std::ostringstream case_stream;
-    REQUIRE(probe->usdt_cases(case_stream));
-
-    std::string cases = case_stream.str();
-    REQUIRE(cases.find("int32_t arg1") != std::string::npos);
-    REQUIRE(cases.find("uint64_t arg2") != std::string::npos);
   }
 }
 #endif  // HAVE_SDT_HEADER
@@ -140,24 +133,6 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
       REQUIRE(probe->num_locations() == 1);
       REQUIRE(probe->num_arguments() == 3);
       REQUIRE(probe->need_enable() == true);
-
-      std::ostringstream thunks_stream;
-      REQUIRE(probe->usdt_thunks(thunks_stream, "ruby_usdt"));
-
-      std::string thunks = thunks_stream.str();
-      REQUIRE(std::count(thunks.begin(), thunks.end(), '\n') == 1);
-      REQUIRE(thunks.find("ruby_usdt_thunk_0") != std::string::npos);
-
-      std::ostringstream case_stream;
-      REQUIRE(probe->usdt_cases(case_stream));
-
-      std::string cases = case_stream.str();
-      REQUIRE(countsubs(cases, "arg1") == 2);
-      REQUIRE(countsubs(cases, "arg2") == 2);
-      REQUIRE(countsubs(cases, "arg3") == 2);
-
-      REQUIRE(countsubs(cases, "uint64_t") == 4);
-      REQUIRE(countsubs(cases, "int32_t") == 2);
     }
 
     SECTION("array creation probe") {
@@ -168,26 +143,6 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
       REQUIRE(probe->num_locations() == 7);
       REQUIRE(probe->num_arguments() == 3);
       REQUIRE(probe->need_enable() == true);
-
-      std::ostringstream thunks_stream;
-      REQUIRE(probe->usdt_thunks(thunks_stream, "ruby_usdt"));
-
-      std::string thunks = thunks_stream.str();
-      REQUIRE(std::count(thunks.begin(), thunks.end(), '\n') == 7);
-      REQUIRE(thunks.find("ruby_usdt_thunk_0") != std::string::npos);
-      REQUIRE(thunks.find("ruby_usdt_thunk_6") != std::string::npos);
-      REQUIRE(thunks.find("ruby_usdt_thunk_7") == std::string::npos);
-
-      std::ostringstream case_stream;
-      REQUIRE(probe->usdt_cases(case_stream));
-
-      std::string cases = case_stream.str();
-      REQUIRE(countsubs(cases, "arg1") == 8);
-      REQUIRE(countsubs(cases, "arg2") == 8);
-      REQUIRE(countsubs(cases, "arg3") == 8);
-
-      REQUIRE(countsubs(cases, "__loc_id") == 7);
-      REQUIRE(cases.find("int64_t arg1 =") != std::string::npos);
     }
   }