Fix 20fb64c and skip probe rewriter for bpf_probe_read (iovisor#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

@@ -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)