Trace external pointers through maps

The bcc rewriter currently traces external pointers using ProbeVisitor in order to replace dereferences with calls to bpf_probe_read. It is, however, unable to trace them through maps. This commit remedy this for simple (yet common) cases. Through a first traversal of the Clang AST, MapVisitor looks for calls to update (and insert) to find maps with an external pointer as value. When ProbeVisitor runs, in a second time, it looks for calls to lookup (and lookup_or_init). If the map was registered as having an external pointer as value, the l-value of the lookup assignment is marked as being an external pointer. Two traversals of the Clang AST are needed because the update of a map may happen after the lookup in the AST. Therefore, the first traversal makes sure we inspect all updates before acting on lookups. To implement this two-stage traversal without parsing the AST twice, ProbeConsumer and BTypeConsumer now implement HandleTranslationUnit, which is called after a whole translation unit has been parsed.
sjtuhjh · Oct 26, 2017 · eae0acf · eae0acf
1 parent fbbe6d6
commit eae0acf
Show file tree

Hide file tree

Showing 7 changed files with 159 additions and 79 deletions.
diff --git a/examples/tracing/tcpv4connect.py b/examples/tracing/tcpv4connect.py
@@ -55,11 +55,9 @@
 
 	// pull in details
 	struct sock *skp = *skpp;
-	u32 saddr = 0, daddr = 0;
-	u16 dport = 0;
-	bpf_probe_read(&saddr, sizeof(saddr), &skp->__sk_common.skc_rcv_saddr);
-	bpf_probe_read(&daddr, sizeof(daddr), &skp->__sk_common.skc_daddr);
-	bpf_probe_read(&dport, sizeof(dport), &skp->__sk_common.skc_dport);
+	u32 saddr = skp->__sk_common.skc_rcv_saddr;
+	u32 daddr = skp->__sk_common.skc_daddr;
+	u16 dport = skp->__sk_common.skc_dport;
 
 	// output
 	bpf_trace_printk("trace_tcp4connect %x %x %d\\n", saddr, daddr, ntohs(dport));

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
@@ -108,7 +108,36 @@ class ProbeSetter : public RecursiveASTVisitor<ProbeSetter> {
   set<Decl *> *ptregs_;
 };
 
-ProbeVisitor::ProbeVisitor(ASTContext &C, Rewriter &rewriter) : C(C), rewriter_(rewriter) {}
+// Traces maps with external pointers as values. 
+class MapVisitor : public RecursiveASTVisitor<MapVisitor> {
+ public:
+  explicit MapVisitor(set<Decl *> &m) : m_(m) {}
+  bool VisitCallExpr(CallExpr *Call) {
+    if (MemberExpr *Memb = dyn_cast<MemberExpr>(Call->getCallee()->IgnoreImplicit())) {
+      StringRef memb_name = Memb->getMemberDecl()->getName();
+      if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
+        if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
+          if (!A->getName().startswith("maps"))
+            return true;
+
+          if (memb_name == "update" || memb_name == "insert") {
+            if (ProbeChecker(Call->getArg(1), ptregs_).needs_probe()) {
+              m_.insert(Ref->getDecl());
+            }
+          }
+        }
+      }
+    }
+    return true;
+  }
+  void set_ptreg(Decl *D) { ptregs_.insert(D); }
+ private:
+  set<Decl *> &m_;
+  set<clang::Decl *> ptregs_;
+};
+
+ProbeVisitor::ProbeVisitor(ASTContext &C, Rewriter &rewriter, set<Decl *> &m) :
+  C(C), rewriter_(rewriter), m_(m) {}
 
 bool ProbeVisitor::VisitVarDecl(VarDecl *Decl) {
   if (Expr *E = Decl->getInit()) {
@@ -141,6 +170,25 @@ bool ProbeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   if (ProbeChecker(E->getRHS(), ptregs_).is_transitive()) {
     ProbeSetter setter(&ptregs_);
     setter.TraverseStmt(E->getLHS());
+  } else if (E->isAssignmentOp() && E->getRHS()->getStmtClass() == Stmt::CallExprClass) {
+    CallExpr *Call = dyn_cast<CallExpr>(E->getRHS());
+    if (MemberExpr *Memb = dyn_cast<MemberExpr>(Call->getCallee()->IgnoreImplicit())) {
+      StringRef memb_name = Memb->getMemberDecl()->getName();
+      if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
+        if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
+          if (!A->getName().startswith("maps"))
+            return true;
+
+          if (memb_name == "lookup" || memb_name == "lookup_or_init") {
+            if (m_.find(Ref->getDecl()) != m_.end()) {
+            // Retrieved an external pointer from a map, mark LHS as external pointer.
+              ProbeSetter setter(&ptregs_);
+              setter.TraverseStmt(E->getLHS());
+            }
+          }
+        }
+      }
+    }
   }
   return true;
 }
@@ -752,19 +800,51 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
   return true;
 }
 
