Rewrite array accesses as part of dereferences

Stops at any array accesses on external pointers and tries to rewrite
both the array access and the member dereference if any, in one shot.
With this commit, the following C code is rewritten properly into a
single bpf_probe_read call.

int test(struct pt_regs *ctx, const struct qstr *name) {
    return name->name[1];
}

Based on Yonghong Song's code.

2 files changed, 60 insertions, 0 deletions

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
index 593859a6..3c3c9fee 100644
--- a/src/cc/frontends/clang/b_frontend_action.cc
+++ b/src/cc/frontends/clang/b_frontend_action.cc
@@ -25,6 +25,7 @@
 #include <clang/Frontend/CompilerInstance.h>
 #include <clang/Frontend/MultiplexConsumer.h>
 #include <clang/Rewrite/Core/Rewriter.h>
+#include <clang/Lex/Lexer.h>
 
 #include "b_frontend_action.h"
 #include "bpf_module.h"
@@ -487,7 +488,64 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   rewriter_.ReplaceText(expansionRange(SourceRange(member, E->getLocEnd())), post);
   return true;
 }
+bool ProbeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *E) {
+  if (memb_visited_.find(E) != memb_visited_.end()) return true;
+  if (!ProbeChecker(E, ptregs_, track_helpers_).needs_probe())
+    return true;
+
+  // Parent expr has addrof, skip the rewrite.
+  if (is_addrof_)
+    return true;
 
+  if (!rewriter_.isRewritable(E->getLocStart()))
+    return true;
+
+  Expr *base = E->getBase();
+  Expr *idx = E->getIdx();
+  memb_visited_.insert(E);
+
+  string pre, lbracket, rbracket;
+  LangOptions opts;
+  SourceLocation lbracket_start, lbracket_end;
+  SourceRange lbracket_range;
+  pre = "({ typeof(" + E->getType().getAsString() + ") _val; __builtin_memset(&_val, 0, sizeof(_val));";
+  pre += " bpf_probe_read(&_val, sizeof(_val), (u64)(";
+  if (isMemberDereference(base)) {
+    pre += "&";
+    // If the base of the array subscript is a member dereference, we'll rewrite
+    // both at the same time.
+    addrof_stmt_ = base;
+    is_addrof_ = true;
+  }
+  rewriter_.InsertText(expansionLoc(base->getLocStart()), pre);
+
+  /* Replace left bracket and any space around it.  Since Clang doesn't provide
+   * a method to retrieve the left bracket, replace everything from the end of
+   * the base to the start of the index. */
+  lbracket = ") + (";
+  lbracket_start = Lexer::getLocForEndOfToken(base->getLocEnd(), 1,
+                                              rewriter_.getSourceMgr(),
+                                              opts).getLocWithOffset(1);
+  lbracket_end = idx->getLocStart().getLocWithOffset(-1);
+  lbracket_range = expansionRange(SourceRange(lbracket_start, lbracket_end));
+  rewriter_.ReplaceText(lbracket_range, lbracket);
+
+  rbracket = ")); _val; })";
+  rewriter_.ReplaceText(expansionLoc(E->getRBracketLoc()), 1, rbracket);
+
+  return true;
+}
+
+bool ProbeVisitor::isMemberDereference(Expr *E) {
+  if (E->IgnoreParenCasts()->getStmtClass() != Stmt::MemberExprClass)
+    return false;
+  for (MemberExpr *M = dyn_cast<MemberExpr>(E->IgnoreParenCasts()); M;
+       M = dyn_cast<MemberExpr>(M->getBase()->IgnoreParenCasts())) {
+    if (M->isArrow())
+      return true;
+  }
+  return false;
+}
 bool ProbeVisitor::IsContextMemberExpr(Expr *E) {
   if (!E->getType()->isPointerType())
     return false;
diff --git a/src/cc/frontends/clang/b_frontend_action.h b/src/cc/frontends/clang/b_frontend_action.h
index c4e94b86..703f80da 100644
--- a/src/cc/frontends/clang/b_frontend_action.h
+++ b/src/cc/frontends/clang/b_frontend_action.h
@@ -102,11 +102,13 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   bool VisitBinaryOperator(clang::BinaryOperator *E);
   bool VisitUnaryOperator(clang::UnaryOperator *E);
   bool VisitMemberExpr(clang::MemberExpr *E);
+  bool VisitArraySubscriptExpr(clang::ArraySubscriptExpr *E);
   void set_ptreg(std::tuple<clang::Decl *, int> &pt) { ptregs_.insert(pt); }
   void set_ctx(clang::Decl *D) { ctx_ = D; }
   std::set<std::tuple<clang::Decl *, int>> get_ptregs() { return ptregs_; }
  private:
   bool assignsExtPtr(clang::Expr *E, int *nbAddrOf);
+  bool isMemberDereference(clang::Expr *E);
   bool IsContextMemberExpr(clang::Expr *E);
   clang::SourceRange expansionRange(clang::SourceRange range);
   clang::SourceLocation expansionLoc(clang::SourceLocation loc);