usdt: resolve usdt arguments

This patch introduces initial support for USDT arguments through the
argX builtins.  For example:
`./bpftrace -e 'usdt:bin:probe { printf("%s\n", str(arg0)); }'`

Ref: https://github.com/iovisor/bpftrace/issues/33

src/ast/codegen_llvm.cpp

@@ -99,6 +99,11 @@ void CodegenLLVM::visit(Builtin &builtin)
     else // argX
     {
       int arg_num = atoi(builtin.ident.substr(3).c_str());
+      if (probetype(current_attach_point_->provider) == ProbeType::usdt) {
+        expr_ = b_.CreateUSDTReadArgument(ctx_, current_attach_point_, 
+                                          arg_num, builtin);
+        return;
+      }
       offset = arch::arg_offset(arg_num);
     }
 
@@ -919,6 +924,7 @@ void CodegenLLVM::visit(Probe &probe)
 
   // needed for uaddr() call:
   for (auto &attach_point : *probe.attach_points) {
+    current_attach_point_ = attach_point;
     path_ = attach_point->target;
     // TODO: semantic analyser should ensure targets are equal when uaddr() is used
     break;
@@ -960,6 +966,7 @@ void CodegenLLVM::visit(Probe &probe)
     int starting_time_id_ = time_id_;
 
     for (auto &attach_point : *probe.attach_points) {
+      current_attach_point_ = attach_point;
       std::string file_name;
       switch (probetype(attach_point->provider))
       {
@@ -1000,6 +1007,7 @@ void CodegenLLVM::visit(Probe &probe)
     }
   }
   bpftrace_.add_probe(probe);
+  current_attach_point_ = nullptr;
 }
 
 void CodegenLLVM::visit(Program &program)

src/ast/codegen_llvm.h

@@ -66,6 +66,7 @@ private:
   DataLayout layout_;
   Value *expr_ = nullptr;
   Value *ctx_;
+  AttachPoint *current_attach_point_ = nullptr;
   BPFtrace &bpftrace_;
   std::string probefull_;
   std::string path_;

src/ast/irbuilderbpf.cpp

+#include <iostream>
+
 #include "irbuilderbpf.h"
 #include "libbpf.h"
+#include "bcc_usdt.h"
+#include "arch/arch.h"
+#include "utils-inl.h"
 
 #include <llvm/IR/Module.h>
 
@@ -262,6 +267,62 @@ CallInst *IRBuilderBPF::CreateProbeReadStr(Value *dst, size_t size, Value *src)
   return CreateCall(probereadstr_func, {dst, getInt64(size), src}, "map_read_str");
 }
 