+// First traversal of AST to retrieve maps with external pointers.
+class MapConsumer : public clang::ASTConsumer {
+ public:
+  explicit MapConsumer(set<Decl *> &m) : visitor_(m) {}
+  bool HandleTopLevelDecl(DeclGroupRef Group) {
+    for (auto D : Group) {
+      if (FunctionDecl *F = dyn_cast<FunctionDecl>(D)) {
+        if (F->isExternallyVisible() && F->hasBody()) {
+          for (auto arg : F->parameters()) {
+            if (arg != F->getParamDecl(0) && !arg->getType()->isFundamentalType()) {
+              visitor_.set_ptreg(arg);
+            }
+          }
+          visitor_.TraverseDecl(D);
+        }
+      }
+    }
+    return true;
+  }
+ private:
+  MapVisitor visitor_;
+};
+
 BTypeConsumer::BTypeConsumer(ASTContext &C, BFrontendAction &fe) : visitor_(C, fe) {}
 
-bool BTypeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
-  for (auto D : Group)
-    visitor_.TraverseDecl(D);
-  return true;
+void BTypeConsumer::HandleTranslationUnit(ASTContext &Context) {
+  DeclContext::decl_iterator it;
+  DeclContext *DC = TranslationUnitDecl::castToDeclContext(Context.getTranslationUnitDecl());
+  for (it = DC->decls_begin(); it != DC->decls_end(); it++) {
+    visitor_.TraverseDecl(*it);
+  }
 }
 
-ProbeConsumer::ProbeConsumer(ASTContext &C, Rewriter &rewriter)
-    : visitor_(C, rewriter) {}
+ProbeConsumer::ProbeConsumer(ASTContext &C, Rewriter &rewriter, set<Decl *> &m)
+    : visitor_(C, rewriter, m) {}
 
-bool ProbeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
-  for (auto D : Group) {
+/**
+ * ProbeVisitor's traversal runs after an entire translation unit has been parsed.
+ * to make sure maps with external pointers have been identified.
+ */
+void ProbeConsumer::HandleTranslationUnit(ASTContext &Context) {
+  DeclContext::decl_iterator it;
+  DeclContext *DC = TranslationUnitDecl::castToDeclContext(Context.getTranslationUnitDecl());
+  for (it = DC->decls_begin(); it != DC->decls_end(); it++) {
+    Decl *D = *it;
     if (FunctionDecl *F = dyn_cast<FunctionDecl>(D)) {
       if (F->isExternallyVisible() && F->hasBody()) {
         for (auto arg : F->parameters()) {
@@ -775,7 +855,6 @@ bool ProbeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
       }
     }
   }
-  return true;
 }
 
 BFrontendAction::BFrontendAction(llvm::raw_ostream &os, unsigned flags,
@@ -810,7 +889,8 @@ void BFrontendAction::EndSourceFileAction() {
 unique_ptr<ASTConsumer> BFrontendAction::CreateASTConsumer(CompilerInstance &Compiler, llvm::StringRef InFile) {
   rewriter_->setSourceMgr(Compiler.getSourceManager(), Compiler.getLangOpts());
   vector<unique_ptr<ASTConsumer>> consumers;
-  consumers.push_back(unique_ptr<ASTConsumer>(new ProbeConsumer(Compiler.getASTContext(), *rewriter_)));
+  consumers.push_back(unique_ptr<ASTConsumer>(new MapConsumer(m_)));
+  consumers.push_back(unique_ptr<ASTConsumer>(new ProbeConsumer(Compiler.getASTContext(), *rewriter_, m_)));
   consumers.push_back(unique_ptr<ASTConsumer>(new BTypeConsumer(Compiler.getASTContext(), *this)));
   return unique_ptr<ASTConsumer>(new MultiplexConsumer(std::move(consumers)));
 }

diff --git a/src/cc/frontends/clang/b_frontend_action.h b/src/cc/frontends/clang/b_frontend_action.h
@@ -77,7 +77,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
 // Do a depth-first search to rewrite all pointers that need to be probed
 class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
  public:
-  explicit ProbeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter);
+  explicit ProbeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter, std::set<clang::Decl *> &m);
   bool VisitVarDecl(clang::VarDecl *Decl);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitBinaryOperator(clang::BinaryOperator *E);
@@ -94,22 +94,23 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   std::set<clang::Decl *> fn_visited_;
   std::set<clang::Expr *> memb_visited_;
   std::set<clang::Decl *> ptregs_;
+  std::set<clang::Decl *> &m_;
 };
 
 // A helper class to the frontend action, walks the decls
 class BTypeConsumer : public clang::ASTConsumer {
  public:
   explicit BTypeConsumer(clang::ASTContext &C, BFrontendAction &fe);
-  bool HandleTopLevelDecl(clang::DeclGroupRef Group) override;
+  void HandleTranslationUnit(clang::ASTContext &Context) override;
  private:
   BTypeVisitor visitor_;
 };
 
 // A helper class to the frontend action, walks the decls
 class ProbeConsumer : public clang::ASTConsumer {
  public:
-  ProbeConsumer(clang::ASTContext &C, clang::Rewriter &rewriter);
-  bool HandleTopLevelDecl(clang::DeclGroupRef Group) override;
+  ProbeConsumer(clang::ASTContext &C, clang::Rewriter &rewriter, std::set<clang::Decl *> &map);
+  void HandleTranslationUnit(clang::ASTContext &Context) override;
  private:
   ProbeVisitor visitor_;
 };
@@ -146,6 +147,7 @@ class BFrontendAction : public clang::ASTFrontendAction {
   std::map<std::string, clang::SourceRange> func_range_;
   FuncSource &func_src_;
   std::string &mod_src_;
+  std::set<clang::Decl *> m_;
 };
 
 }  // namespace visitor
