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 @@ int kretprobe__tcp_v4_connect(struct pt_regs *ctx)
 
 	// 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,13 +94,14 @@ 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_;
 };
@@ -108,8 +109,8 @@ class BTypeConsumer : public clang::ASTConsumer {
 // 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 @@ int process(struct xdp_md *ctx) {
         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 @@ int trace_completion(struct pt_regs *ctx)
     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 @@ static int trace_connect_return(struct pt_regs *ctx, short ipver)
 
     // 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 @@ static int trace_connect_return(struct pt_regs *ctx, short ipver)
         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));

tools/tcptracer.py

@@ -107,20 +107,15 @@ BPF_HASH(connectsock, u64, struct sock *);
 
 static int read_ipv4_tuple(struct ipv4_tuple_t *tuple, struct sock *skp)
 {
-  u32 saddr = 0, daddr = 0, net_ns_inum = 0;
-  u16 sport = 0, dport = 0;
-  possible_net_t skc_net;
-
-  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(&sport, sizeof(sport),
-                 &((struct inet_sock *)skp)->inet_sport);
-  bpf_probe_read(&dport, sizeof(dport), &skp->__sk_common.skc_dport);
+  u32 net_ns_inum = 0;
+  u32 saddr = skp->__sk_common.skc_rcv_saddr;
+  u32 daddr = skp->__sk_common.skc_daddr;
+  struct inet_sock *sockp = (struct inet_sock *)skp;
+  u16 sport = sockp->inet_sport;
+  u16 dport = skp->__sk_common.skc_dport;
 #ifdef CONFIG_NET_NS
-  bpf_probe_read(&skc_net, sizeof(skc_net), &skp->__sk_common.skc_net);
+  possible_net_t skc_net = skp->__sk_common.skc_net;
   bpf_probe_read(&net_ns_inum, sizeof(net_ns_inum), &skc_net.net->ns.inum);
-#else
-  net_ns_inum = 0;
 #endif
 
   ##FILTER_NETNS##
@@ -142,23 +137,18 @@ static int read_ipv4_tuple(struct ipv4_tuple_t *tuple, struct sock *skp)
 static int read_ipv6_tuple(struct ipv6_tuple_t *tuple, struct sock *skp)
 {
   u32 net_ns_inum = 0;
-  u16 sport = 0, dport = 0;
   unsigned __int128 saddr = 0, daddr = 0;
-  possible_net_t skc_net;
-
-  bpf_probe_read(&sport, sizeof(sport),
-                 &((struct inet_sock *)skp)->inet_sport);
-  bpf_probe_read(&dport, sizeof(dport), &skp->__sk_common.skc_dport);
-  bpf_probe_read(&saddr, sizeof(saddr),
-                 &skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
-  bpf_probe_read(&daddr, sizeof(daddr),
-                 &skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
+  struct inet_sock *sockp = (struct inet_sock *)skp;
+  u16 sport = sockp->inet_sport;
+  u16 dport = skp->__sk_common.skc_dport;
 #ifdef CONFIG_NET_NS
-  bpf_probe_read(&skc_net, sizeof(skc_net), &skp->__sk_common.skc_net);
+  possible_net_t skc_net = skp->__sk_common.skc_net;
   bpf_probe_read(&net_ns_inum, sizeof(net_ns_inum), &skc_net.net->ns.inum);
-#else
-  net_ns_inum = 0;
 #endif
+  bpf_probe_read(&saddr, sizeof(saddr),
+                 skp->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+  bpf_probe_read(&daddr, sizeof(daddr),
+                 skp->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
 
   ##FILTER_NETNS##
 
@@ -178,10 +168,7 @@ static int read_ipv6_tuple(struct ipv6_tuple_t *tuple, struct sock *skp)
 
 static bool check_family(struct sock *sk, u16 expected_family) {
   u64 zero = 0;
-  u16 family = 0;
-
-  bpf_probe_read(&family, sizeof(family), &sk->__sk_common.skc_family);
-
+  u16 family = sk->__sk_common.skc_family;
   return family == expected_family;
 }
 
@@ -279,15 +266,12 @@ int trace_connect_v6_return(struct pt_regs *ctx)
   return 0;
 }
 
-int trace_tcp_set_state_entry(struct pt_regs *ctx, struct sock *sk, int state)
+int trace_tcp_set_state_entry(struct pt_regs *ctx, struct sock *skp, int state)
 {
   if (state != TCP_ESTABLISHED && state != TCP_CLOSE) {
       return 0;
   }
 
-  struct sock *skp;
-  bpf_probe_read(&skp, sizeof(struct sock *), &sk);
-
   u8 ipver = 0;
   if (check_family(skp, AF_INET)) {
       ipver = 4;
@@ -367,18 +351,13 @@ int trace_tcp_set_state_entry(struct pt_regs *ctx, struct sock *sk, int state)
   return 0;
 }
 
-int trace_close_entry(struct pt_regs *ctx, struct sock *sk)
+int trace_close_entry(struct pt_regs *ctx, struct sock *skp)
 {
   u64 pid = bpf_get_current_pid_tgid();
 
   ##FILTER_PID##
 
-  // pull in details
-  struct sock *skp;
-  bpf_probe_read(&skp, sizeof(struct sock *), &sk);
-
-  u8 oldstate = 0;
-  bpf_probe_read(&oldstate, sizeof(oldstate), (u8 *)&skp->sk_state);
+  u8 oldstate = skp->sk_state;
   // Don't generate close events for connections that were never
   // established in the first place.
   if (oldstate == TCP_SYN_SENT ||
@@ -500,9 +479,9 @@ int trace_accept_return(struct pt_regs *ctx)
       evt6.ip = ipver;
 
       bpf_probe_read(&evt6.saddr, sizeof(evt6.saddr),
-                     &newsk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+                     newsk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
       bpf_probe_read(&evt6.daddr, sizeof(evt6.daddr),
-                     &newsk->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
+                     newsk->__sk_common.skc_v6_daddr.in6_u.u6_addr32);
 
       evt6.sport = lport;
       evt6.dport = ntohs(dport);