cc: Implement symbol resolution for USDT

src/cc/usdt.cc

@@ -40,28 +40,18 @@ Probe::Probe(const char *bin_path, const char *provider, const char *name,
       name_(name),
       semaphore_(semaphore) {}
 
-const std::string &Probe::usdt_thunks(const std::string &prefix) {
-  if (!gen_thunks_.empty())
-    return gen_thunks_;
-
-  std::ostringstream stream;
+bool Probe::usdt_thunks(std::ostream &stream, const std::string &prefix) {
   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);
   }
-
-  gen_thunks_ = stream.str();
-  return gen_thunks_;
+  return true;
 }
 
-const std::string &Probe::usdt_cases(const optional<int> &pid) {
-  if (!gen_cases_.empty())
-    return gen_cases_;
-
-  std::ostringstream stream;
-  size_t arg_count = locations_[0].arguments_.size();
+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;
@@ -80,14 +70,13 @@ const std::string &Probe::usdt_cases(const optional<int> &pid) {
 
     for (size_t arg_n = 0; arg_n < location.arguments_.size(); ++arg_n) {
       Argument *arg = location.arguments_[arg_n];
-      arg->assign_to_local(stream, tfm::format("arg%d", arg_n + 1), bin_path_,
-                           pid);
+      if (!arg->assign_to_local(stream, tfm::format("arg%d", arg_n + 1),
+                                bin_path_, pid))
+        return false;
     }
     stream << "}\n";
   }
-
-  gen_cases_ = stream.str();
-  return gen_cases_;
+  return true;
 }
 
 void Probe::add_location(uint64_t addr, const char *fmt) {

src/cc/usdt.h

@@ -35,15 +35,15 @@ private:
   optional<std::string> deref_ident_;
   optional<std::string> register_name_;
 
-  uint64_t get_global_address(const std::string &binpath,
-                              const optional<int> &pid) const;
+  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();
   ~Argument();
 
-  void assign_to_local(std::ostream &stream, const std::string &local_name,
+  bool assign_to_local(std::ostream &stream, const std::string &local_name,
                        const std::string &binpath,
                        const optional<int> &pid = nullopt) const;
 
@@ -103,19 +103,16 @@ class Probe {
 
   std::vector<Location> locations_;
 
-  std::string gen_thunks_;
-  std::string gen_cases_;
-
 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 location_count() const { return locations_.size(); }
+  size_t num_locations() const { return locations_.size(); }
 
-  const std::string &usdt_thunks(const std::string &prefix);
-  const std::string &usdt_cases(const optional<int> &pid);
+  bool usdt_thunks(std::ostream &stream, const std::string &prefix);
+  bool usdt_cases(std::ostream &stream, const optional<int> &pid = nullopt);
 
   friend class Context;
 };
@@ -131,5 +128,7 @@ class Context {
 public:
   Context(const std::string &bin_path);
   Context(int pid);
+
+  size_t num_probes() const { return probes_.size(); }
 };
 }

src/cc/usdt_args.cc

@@ -15,9 +15,13 @@
  */
 #include <unordered_map>
 
+#include "syms.h"
 #include "usdt.h"
 #include "vendor/tinyformat.hpp"
 
+#include "bcc_elf.h"
+#include "bcc_syms.h"
+
 namespace USDT {
 
 Argument::Argument() {}
@@ -49,12 +53,25 @@ void Argument::normalize_register_name(std::string *normalized) const {
     normalized->assign(it->second);
 }
 
-uint64_t Argument::get_global_address(const std::string &binpath,
-                                      const optional<int> &pid) const {
-  return 0x0;
+bool Argument::get_global_address(uint64_t *address, const std::string &binpath,
+                                  const optional<int> &pid) const {
+  if (pid) {
+    return ProcSyms(*pid).resolve_name(binpath.c_str(), deref_ident_->c_str(),
+                                       address);
+  }
+
+  if (bcc_elf_is_shared_obj(binpath.c_str()) == 0) {
+    struct bcc_symbol sym = {deref_ident_->c_str(), binpath.c_str(), 0x0};
+    if (!bcc_find_symbol_addr(&sym) && sym.offset) {
+      *address = sym.offset;
+      return true;
+    }
+  }
+
+  return false;
 }
 
-void Argument::assign_to_local(std::ostream &stream,
+bool Argument::assign_to_local(std::ostream &stream,
                                const std::string &local_name,
                                const std::string &binpath,
                                const optional<int> &pid) const {
@@ -63,12 +80,12 @@ void Argument::assign_to_local(std::ostream &stream,
 
   if (constant_) {
     tfm::format(stream, "%s = %d;\n", local_name, *constant_);
-    return;
+    return true;
   }
 
   if (!deref_offset_) {
     tfm::format(stream, "%s = (%s)ctx->%s;\n", local_name, ctype(), regname);
-    return;
+    return true;
   }
 
   if (deref_offset_ && !deref_ident_) {
@@ -78,16 +95,24 @@ void Argument::assign_to_local(std::ostream &stream,
                 "    bpf_probe_read(&%s, sizeof(%s), (void *)__temp);\n"
                 "}\n",
                 regname, *deref_offset_, local_name, local_name);
-    return;
+    return true;
+  }
+
+  if (deref_offset_ && deref_ident_) {
+    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);
+    return true;
   }
 
-  tfm::format(stream,
-              "{\n"
-              "    u64 __temp = 0x%xull + %d;\n"
-              "    bpf_probe_read(&%s, sizeof(%s), (void *)__temp);\n"
-              "}\n",
-              get_global_address(binpath, pid), *deref_offset_, local_name,
-              local_name);
+  return false;
 }
 
 ssize_t ArgumentParser::parse_number(ssize_t pos, optional<int> *result) {