Merge pull request #218 from iovisor/bblanco_dev

Translate multiple pointer dereference into bpr_probe_read
iovisor · Sep 15, 2015 · 2d75aa0 · 2d75aa0
2 parents 4afb3f8 + 70fa0a1
commit 2d75aa0
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 23 deletions.
diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
@@ -90,6 +90,27 @@ bool BMapDeclVisitor::VisitBuiltinType(const BuiltinType *T) {
   return true;
 }
 
+class BProbeChecker : public clang::RecursiveASTVisitor<BProbeChecker> {
+ public:
+  bool VisitDeclRefExpr(clang::DeclRefExpr *E) {
+    if (E->getDecl()->hasAttr<UnavailableAttr>())
+      return false;
+    return true;
+  }
+};
+
+// Visit a piece of the AST and mark it as needing probe reads
+class BProbeSetter : public clang::RecursiveASTVisitor<BProbeSetter> {
+ public:
+  explicit BProbeSetter(ASTContext &C) : C(C) {}
+  bool VisitDeclRefExpr(clang::DeclRefExpr *E) {
+    E->getDecl()->addAttr(UnavailableAttr::CreateImplicit(C, "ptregs"));
+    return true;
+  }
+ private:
+  ASTContext &C;
+};
+
 BTypeVisitor::BTypeVisitor(ASTContext &C, Rewriter &rewriter, vector<TableDesc> &tables)
     : C(C), rewriter_(rewriter), out_(llvm::errs()), tables_(tables) {
 }
@@ -111,6 +132,7 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
       }
       fn_args_.push_back(arg);
       if (fn_args_.size() > 1) {
+        arg->addAttr(UnavailableAttr::CreateImplicit(C, "ptregs"));
         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) + ";";
@@ -260,30 +282,34 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
   return true;
 }
 
-bool BTypeVisitor::TraverseMemberExpr(MemberExpr *E) {
-  for (auto child : E->children())
-    if (!TraverseStmt(child))
-      return false;
-  if (!WalkUpFromMemberExpr(E))
-    return false;
-  return true;
-}
-
 bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
-  if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(E->getBase()->IgnoreImplicit())) {
-    auto it = std::find(fn_args_.begin() + 1, fn_args_.end(), Ref->getDecl());
-    if (it != fn_args_.end()) {
-      FieldDecl *F = dyn_cast<FieldDecl>(E->getMemberDecl());
-      string base_type = Ref->getType()->getPointeeType().getAsString();
-      string pre, post;
-      pre = "({ " + E->getType().getAsString() + " _val; memset(&_val, 0, sizeof(_val));";
-      pre += " bpf_probe_read(&_val, sizeof(_val), (u64)";
-      post = " + offsetof(" + base_type + ", " + F->getName().str() + ")";
-      post += "); _val; })";
-      rewriter_.InsertText(E->getLocStart(), pre);
-      rewriter_.ReplaceText(SourceRange(E->getOperatorLoc(), E->getLocEnd()), post);
-    }
+  if (visited_.find(E) != visited_.end()) return true;
+
+  // Checks to see if the expression references something that needs to be run
+  // through bpf_probe_read.
+  BProbeChecker checker;
+  if (checker.TraverseStmt(E))
+    return true;
+
+  Expr *base;
+  SourceLocation rhs_start, op;
+  for (MemberExpr *M = E; M; M = dyn_cast<MemberExpr>(M->getBase())) {
+    visited_.insert(M);
+    rhs_start = M->getLocEnd();
+    base = M->getBase();
+    op = M->getOperatorLoc();
+    if (M->isArrow())
+      break;
   }
+  string rhs = rewriter_.getRewrittenText(SourceRange(rhs_start, E->getLocEnd()));
+  string base_type = base->getType()->getPointeeType().getAsString();
+  string pre, post;
+  pre = "({ " + E->getType().getAsString() + " _val; memset(&_val, 0, sizeof(_val));";
+  pre += " bpf_probe_read(&_val, sizeof(_val), (u64)";
+  post = " + offsetof(" + base_type + ", " + rhs + ")";
+  post += "); _val; })";
+  rewriter_.InsertText(E->getLocStart(), pre);
+  rewriter_.ReplaceText(SourceRange(op, E->getLocEnd()), post);
   return true;
 }
 
@@ -314,6 +340,12 @@ bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
       }
     }
   }
+  // copy probe attribute from RHS to LHS if present
+  BProbeChecker checker;
+  if (!checker.TraverseStmt(E->getRHS())) {
+    BProbeSetter setter(C);
+    setter.TraverseStmt(E->getLHS());
+  }
   return true;
 }
 bool BTypeVisitor::VisitImplicitCastExpr(ImplicitCastExpr *E) {
@@ -421,6 +453,11 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       }
     }
   }
+  if (Expr *E = Decl->getInit()) {
+    BProbeChecker checker;
+    if (!checker.TraverseStmt(E))
+      Decl->addAttr(UnavailableAttr::CreateImplicit(C, "ptregs"));
+  }
   return true;
 }
 

diff --git a/src/cc/frontends/clang/b_frontend_action.h b/src/cc/frontends/clang/b_frontend_action.h
@@ -16,6 +16,7 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -62,7 +63,6 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   explicit BTypeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter,
                         std::vector<TableDesc> &tables);
   bool TraverseCallExpr(clang::CallExpr *Call);
-  bool TraverseMemberExpr(clang::MemberExpr *E);
   bool VisitFunctionDecl(clang::FunctionDecl *D);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitVarDecl(clang::VarDecl *Decl);
@@ -76,6 +76,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   llvm::raw_ostream &out_;  /// for debugging
   std::vector<TableDesc> &tables_;  /// store the open FDs
   std::vector<clang::ParmVarDecl *> fn_args_;
+  std::set<clang::Expr *> visited_;
 };
 
 // A helper class to the frontend action, walks the decls

diff --git a/tests/cc/test_clang.py b/tests/cc/test_clang.py
@@ -130,5 +130,40 @@ def test_bpf_hash(self):
 """
         b = BPF(text=text, debug=0)
 
+    def test_consecutive_probe_read(self):
+        text = """
+#include <linux/fs.h>
+#include <linux/mount.h>
+BPF_HASH(table1, struct super_block *);
+int trace_entry(struct pt_regs *ctx, struct file *file) {
+    if (!file) return 0;
+    struct vfsmount *mnt = file->f_path.mnt;
+    if (mnt) {
+        struct super_block *k = mnt->mnt_sb;
+        u64 zero = 0;
+        table1.update(&k, &zero);
+        k = mnt->mnt_sb;
+        table1.update(&k, &zero);
+    }
+
+    return 0;
+}
+"""
+        b = BPF(text=text, debug=0)
+        fn = b.load_func("trace_entry", BPF.KPROBE)
+
+    def test_nested_probe_read(self):
+        text = """
+#include <linux/fs.h>
+int trace_entry(struct pt_regs *ctx, struct file *file) {
+    if (!file) return 0;
+    const char *name = file->f_path.dentry->d_name.name;
+    bpf_trace_printk("%s\\n", name);
+    return 0;
+}
+"""
+        b = BPF(text=text, debug=0)
+        fn = b.load_func("trace_entry", BPF.KPROBE)
+
 if __name__ == "__main__":
     main()