Detect external pointers from context pointers

The bcc rewriter is currently unable to detect external pointers
(i.e., to a memory address that requires calls to bpf_probe_read) if
they are not declared as arguments, e.g., if they are retrieved
through the context argument.
For example, although the two following examples translate to the
same C code in the end (the bcc rewriter translates the first into
the second), the sk pointer is recognized as an external pointer only
in the first example.

int test1(struct pt_regs *ctx, struct sock *sk) {
    // sk is correctly recognized as an external pointer.
}
int test2(struct pt_regs *ctx) {
    struct sock *sk = (struct sock *)ctx->di;
    // sk is not recognized as an external pointer.
}

This commit fixes that by detecting member dereferences of the
context argument (i.e., the first argument of externally visible
functions). It also works for the TRACEPOINT_PROBE macro.

src/cc/frontends/clang/b_frontend_action.cc

@@ -156,9 +156,10 @@ ProbeVisitor::ProbeVisitor(ASTContext &C, Rewriter &rewriter, set<Decl *> &m) :
 
 bool ProbeVisitor::VisitVarDecl(VarDecl *Decl) {
   if (Expr *E = Decl->getInit()) {
-    if (ProbeChecker(E, ptregs_).is_transitive())
+    if (ProbeChecker(E, ptregs_).is_transitive() || IsContextMemberExpr(E)) {
       set_ptreg(Decl);
     }
+  }
   return true;
 }
 bool ProbeVisitor::VisitCallExpr(CallExpr *Call) {
@@ -185,7 +186,7 @@ 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) {
+  } else if (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();
@@ -204,6 +205,9 @@ bool ProbeVisitor::VisitBinaryOperator(BinaryOperator *E) {
         }
       }
     }
+  } else if (IsContextMemberExpr(E->getRHS())) {
+    ProbeSetter setter(&ptregs_);
+    setter.TraverseStmt(E->getLHS());
   }
   return true;
 }
@@ -262,6 +266,40 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   return true;
 }
 
+bool ProbeVisitor::IsContextMemberExpr(Expr *E) {
+  if (!E->getType()->isPointerType())
+    return false;
+
+  MemberExpr *Memb = dyn_cast<MemberExpr>(E->IgnoreParenCasts());
+  Expr *base;
+  SourceLocation rhs_start, member;
+  bool found = false;
+  MemberExpr *M;
+  for (M = Memb; M; M = dyn_cast<MemberExpr>(M->getBase())) {
+    memb_visited_.insert(M);
+    rhs_start = M->getLocEnd();
+    base = M->getBase();
+    member = M->getMemberLoc();
+    if (M->isArrow()) {
+      found = true;
+      break;
+    }
+  }
+  if (!found) {
+    return false;
+  }
+  if (member.isInvalid()) {
+    return false;
+  }
+
+  if (DeclRefExpr *base_expr = dyn_cast<DeclRefExpr>(base->IgnoreImplicit())) {
+    if (base_expr->getDecl() == ctx_) {
+      return true;
+    }
+  }
+  return false;
+}
+
 SourceRange
 ProbeVisitor::expansionRange(SourceRange range) {
   return rewriter_.getSourceMgr().getExpansionRange(range);
@@ -862,9 +900,12 @@ void BTypeConsumer::HandleTranslationUnit(ASTContext &Context) {
     if (FunctionDecl *F = dyn_cast<FunctionDecl>(D)) {
       if (fe_.is_rewritable_ext_func(F)) {
         for (auto arg : F->parameters()) {
-          if (arg != F->getParamDecl(0) && !arg->getType()->isFundamentalType())
+          if (arg == F->getParamDecl(0)) {
+            probe_visitor_.set_ctx(arg);
+          } else if (!arg->getType()->isFundamentalType()) {
             probe_visitor_.set_ptreg(arg);
           }
+        }
         probe_visitor_.TraverseDecl(D);
       }
     }

src/cc/frontends/clang/b_frontend_action.h

@@ -95,7 +95,9 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   bool VisitUnaryOperator(clang::UnaryOperator *E);
   bool VisitMemberExpr(clang::MemberExpr *E);
   void set_ptreg(clang::Decl *D) { ptregs_.insert(D); }
+  void set_ctx(clang::Decl *D) { ctx_ = D; }
  private:
+  bool IsContextMemberExpr(clang::Expr *E);
   clang::SourceRange expansionRange(clang::SourceRange range);
   template <unsigned N>
   clang::DiagnosticBuilder error(clang::SourceLocation loc, const char (&fmt)[N]);
@@ -106,6 +108,7 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   std::set<clang::Expr *> memb_visited_;
   std::set<clang::Decl *> ptregs_;
   std::set<clang::Decl *> &m_;
+  clang::Decl *ctx_;
 };
 
 // A helper class to the frontend action, walks the decls