diff --git a/tests/python/test_clang.py b/tests/python/test_clang.py
@@ -461,6 +461,36 @@ def test_update_macro_arg(self):
         t = b["act"]
         self.assertEquals(len(t), 32);
 
+    def test_ext_ptr_maps(self):
+        bpf_text = """
+#include <uapi/linux/ptrace.h>
+#include <net/sock.h>
+#include <bcc/proto.h>
+
+BPF_HASH(currsock, u32, struct sock *);
+
+int trace_entry(struct pt_regs *ctx, struct sock *sk,
+    struct sockaddr *uaddr, int addr_len) {
+    u32 pid = bpf_get_current_pid_tgid();
+    currsock.update(&pid, &sk);
+    return 0;
+};
+
+int trace_exit(struct pt_regs *ctx) {
+    u32 pid = bpf_get_current_pid_tgid();
+    struct sock **skpp;
+    skpp = currsock.lookup(&pid);
+    if (skpp) {
+        struct sock *skp = *skpp;
+        return skp->__sk_common.skc_dport;
+    }
+    return 0;
+}
+        """
+        b = BPF(text=bpf_text)
+        b.load_func("trace_entry", BPF.KPROBE)
+        b.load_func("trace_exit", BPF.KPROBE)
+
     def test_bpf_dins_pkt_rewrite(self):
         text = """
 #include <bcc/proto.h>

diff --git a/tools/hardirqs.py b/tools/hardirqs.py
@@ -86,15 +86,9 @@
     if (tsp == 0 || descp == 0) {
         return 0;   // missed start
     }
-    // Note: descp is a value from map, so '&' can be done without
-    // probe_read, but the next level irqaction * needs a probe read.
-    // Do these steps first after reading the map, otherwise some of these
-    // pointers may get pushed onto the stack and verifier will fail.
-    struct irqaction *action = 0;
-    bpf_probe_read(&action, sizeof(action), &(*descp)->action);
-    const char **namep = &action->name;
-    char *name = 0;
-    bpf_probe_read(&name, sizeof(name), namep);
+    struct irq_desc *desc = *descp;
+    struct irqaction *action = desc->action;
+    char *name = (char *)action->name;
     delta = bpf_ktime_get_ns() - *tsp;
 
     // store as sum or histogram

diff --git a/tools/tcpconnect.py b/tools/tcpconnect.py
@@ -108,18 +108,15 @@
 
     // pull in details
     struct sock *skp = *skpp;
-    u16 dport = 0;
-    bpf_probe_read(&dport, sizeof(dport), &skp->__sk_common.skc_dport);
+    u16 dport = skp->__sk_common.skc_dport;
 
     FILTER_PORT
 
     if (ipver == 4) {
         struct ipv4_data_t data4 = {.pid = pid, .ip = ipver};
         data4.ts_us = bpf_ktime_get_ns() / 1000;
-        bpf_probe_read(&data4.saddr, sizeof(u32),
-            &skp->__sk_common.skc_rcv_saddr);
-        bpf_probe_read(&data4.daddr, sizeof(u32),
-            &skp->__sk_common.skc_daddr);
+        data4.saddr = skp->__sk_common.skc_rcv_saddr;
+        data4.daddr = skp->__sk_common.skc_daddr;
         data4.dport = ntohs(dport);
         bpf_get_current_comm(&data4.task, sizeof(data4.task));
         ipv4_events.perf_submit(ctx, &data4, sizeof(data4));
@@ -128,9 +125,9 @@
         struct ipv6_data_t data6 = {.pid = pid, .ip = ipver};
         data6.ts_us = bpf_ktime_get_ns() / 1000;
         bpf_probe_read(&data6.saddr, sizeof(data6.saddr),
-            &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+            skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
         bpf_probe_read(&data6.daddr, sizeof(data6.daddr),
-            &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
+            skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
         data6.dport = ntohs(dport);
         bpf_get_current_comm(&data6.task, sizeof(data6.task));
         ipv6_events.perf_submit(ctx, &data6, sizeof(data6));