Fix 20fb64c and skip probe rewriter for bpf_probe_read

20fb64c stops the whole AST traversal if it meets a bpf_probe_read call.  I
think the original intent was to simply not rewrite the third argument, so this
commit fixes it by remembering the third argument on bpf_probe_read call
traversals and overriding TraverseStmt to skip the traversal of that argument
when we meet it later.

src/cc/frontends/clang/b_frontend_action.cc

@@ -270,14 +270,20 @@ bool ProbeVisitor::VisitVarDecl(VarDecl *D) {
   return true;
 }
 
+bool ProbeVisitor::TraverseStmt(Stmt *S) {
+  if (whitelist_.find(S) != whitelist_.end())
+    return true;
+  return RecursiveASTVisitor<ProbeVisitor>::TraverseStmt(S);
+}
+
 bool ProbeVisitor::VisitCallExpr(CallExpr *Call) {
   // Skip bpf_probe_read for the third argument if it is an AddrOf.
   if (VarDecl *V = dyn_cast<VarDecl>(Call->getCalleeDecl())) {
     if (V->getName() == "bpf_probe_read" && Call->getNumArgs() >= 3) {
       const Expr *E = Call->getArg(2)->IgnoreParenCasts();
       if (const UnaryOperator *UnaryExpr = dyn_cast<UnaryOperator>(E)) {
         if (UnaryExpr->getOpcode() == UO_AddrOf)
-          return false;
+          whitelist_.insert(E);
       }
       return true;
     }

src/cc/frontends/clang/b_frontend_action.h

@@ -91,6 +91,7 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   explicit ProbeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter,
                         std::set<clang::Decl *> &m, bool track_helpers);
   bool VisitVarDecl(clang::VarDecl *Decl);
+  bool TraverseStmt(clang::Stmt *S);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitBinaryOperator(clang::BinaryOperator *E);
   bool VisitUnaryOperator(clang::UnaryOperator *E);
@@ -109,6 +110,7 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   clang::Rewriter &rewriter_;
   std::set<clang::Decl *> fn_visited_;
   std::set<clang::Expr *> memb_visited_;
+  std::set<const clang::Stmt *> whitelist_;
   std::set<std::tuple<clang::Decl *, int>> ptregs_;
   std::set<clang::Decl *> &m_;
   clang::Decl *ctx_;

tests/python/test_clang.py

@@ -76,7 +76,7 @@ def test_probe_read2(self):
         b = BPF(text=text, debug=0)
         fn = b.load_func("count_foo", BPF.KPROBE)
 
-    def test_probe_read3(self):
+    def test_probe_read_whitelist1(self):
         text = """
 #define KBUILD_MODNAME "foo"
 #include <net/tcp.h>
@@ -90,6 +90,24 @@ def test_probe_read3(self):
     bpf_probe_read(&val, sizeof(val), &(TCP_SKB_CB(skb)->tcp_gso_size));
     return val;
 }
+"""
+        b = BPF(text=text)
+        fn = b.load_func("count_tcp", BPF.KPROBE)
+
+    def test_probe_read_whitelist2(self):
+        text = """
+#define KBUILD_MODNAME "foo"
+#include <net/tcp.h>
+int count_tcp(struct pt_regs *ctx, struct sk_buff *skb) {
+    // The below define is in net/tcp.h:
+    //    #define TCP_SKB_CB(__skb) ((struct tcp_skb_cb *)&((__skb)->cb[0]))
+    // Note that it has AddrOf in the macro, which will cause current rewriter
+    // failing below statement
+    // return TCP_SKB_CB(skb)->tcp_gso_size;
+    u16 val = 0;
+    bpf_probe_read(&val, sizeof(val), &(TCP_SKB_CB(skb)->tcp_gso_size));
+    return val + skb->protocol;
+}
 """
         b = BPF(text=text)
         fn = b.load_func("count_tcp", BPF.KPROBE)