frontends/clang: Safety check for invalid opLoc in ProbeVisitor (#667)

As reported in #664:
In cases where the AST contains a MemberExpr with an invalid opLoc,
rewrite will fail but prior to this change would segfault. Fail a little
more gracefully, even though the program will still be rejected.
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/cc/frontends/clang/b_frontend_action.cc

@@ -158,7 +158,7 @@ class ProbeSetter : public RecursiveASTVisitor<ProbeSetter> {
   set<Decl *> *ptregs_;
 };
 
-ProbeVisitor::ProbeVisitor(Rewriter &rewriter) : rewriter_(rewriter) {}
+ProbeVisitor::ProbeVisitor(ASTContext &C, Rewriter &rewriter) : C(C), rewriter_(rewriter) {}
 
 bool ProbeVisitor::VisitVarDecl(VarDecl *Decl) {
   if (Expr *E = Decl->getInit()) {
@@ -234,6 +234,10 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   }
   if (!found)
     return true;
+  if (op.isInvalid()) {
+    error(base->getLocEnd(), "internal error: opLoc is invalid while preparing probe rewrite");
+    return false;
+  }
   string rhs = rewriter_.getRewrittenText(SourceRange(rhs_start, E->getLocEnd()));
   string base_type = base->getType()->getPointeeType().getAsString();
   string pre, post;
@@ -246,6 +250,13 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   return true;
 }
 
+template <unsigned N>
+DiagnosticBuilder ProbeVisitor::error(SourceLocation loc, const char (&fmt)[N]) {
+  unsigned int diag_id = C.getDiagnostics().getCustomDiagID(DiagnosticsEngine::Error, fmt);
+  return C.getDiagnostics().Report(loc, diag_id);
+}
+
+
 BTypeVisitor::BTypeVisitor(ASTContext &C, Rewriter &rewriter, vector<TableDesc> &tables)
     : C(C), diag_(C.getDiagnostics()), rewriter_(rewriter), out_(llvm::errs()), tables_(tables) {
 }
@@ -697,7 +708,7 @@ bool BTypeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
 }
 
 ProbeConsumer::ProbeConsumer(ASTContext &C, Rewriter &rewriter)
-    : visitor_(rewriter) {}
+    : visitor_(C, rewriter) {}
 
 bool ProbeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
   for (auto D : Group) {

src/cc/frontends/clang/b_frontend_action.h

@@ -89,7 +89,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::Rewriter &rewriter);
+  explicit ProbeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter);
   bool VisitVarDecl(clang::VarDecl *Decl);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitBinaryOperator(clang::BinaryOperator *E);
@@ -97,6 +97,10 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   bool VisitMemberExpr(clang::MemberExpr *E);
   void set_ptreg(clang::Decl *D) { ptregs_.insert(D); }
  private:
+  template <unsigned N>
+  clang::DiagnosticBuilder error(clang::SourceLocation loc, const char (&fmt)[N]);
+
+  clang::ASTContext &C;
   clang::Rewriter &rewriter_;
   std::set<clang::Decl *> fn_visited_;
   std::set<clang::Expr *> memb_visited_;

tests/python/test_clang.py

@@ -212,6 +212,20 @@ int trace_entry2(struct pt_regs *ctx, int unused, struct file *file) {
         fn = b.load_func("trace_entry1", BPF.KPROBE)
         fn = b.load_func("trace_entry2", BPF.KPROBE)
 
+    def test_probe_unnamed_union_deref(self):
+        text = """
+#include <linux/mm_types.h>
+int trace(struct pt_regs *ctx, struct page *page) {
+    void *p = page->mapping;
+    return p != NULL;
+}
+"""
+        # depending on llvm, compile may pass/fail, but at least shouldn't crash
+        try:
+            b = BPF(text=text)
+        except:
+            pass
+
     def test_probe_struct_assign(self):
         b = BPF(text = """
 #include <uapi/linux/ptrace.h>