From eebd4856b02fde16ee1d4f4ba02a9045bcb5b5e5 Mon Sep 17 00:00:00 2001
From: Paul Chaignon <paul.chaignon@gmail.com>
Date: Wed, 13 Jun 2018 03:55:52 +0200
Subject: [PATCH] Fix 20fb64c and skip probe rewriter for bpf_probe_read
 (#1812)

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 |  8 +++++++-
 src/cc/frontends/clang/b_frontend_action.h  |  2 ++
 tests/python/test_clang.py                  | 20 +++++++++++++++++++-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
index 9baec7234047..d7423f829e7e 100644
--- a/src/cc/frontends/clang/b_frontend_action.cc
+++ b/src/cc/frontends/clang/b_frontend_action.cc
@@ -270,6 +270,12 @@ 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())) {
@@ -277,7 +283,7 @@ bool ProbeVisitor::VisitCallExpr(CallExpr *Call) {
       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;
     }
diff --git a/src/cc/frontends/clang/b_frontend_action.h b/src/cc/frontends/clang/b_frontend_action.h
index 7dc373c6b00b..f5c3b442f991 100644
--- a/src/cc/frontends/clang/b_frontend_action.h
+++ b/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_;
diff --git a/tests/python/test_clang.py b/tests/python/test_clang.py
index 35cabb2de061..e291a9f16e69 100755
--- a/tests/python/test_clang.py
+++ b/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)