cc: Wrap the USDT probe context in a C API

src/cc/bcc_usdt.h

+/*
+ * Copyright (c) 2016 GitHub, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef LIBBCC_USDT_H
+#define LIBBCC_USDT_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stdint.h>
+
+void *bcc_usdt_new_frompid(int pid);
+void *bcc_usdt_new_frompath(const char *path);
+void bcc_usdt_close(void *usdt);
+
+int bcc_usdt_enable_probe(void *, const char *, const char *);
+char *bcc_usdt_genargs(void *);
+
+typedef void (*bcc_usdt_uprobe_cb)(const char *, const char *, uint64_t, int);
+void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback);
+
+#ifdef __cplusplus
+}
+#endif
+#endif

src/cc/export/helpers.h

@@ -434,5 +434,10 @@ 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); }
+
 #endif
 )********"

src/cc/usdt.cc

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 #include <sstream>
+#include <cstring>
 
 #include <fcntl.h>
 #include <sys/types.h>
@@ -29,12 +30,10 @@ namespace USDT {
 Probe::Location::Location(uint64_t addr, const char *arg_fmt) : address_(addr) {
   ArgumentParser_x64 parser(arg_fmt);
   while (!parser.done()) {
-    Argument *arg = new Argument();
-    if (!parser.parse(arg)) {
-      delete arg;  // TODO: report error
+    Argument arg;
+    if (!parser.parse(&arg))
       continue;
-    }
-    arguments_.push_back(arg);
+    arguments_.push_back(std::move(arg));
   }
 }
 
@@ -51,6 +50,16 @@ bool Probe::in_shared_object() {
   return in_shared_object_.value();
 }
 
+bool Probe::resolve_global_address(uint64_t *global, const uint64_t addr, optional<int> pid) {
+  if (in_shared_object()) {
+    return (pid &&
+	bcc_resolve_global_addr(*pid, bin_path_.c_str(), addr, global) == 0);
+  }
+
+  *global = addr;
+  return true;
+}
+
 bool Probe::lookup_semaphore_addr(uint64_t *address, int pid) {
   auto it = semaphores_.find(pid);
   if (it != semaphores_.end()) {
@@ -58,12 +67,8 @@ bool Probe::lookup_semaphore_addr(uint64_t *address, int pid) {
     return true;
   }
 
-  if (in_shared_object()) {
-    uint64_t load_address = 0x0;  // TODO
-    *address = load_address + semaphore_;
-  } else {
-    *address = semaphore_;
-  }
+  if (!resolve_global_address(address, semaphore_, pid))
+    return false;
 
   semaphores_[pid] = *address;
   return true;
@@ -100,10 +105,12 @@ bool Probe::add_to_semaphore(int pid, int16_t val) {
 }
 
 bool Probe::enable(int pid) {
+  if (enabled_semaphores_.find(pid) != enabled_semaphores_.end())
+    return true;
+
   if (!add_to_semaphore(pid, +1))
     return false;
 
-  // TODO: what happens if we enable this twice?
   enabled_semaphores_.emplace(pid, std::move(ProcStat(pid)));
   return true;
 }
@@ -137,15 +144,8 @@ bool Probe::usdt_cases(std::ostream &stream, const optional<int> &pid) {
   const size_t arg_count = locations_[0].arguments_.size();
 
   for (size_t arg_n = 0; arg_n < arg_count; ++arg_n) {
-    Argument *largest = nullptr;
-    for (Location &location : locations_) {
-      Argument *candidate = location.arguments_[arg_n];
-      if (!largest ||
-          std::abs(candidate->arg_size()) > std::abs(largest->arg_size()))
-        largest = candidate;
-    }
-
-    tfm::format(stream, "%s arg%d = 0;\n", largest->ctype(), arg_n + 1);
+    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) {
@@ -153,8 +153,8 @@ bool Probe::usdt_cases(std::ostream &stream, const optional<int> &pid) {
     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),
+      Argument &arg = location.arguments_[arg_n];
+      if (!arg.assign_to_local(stream, tfm::format("arg%d", arg_n + 1),
                                 bin_path_, pid))
         return false;
     }
@@ -163,6 +163,61 @@ bool Probe::usdt_cases(std::ostream &stream, const optional<int> &pid) {
   return true;
 }
 
+std::string Probe::largest_arg_type(size_t arg_n) {
+  Argument *largest = nullptr;
+  for (Location &location : locations_) {
+    Argument *candidate = &location.arguments_[arg_n];
+    if (!largest ||
+	std::abs(candidate->arg_size()) > std::abs(largest->arg_size()))
+      largest = candidate;
+  }
+
+  assert(largest);
+  return largest->ctype();
+}
+
+bool Probe::usdt_getarg(std::ostream &stream, const optional<int> &pid) {
+  const size_t arg_count = locations_[0].arguments_.size();
+
+  if (arg_count == 0)
+    return true;
+
+  for (size_t arg_n = 0; arg_n < arg_count; ++arg_n) {
+    std::string ctype = largest_arg_type(arg_n);
+    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);
+
+    if (locations_.size() == 1) {
+      Location &location = locations_.front();
+      stream << "  ";
+      if (!location.arguments_[arg_n].assign_to_local(
+	  stream, "result", bin_path_, pid))
+	return false;
+      stream << "\n";
+    } else {
+      stream << "  switch(ctx->ip) {\n";
+      for (Location &location : locations_) {
+	uint64_t global_address;
+
+	if (!resolve_global_address(&global_address, location.address_, pid))
+	  return false;
+
+	tfm::format(stream, "  case 0x%xULL: ", global_address);
+	if (!location.arguments_[arg_n].assign_to_local(
+	    stream, "result", bin_path_, pid))
+	  return false;
+
+	stream << " break;\n";
+      }
+      stream << "  }\n";
+    }
+    stream << "  return result;\n}\n";
+  }
+  return true;
+}
+
 void Probe::add_location(uint64_t addr, const char *fmt) {
   locations_.emplace_back(addr, fmt);
 }
@@ -210,7 +265,7 @@ std::string Context::resolve_bin_path(const std::string &bin_path) {
   return result;
 }
 
-Probe *Context::find_probe(const std::string &probe_name) {
+Probe *Context::get(const std::string &probe_name) const {
   for (Probe *p : probes_) {
     if (p->name_ == probe_name)
       return p;
@@ -218,6 +273,41 @@ Probe *Context::find_probe(const std::string &probe_name) {
   return nullptr;
 }
 
+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_))
+      return false;
+  }
+  return true;
+}
+
+bool Context::enable_probe(const std::string &probe_name, const std::string &fn_name) {
+  Probe *p = get(probe_name);
+  if (!p)
+    return false;
+
+  if (p->need_enable()) {
+    if (!pid_ || !p->enable(pid_.value()))
+      return false;
+  }
+
+  uprobes_.emplace_back(p, fn_name);
+  return true;
+}
+
+void Context::each_uprobe(each_uprobe_cb callback) {
+  for (auto &p : uprobes_) {
+    for (Probe::Location &loc : p.first->locations_) {
+      callback(
+	  p.first->bin_path_.c_str(),
+	  p.second.c_str(),
+	  loc.address_,
+	  pid_.value_or(-1));
+    }
+  }
+}
+
 Context::Context(const std::string &bin_path) : loaded_(false) {
   std::string full_path = resolve_bin_path(bin_path);
   if (!full_path.empty()) {
@@ -226,8 +316,63 @@ Context::Context(const std::string &bin_path) : loaded_(false) {
   }
 }
 
-Context::Context(int pid) : loaded_(false) {
+Context::Context(int pid) : pid_(pid), loaded_(false) {
   if (bcc_procutils_each_module(pid, _each_module, this) == 0)
     loaded_ = true;
 }
+
+Context::~Context() {
+  for (Probe *p : probes_) {
+    if (pid_ && p->enabled())
+      p->disable(pid_.value());
+    delete p;
+  }
+}
+
+}
+
+extern "C" {
+#include "bcc_usdt.h"
+
+void *bcc_usdt_new_frompid(int pid) {
+  USDT::Context *ctx = new USDT::Context(pid);
+  if (!ctx->loaded()) {
+    delete ctx;
+    return nullptr;
+  }
+  return static_cast<void *>(ctx);
+}
+
+void *bcc_usdt_new_frompath(const char *path) {
+  USDT::Context *ctx = new USDT::Context(path);
+  if (!ctx->loaded()) {
+    delete ctx;
+    return nullptr;
+  }
+  return static_cast<void *>(ctx);
+}
+
+void bcc_usdt_close(void *usdt) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  delete ctx;
+}
+
+int bcc_usdt_enable_probe(void *usdt, const char *probe_name, const char *fn_name) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  return ctx->enable_probe(probe_name, fn_name) ? 0 : -1;
+}
+
+char *bcc_usdt_genargs(void *usdt) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  std::ostringstream stream;
+  if (!ctx->generate_usdt_args(stream))
+    return nullptr;
+  return strdup(stream.str().c_str());
+}
+
+void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  ctx->each_uprobe(callback);
+}
+
 }

src/cc/usdt.h

@@ -122,7 +122,7 @@ class Probe {
 
   struct Location {
     uint64_t address_;
-    std::vector<Argument *> arguments_;
+    std::vector<Argument> arguments_;
     Location(uint64_t addr, const char *arg_fmt);
   };
 
@@ -131,7 +131,10 @@ class Probe {
   std::unordered_map<int, ProcStat> enabled_semaphores_;
   optional<bool> in_shared_object_;
 
+  std::string largest_arg_type(size_t arg_n);
+
   bool add_to_semaphore(int pid, int16_t val);
+  bool resolve_global_address(uint64_t *global, const uint64_t addr, optional<int> pid);
   bool lookup_semaphore_addr(uint64_t *address, int pid);
   void add_location(uint64_t addr, const char *fmt);
 
@@ -142,12 +145,17 @@ public:
   size_t num_locations() const { return locations_.size(); }
   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 need_enable() const { return semaphore_ != 0x0; }
   bool enable(int pid);
   bool disable(int pid);
+  bool enabled() const { return !enabled_semaphores_.empty(); }
+
 
   bool in_shared_object();
   const std::string &name() { return name_; }
@@ -159,6 +167,8 @@ public:
 
 class Context {
   std::vector<Probe *> probes_;
+  std::vector<std::pair<Probe *, std::string>> uprobes_;
+  optional<int> pid_;
   bool loaded_;
 
   static void _each_probe(const char *binpath, const struct bcc_elf_usdt *probe,
@@ -171,9 +181,19 @@ class Context {
 public:
   Context(const std::string &bin_path);
   Context(int pid);
+  ~Context();
 
+  optional<int> pid() const { return pid_; }
   bool loaded() const { return loaded_; }
   size_t num_probes() const { return probes_.size(); }
-  Probe *find_probe(const std::string &probe_name);
+
+  Probe *get(const std::string &probe_name) const;
+  Probe *get(int pos) const { return probes_[pos]; }
+
+  bool enable_probe(const std::string &probe_name, const std::string &fn_name);
+  bool generate_usdt_args(std::ostream &stream);
+
+  typedef void (*each_uprobe_cb)(const char *, const char *, uint64_t, int);
+  void each_uprobe(each_uprobe_cb callback);
 };
 }

src/cc/usdt_args.cc

@@ -55,37 +55,35 @@ bool Argument::assign_to_local(std::ostream &stream,
                                const std::string &binpath,
                                const optional<int> &pid) const {
   if (constant_) {
-    tfm::format(stream, "%s = %d;\n", local_name, *constant_);
+    tfm::format(stream, "%s = %d;", local_name, *constant_);
     return true;
   }
 
   if (!deref_offset_) {
-    tfm::format(stream, "%s = (%s)ctx->%s;\n", local_name, ctype(),
+    tfm::format(stream, "%s = (%s)ctx->%s;", local_name, ctype(),
                 *register_name_);
     return true;
   }
 
   if (deref_offset_ && !deref_ident_) {
     tfm::format(stream,
-                "{\n"
-                "    u64 __temp = ctx->%s + (%d);\n"
-                "    bpf_probe_read(&%s, sizeof(%s), (void *)__temp);\n"
-                "}\n",
-                *register_name_, *deref_offset_, local_name, local_name);
+                "{ u64 __addr = ctx->%s + (%d); %s __res = 0x0; "
+                "bpf_probe_read(&__res, sizeof(__res), (void *)__addr); "
+				"%s = __res; }",
+                *register_name_, *deref_offset_, ctype(), local_name);
     return true;
   }
 
-  if (deref_offset_ && deref_ident_) {
+  if (deref_offset_ && deref_ident_ && *register_name_ == "ip") {
     uint64_t global_address;
     if (!get_global_address(&global_address, binpath, pid))
       return false;
 
     tfm::format(stream,
-                "{\n"
-                "    u64 __temp = 0x%xull + %d;\n"
-                "    bpf_probe_read(&%s, sizeof(%s), (void *)__temp);\n"
-                "}\n",
-                global_address, *deref_offset_, local_name, local_name);
+                "{ u64 __addr = 0x%xull + %d; %s __res = 0x0; "
+                "bpf_probe_read(&__res, sizeof(__res), (void *)__addr); "
+				"%s = __res; }",
+                global_address, *deref_offset_, ctype(), local_name);
     return true;
   }
 

tests/cc/test_usdt_probes.cc

@@ -39,7 +39,7 @@ TEST_CASE("test finding a probe in our own process", "[usdt]") {
   REQUIRE(ctx.num_probes() >= 1);
 
   SECTION("our test probe") {
-    USDT::Probe *probe = ctx.find_probe("sample_probe_1");
+    USDT::Probe *probe = ctx.get("sample_probe_1");
     REQUIRE(probe != nullptr);
 
     REQUIRE(probe->in_shared_object() == false);
@@ -115,7 +115,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     mri_probe_count = ctx.num_probes();
 
     SECTION("GC static probe") {
-      USDT::Probe *probe = ctx.find_probe("gc__mark__begin");
+      USDT::Probe *probe = ctx.get("gc__mark__begin");
       REQUIRE(probe != nullptr);
 
       REQUIRE(probe->in_shared_object() == true);
@@ -129,7 +129,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     }
 
     SECTION("object creation probe") {
-      USDT::Probe *probe = ctx.find_probe("object__create");
+      USDT::Probe *probe = ctx.get("object__create");
       REQUIRE(probe != nullptr);
 
       REQUIRE(probe->in_shared_object() == true);
@@ -161,7 +161,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     }
 
     SECTION("array creation probe") {
-      USDT::Probe *probe = ctx.find_probe("array__create");
+      USDT::Probe *probe = ctx.get("array__create");
       REQUIRE(probe != nullptr);
       REQUIRE(probe->name() == "array__create");
 
@@ -203,7 +203,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
     REQUIRE(ctx.num_probes() >= mri_probe_count);
 
     SECTION("get probe in running process") {
-      USDT::Probe *probe = ctx.find_probe("gc__mark__begin");
+      USDT::Probe *probe = ctx.get("gc__mark__begin");
       REQUIRE(probe != nullptr);
 
       REQUIRE(probe->in_shared_object() == true);