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.

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));

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)));
 }

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

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>

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

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));