Merge pull request #196 from iovisor/bblanco_dev2

Add debug and fix the inline replace of kprobe args

src/cc/bpf_module.cc

@@ -295,7 +295,7 @@ unique_ptr<ExecutionEngine> BPFModule::finalize_rw(unique_ptr<Module> m) {
 
 // load an entire c file as a module
 int BPFModule::load_cfile(const string &file, bool in_memory) {
-  clang_loader_ = make_unique<ClangLoader>(&*ctx_);
+  clang_loader_ = make_unique<ClangLoader>(&*ctx_, flags_);
   if (clang_loader_->parse(&mod_, &tables_, file, in_memory))
     return -1;
   return 0;
@@ -307,7 +307,7 @@ int BPFModule::load_cfile(const string &file, bool in_memory) {
 // Load in a pre-built list of functions into the initial Module object, then
 // build an ExecutionEngine.
 int BPFModule::load_includes(const string &tmpfile) {
-  clang_loader_ = make_unique<ClangLoader>(&*ctx_);
+  clang_loader_ = make_unique<ClangLoader>(&*ctx_, flags_);
   if (clang_loader_->parse(&mod_, &tables_, tmpfile, false))
     return -1;
   return 0;
@@ -652,7 +652,7 @@ int BPFModule::load_b(const string &filename, const string &proto_filename) {
   if (int rc = load_includes(BCC_INSTALL_PREFIX "/share/bcc/include/bcc/helpers.h"))
     return rc;
 
-  b_loader_.reset(new BLoader);
+  b_loader_.reset(new BLoader(flags_));
   if (int rc = b_loader_->parse(&*mod_, filename, proto_filename, &tables_))
     return rc;
   if (int rc = annotate())

src/cc/frontends/b/loader.cc

@@ -27,7 +27,8 @@ using std::vector;
 
 namespace ebpf {
 
-BLoader::BLoader() {
+BLoader::BLoader(unsigned flags) : flags_(flags) {
+  (void)flags_;
 }
 
 BLoader::~BLoader() {

src/cc/frontends/b/loader.h

@@ -35,11 +35,12 @@ class CodegenLLVM;
 
 class BLoader {
  public:
-  BLoader();
+  explicit BLoader(unsigned flags);
   ~BLoader();
   int parse(llvm::Module *mod, const std::string &filename, const std::string &proto_filename,
             std::unique_ptr<std::vector<TableDesc>> *tables);
  private:
+  unsigned flags_;
   std::unique_ptr<cc::Parser> parser_;
   std::unique_ptr<cc::Parser> proto_parser_;
   std::unique_ptr<cc::CodegenLLVM> codegen_;

src/cc/frontends/clang/b_frontend_action.cc

@@ -102,6 +102,7 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
     rewriter_.InsertText(D->getLocStart(), attr);
     // remember the arg names of the current function...first one is the ctx
     fn_args_.clear();
+    string preamble = "{";
     for (auto arg : D->params()) {
       if (arg->getName() == "") {
         C.getDiagnostics().Report(arg->getLocEnd(), diag::err_expected)
@@ -109,7 +110,15 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
         return false;
       }
       fn_args_.push_back(arg);
+      if (fn_args_.size() > 1) {
+        size_t d = fn_args_.size() - 2;
+        const char *reg = calling_conv_regs[d];
+        preamble += arg->getName().str() + " = " + fn_args_[0]->getName().str() + "->" + string(reg) + ";";
+      }
     }
+    // for each trace argument, convert the variable from ptregs to something on stack
+    if (CompoundStmt *S = dyn_cast<CompoundStmt>(D->getBody()))
+      rewriter_.ReplaceText(S->getLBracLoc(), 1, preamble);
   }
   return true;
 }
@@ -278,23 +287,6 @@ bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
   return true;
 }
 
-bool BTypeVisitor::VisitDeclRefExpr(DeclRefExpr *E) {
-  auto it = std::find(fn_args_.begin() + 1, fn_args_.end(), E->getDecl());
-  if (it != fn_args_.end()) {
-    if (!rewriter_.isRewritable(E->getLocStart())) {
-      C.getDiagnostics().Report(E->getLocStart(), diag::err_expected)
-          << "use of probe argument not in a macro";
-      return false;
-    }
-    size_t d = std::distance(fn_args_.begin() + 1, it);
-    const char *reg = calling_conv_regs[d];
-    string text = "((u64)" + fn_args_[0]->getName().str() + "->" + string(reg) + ")";
-    rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
-    return true;
-  }
-  return true;
-}
-
 bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   if (!E->isAssignmentOp())
     return true;
@@ -442,13 +434,14 @@ bool BTypeConsumer::HandleTopLevelDecl(DeclGroupRef D) {
   return true;
 }
 
-BFrontendAction::BFrontendAction(llvm::raw_ostream &os)
-    : rewriter_(new Rewriter), os_(os), tables_(new vector<TableDesc>) {
+BFrontendAction::BFrontendAction(llvm::raw_ostream &os, unsigned flags)
+    : os_(os), flags_(flags), rewriter_(new Rewriter), tables_(new vector<TableDesc>) {
 }
 
 void BFrontendAction::EndSourceFileAction() {
   // uncomment to see rewritten source
-  //rewriter_->getEditBuffer(rewriter_->getSourceMgr().getMainFileID()).write(llvm::errs());
+  if (flags_ & 0x4)
+    rewriter_->getEditBuffer(rewriter_->getSourceMgr().getMainFileID()).write(llvm::errs());
   rewriter_->getEditBuffer(rewriter_->getSourceMgr().getMainFileID()).write(os_);
   os_.flush();
 }

src/cc/frontends/clang/b_frontend_action.h

@@ -67,7 +67,6 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitVarDecl(clang::VarDecl *Decl);
   bool VisitMemberExpr(clang::MemberExpr *E);
-  bool VisitDeclRefExpr(clang::DeclRefExpr *E);
   bool VisitBinaryOperator(clang::BinaryOperator *E);
   bool VisitImplicitCastExpr(clang::ImplicitCastExpr *E);
 
@@ -96,7 +95,7 @@ class BFrontendAction : public clang::ASTFrontendAction {
  public:
   // Initialize with the output stream where the new source file contents
   // should be written.
-  explicit BFrontendAction(llvm::raw_ostream &os);
+  BFrontendAction(llvm::raw_ostream &os, unsigned flags);
 
   // Called by clang when the AST has been completed, here the output stream
   // will be flushed.
@@ -108,8 +107,9 @@ class BFrontendAction : public clang::ASTFrontendAction {
   // take ownership of the table-to-fd mapping data structure
   std::unique_ptr<std::vector<TableDesc>> take_tables() { return move(tables_); }
  private:
-  std::unique_ptr<clang::Rewriter> rewriter_;
   llvm::raw_ostream &os_;
+  unsigned flags_;
+  std::unique_ptr<clang::Rewriter> rewriter_;
   std::unique_ptr<std::vector<TableDesc>> tables_;
 };
 

src/cc/frontends/clang/loader.cc

@@ -58,8 +58,8 @@ using std::vector;
 
 namespace ebpf {
 
-ClangLoader::ClangLoader(llvm::LLVMContext *ctx)
-    : ctx_(ctx)
+ClangLoader::ClangLoader(llvm::LLVMContext *ctx, unsigned flags)
+    : ctx_(ctx), flags_(flags)
 {}
 
 ClangLoader::~ClangLoader() {}
@@ -159,7 +159,7 @@ int ClangLoader::parse(unique_ptr<llvm::Module> *mod, unique_ptr<vector<TableDes
   // capture the rewritten c file
   string out_str;
   llvm::raw_string_ostream os(out_str);
-  BFrontendAction bact(os);
+  BFrontendAction bact(os, flags_);
   if (!compiler1.ExecuteAction(bact))
     return -1;
   // this contains the open FDs

src/cc/frontends/clang/loader.h

@@ -36,12 +36,13 @@ class CodegenLLVM;
 
 class ClangLoader {
  public:
-  explicit ClangLoader(llvm::LLVMContext *ctx);
+  explicit ClangLoader(llvm::LLVMContext *ctx, unsigned flags);
   ~ClangLoader();
   int parse(std::unique_ptr<llvm::Module> *mod, std::unique_ptr<std::vector<TableDesc>> *tables,
             const std::string &file, bool in_memory);
  private:
   llvm::LLVMContext *ctx_;
+  unsigned flags_;
 };
 
 }  // namespace ebpf

tests/cc/test_clang.py

@@ -46,6 +46,28 @@ int count_foo(struct pt_regs *ctx, unsigned long a, unsigned long b) {
         b = BPF(text=text, debug=0)
         fn = b.load_func("count_foo", BPF.KPROBE)
 
+    def test_probe_read_keys(self):
+        text = """
+#include <uapi/linux/ptrace.h>
+#include <linux/blkdev.h>
+BPF_HASH(start, struct request *);
+int do_request(struct pt_regs *ctx, struct request *req) {
+    u64 ts = bpf_ktime_get_ns();
+    start.update(&req, &ts);
+    return 0;
+}
+
+int do_completion(struct pt_regs *ctx, struct request *req) {
+    u64 *tsp = start.lookup(&req);
+    if (tsp != 0) {
+        start.delete(&req);
+    }
+    return 0;
+}
+"""
+        b = BPF(text=text, debug=0)
+        fns = b.load_funcs(BPF.KPROBE)
+
     def test_sscanf(self):
         text = """
 BPF_TABLE("hash", int, struct { u64 a; u64 b; u64 c:36; u64 d:28; struct { u32 a; u32 b; } s; }, stats, 10);