+Value *IRBuilderBPF::CreateUSDTReadArgument(Value *ctx, struct bcc_usdt_argument *argument, Builtin &builtin) {
+  // TODO (mmarchini): Handle base + index * scale addressing.
+  // https://github.com/iovisor/bcc/pull/988
+  if (argument->valid & BCC_USDT_ARGUMENT_INDEX_REGISTER_NAME)
+    std::cerr << "index register is not handled yet [" << argument->index_register_name << "]" << std::endl;
+  if (argument->valid & BCC_USDT_ARGUMENT_SCALE)
+    std::cerr << "scale is not handled yet [" << argument->scale << "]" << std::endl;
+  if (argument->valid & BCC_USDT_ARGUMENT_DEREF_IDENT)
+    std::cerr << "defer ident is not handled yet [" << argument->deref_ident << "]" << std::endl;
+
+  if (argument->valid & BCC_USDT_ARGUMENT_CONSTANT)
+    return getInt64(argument->constant);
+
+  Value *result = nullptr;
+  if (argument->valid & BCC_USDT_ARGUMENT_BASE_REGISTER_NAME) {
+    int offset = 0;
+    offset = arch::offset(argument->base_register_name);
+    Value* reg = CreateGEP(ctx, getInt64(offset * sizeof(uintptr_t)), "load_register");
+    AllocaInst *dst = CreateAllocaBPF(builtin.type, builtin.ident);
+    CreateProbeRead(dst, builtin.type.size, reg);
+    result = CreateLoad(dst);
+    if (argument->valid & BCC_USDT_ARGUMENT_DEREF_OFFSET) {
+      Value *ptr = CreateAdd(
+          result,
+          getInt64(argument->deref_offset));
+      CreateProbeRead(dst, builtin.type.size, ptr);
+      result = CreateLoad(dst);
+    }
+    CreateLifetimeEnd(dst);
+  }
+  return result;
+}
+
+Value *IRBuilderBPF::CreateUSDTReadArgument(Value *ctx, AttachPoint *attach_point, int arg_num, Builtin &builtin)
+{
+  struct bcc_usdt_argument argument;
+
+  void *usdt = bcc_usdt_new_frompath(attach_point->target.c_str());
+  if (usdt == nullptr) {
+    std::cerr << "couldn't load " << attach_point->target << std::endl;
+    exit(-1);
+  }
+
+  std::string provider = GetProviderFromPath(attach_point->target);
+  if (bcc_usdt_get_argument(usdt, provider.c_str(), attach_point->func.c_str(), 0, arg_num, &argument) != 0) {
+    std::cerr << "couldn't get argument " << arg_num << " for " << attach_point->target << ":"
+              << provider << ":" << attach_point->func << std::endl;
+    exit(-2);
+  }
+
+  Value *result = CreateUSDTReadArgument(ctx, &argument, builtin);
+
+  bcc_usdt_close(usdt);
+  return result;
+}
+
 CallInst *IRBuilderBPF::CreateGetNs()
 {
   // u64 ktime_get_ns()

src/ast/irbuilderbpf.h

@@ -2,6 +2,7 @@
 
 #include "ast.h"
 #include "bpftrace.h"
+#include "bcc_usdt.h"
 #include "types.h"
 
 #include <llvm/IR/IRBuilder.h>
@@ -32,6 +33,7 @@ public:
   void        CreateProbeRead(AllocaInst *dst, size_t size, Value *src);
   CallInst   *CreateProbeReadStr(AllocaInst *dst, size_t size, Value *src);
   CallInst   *CreateProbeReadStr(Value *dst, size_t size, Value *src);
+  Value      *CreateUSDTReadArgument(Value *ctx, AttachPoint *attach_point, int arg_name, Builtin &builtin);
   CallInst   *CreateGetNs();
   CallInst   *CreateGetPidTgid();
   CallInst   *CreateGetUidGid();
@@ -46,6 +48,8 @@ public:
 private:
   Module &module_;
   BPFtrace &bpftrace_;
+
+  Value      *CreateUSDTReadArgument(Value *ctx, struct bcc_usdt_argument *argument, Builtin &builtin);
 };
 
 } // namespace ast

src/attached_probe.cpp

@@ -12,6 +12,7 @@
 #include "bcc_usdt.h"
 #include "common.h"
 #include "libbpf.h"
+#include "utils-inl.h"
 #include <linux/perf_event.h>
 #include <linux/version.h>
 
@@ -340,11 +341,7 @@ void AttachedProbe::attach_usdt(int pid)
   if (err)
     throw std::runtime_error("Error finding or enabling probe: " + probe_.name);
 
-  std::string provider_name;
-  if ((i = probe_.path.rfind("/")) != std::string::npos)
-     provider_name = probe_.path.substr(i + 1);
-  else
-     provider_name = probe_.path;
+  std::string provider_name = GetProviderFromPath(probe_.path);
 
   err = bcc_usdt_get_location(ctx, provider_name.c_str(), probe_.attach_point.c_str(), 0, &loc);
   if (err)

src/utils-inl.h

+#pragma once
+
+#include "utils.h"
+
+namespace bpftrace {
+
+inline std::string GetProviderFromPath(std::string path) {
+  std::string provider;
+  int i = path.rfind("/");
+  return (i != std::string::npos) ? path.substr(i + 1) : path;
+}
+
+} // namespace bpftrace

src/utils.h

+#pragma once
+
+#include <string>
+
+namespace bpftrace {
+
+inline std::string GetProviderFromPath(std::string path);
+
+} // namespace bpftrace

tests/CMakeLists.txt

@@ -10,6 +10,7 @@ add_executable(bpftrace_test
   main.cpp
   parser.cpp
   semantic_analyser.cpp
+  utils.cpp
   ${CMAKE_SOURCE_DIR}/src/attached_probe.cpp
   ${CMAKE_SOURCE_DIR}/src/bpftrace.cpp
   ${CMAKE_SOURCE_DIR}/src/clang_parser.cpp

tests/utils.cpp

+#include "gtest/gtest.h"
+
+#include "utils-inl.h"
+
+namespace bpftrace {
+namespace test {
+
+TEST(utils, provider_name_from_path)
+{
+  EXPECT_EQ(GetProviderFromPath("/path/to/binary"), "binary");
+}
+
+TEST(ast, provider_name_from_path)
+{
+  EXPECT_EQ(GetProviderFromPath("provider"), "provider");
+}
+
+} // namespace test
+} // namespace bpftrace