cc: Finish USDT implementation

src/cc/usdt.cc

@@ -15,7 +15,12 @@
  */
 #include <sstream>
 
+#include <fcntl.h>
+#include <sys/types.h>
+#include <unistd.h>
+
 #include "bcc_elf.h"
+#include "bcc_proc.h"
 #include "usdt.h"
 #include "vendor/tinyformat.hpp"
 
@@ -40,7 +45,84 @@ Probe::Probe(const char *bin_path, const char *provider, const char *name,
  name_(name),
  semaphore_(semaphore) {}
 
+bool Probe::in_shared_object() {
+ if (!in_shared_object_)
+ in_shared_object_ = (bcc_elf_is_shared_obj(bin_path_.c_str()) == 1);
+ return in_shared_object_.value();
+}
+
+bool Probe::lookup_semaphore_addr(uint64_t *address, int pid) {
+ auto it = semaphores_.find(pid);
+ if (it != semaphores_.end()) {
+ *address = it->second;
+ return true;
+ }
+
+ if (in_shared_object()) {
+ uint64_t load_address = 0x0; // TODO
+ *address = load_address + semaphore_;
+ } else {
+ *address = semaphore_;
+ }
+
+ semaphores_[pid] = *address;
+ return true;
+}
+
+bool Probe::add_to_semaphore(int pid, int16_t val) {
+ uint64_t address;
+ if (!lookup_semaphore_addr(&address, pid))
+ return false;
+
+ std::string procmem = tfm::format("/proc/%d/mem", pid);
+ int memfd = ::open(procmem.c_str(), O_RDWR);
+ if (memfd < 0)
+ return false;
+
+ int16_t original; // TODO: should this be unsigned?
+
+ if (::lseek(memfd, static_cast<off_t>(address), SEEK_SET) < 0 ||
+ ::read(memfd, &original, 2) != 2) {
+ ::close(memfd);
+ return false;
+ }
+
+ original = original + val;
+
+ if (::lseek(memfd, static_cast<off_t>(address), SEEK_SET) < 0 ||
+ ::write(memfd, &original, 2) != 2) {
+ ::close(memfd);
+ return false;
+ }
+
+ ::close(memfd);
+ return true;
+}
+
+bool Probe::enable(int pid) {
+ 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;
+}
+
+bool Probe::disable(int pid) {
+ auto it = enabled_semaphores_.find(pid);
+ if (it == enabled_semaphores_.end())
+ return false;
+
+ bool result = true;
+ if (!it->second.is_stale())
+ result = add_to_semaphore(pid, -1);
+
+ enabled_semaphores_.erase(it);
+ 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,
@@ -51,17 +133,19 @@ bool Probe::usdt_thunks(std::ostream &stream, const std::string &prefix) {
 }
 
 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) {
  Argument *largest = nullptr;
  for (Location &location : locations_) {
  Argument *candidate = location.arguments_[arg_n];
- if (!largest || candidate->arg_size() > largest->arg_size())
+ 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->ctype(), arg_n + 1);
  }
 
  for (size_t loc_n = 0; loc_n < locations_.size(); ++loc_n) {
@@ -89,6 +173,11 @@ void Context::_each_probe(const char *binpath, const struct bcc_elf_usdt *probe,
  ctx->add_probe(binpath, probe);
 }
 
+int Context::_each_module(const char *modpath, uint64_t, uint64_t, void *p) {
+ bcc_elf_foreach_usdt(modpath, _each_probe, p);
+ return 0;
+}
+
 void Context::add_probe(const char *binpath, const struct bcc_elf_usdt *probe) {
  Probe *found_probe = nullptr;
 
@@ -108,11 +197,37 @@ void Context::add_probe(const char *binpath, const struct bcc_elf_usdt *probe) {
  found_probe->add_location(probe->pc, probe->arg_fmt);
 }
 
-void Context::add_probes(const std::string &bin_path) {
- bcc_elf_foreach_usdt(bin_path.c_str(), _each_probe, this);
+std::string Context::resolve_bin_path(const std::string &bin_path) {
+ std::string result;
+
+ if (char *which = bcc_procutils_which(bin_path.c_str())) {
+ result = which;
+ ::free(which);
+ } else if (const char *which_so = bcc_procutils_which_so(bin_path.c_str())) {
+ result = which_so;
+ }
+
+ return result;
+}
+
+Probe *Context::find_probe(const std::string &probe_name) {
+ for (Probe *p : probes_) {
+ if (p->name_ == probe_name)
+ return p;
+ }
+ return nullptr;
 }
 
-Context::Context(const std::string &bin_path) { add_probes(bin_path); }
+Context::Context(const std::string &bin_path) : loaded_(false) {
+ std::string full_path = resolve_bin_path(bin_path);
+ if (!full_path.empty()) {
+ if (bcc_elf_foreach_usdt(full_path.c_str(), _each_probe, this) == 0)
+ loaded_ = true;
+ }
+}
 
-Context::Context(int pid) {}
+Context::Context(int pid) : loaded_(false) {
+ if (bcc_procutils_each_module(pid, _each_module, this) == 0)
+ loaded_ = true;
+}
 }

src/cc/usdt.h

@@ -19,6 +19,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "syms.h"
 #include "vendor/optional.hpp"
 
 namespace USDT {
@@ -37,7 +38,6 @@ class Argument {
 
  bool get_global_address(uint64_t *address, const std::string &binpath,
  const optional<int> &pid) const;
- static const std::unordered_map<std::string, std::string> translations_;
 
 public:
  Argument();
@@ -49,7 +49,6 @@ class Argument {
 
  int arg_size() const { return arg_size_.value_or(sizeof(void *)); }
  std::string ctype() const;
- void normalize_register_name(std::string *normalized) const;
 
  const optional<std::string> &deref_ident() const { return deref_ident_; }
  const optional<std::string> &register_name() const { return register_name_; }
@@ -64,7 +63,7 @@ class ArgumentParser {
  ssize_t cur_pos_;
 
 protected:
- virtual bool validate_register(const std::string &reg, int *reg_size) = 0;
+ virtual bool normalize_register(std::string *reg, int *reg_size) = 0;
 
  ssize_t parse_number(ssize_t pos, optional<int> *number);
  ssize_t parse_identifier(ssize_t pos, optional<std::string> *ident);
@@ -76,14 +75,40 @@ class ArgumentParser {
 
 public:
  bool parse(Argument *dest);
- bool done() { return arg_[cur_pos_] == '\0'; }
+ bool done() { return cur_pos_ < 0 || arg_[cur_pos_] == '\0'; }
 
  ArgumentParser(const char *arg) : arg_(arg), cur_pos_(0) {}
 };
 
 class ArgumentParser_x64 : public ArgumentParser {
- static const std::unordered_map<std::string, int> registers_;
- bool validate_register(const std::string &reg, int *reg_size);
+ enum Register {
+ REG_A,
+ REG_B,
+ REG_C,
+ REG_D,
+ REG_SI,
+ REG_DI,
+ REG_BP,
+ REG_SP,
+ REG_8,
+ REG_9,
+ REG_10,
+ REG_11,
+ REG_12,
+ REG_13,
+ REG_14,
+ REG_15,
+ REG_RIP,
+ };
+
+ struct RegInfo {
+ Register reg;
+ int size;
+ };
+
+ static const std::unordered_map<std::string, RegInfo> registers_;
+ bool normalize_register(std::string *reg, int *reg_size);
+ void reg_to_name(std::string *norm, Register reg);
 
 public:
  ArgumentParser_x64(const char *arg) : ArgumentParser(arg) {}
@@ -102,33 +127,53 @@ class Probe {
  };
 
  std::vector<Location> locations_;
+ std::unordered_map<int, uint64_t> semaphores_;
+ std::unordered_map<int, ProcStat> enabled_semaphores_;
+ optional<bool> in_shared_object_;
+
+ bool add_to_semaphore(int pid, int16_t val);
+ bool lookup_semaphore_addr(uint64_t *address, int pid);
+ void add_location(uint64_t addr, const char *fmt);
 
 public:
  Probe(const char *bin_path, const char *provider, const char *name,
  uint64_t semaphore);
 
- void add_location(uint64_t addr, const char *fmt);
- bool need_enable() const { return semaphore_ != 0x0; }
  size_t num_locations() const { return locations_.size(); }
+ size_t num_arguments() const { return locations_.front().arguments_.size(); }
 
  bool usdt_thunks(std::ostream &stream, const std::string &prefix);
  bool usdt_cases(std::ostream &stream, const optional<int> &pid = nullopt);
 
+ bool need_enable() const { return semaphore_ != 0x0; }
+ bool enable(int pid);
+ bool disable(int pid);
+
+ bool in_shared_object();
+ const std::string &name() { return name_; }
+ const std::string &bin_path() { return bin_path_; }
+ const std::string &provider() { return provider_; }
+
  friend class Context;
 };
 
 class Context {
  std::vector<Probe *> probes_;
+ bool loaded_;
 
  static void _each_probe(const char *binpath, const struct bcc_elf_usdt *probe,
  void *p);
+ static int _each_module(const char *modpath, uint64_t, uint64_t, void *p);
+
  void add_probe(const char *binpath, const struct bcc_elf_usdt *probe);
- void add_probes(const std::string &bin_path);
+ std::string resolve_bin_path(const std::string &bin_path);
 
 public:
  Context(const std::string &bin_path);
  Context(int pid);
 
+ bool loaded() const { return loaded_; }
  size_t num_probes() const { return probes_.size(); }
+ Probe *find_probe(const std::string &probe_name);
 };
 }