avoid probe rewriting of p->m in &(p->m) (#1831)

Fix issue #1830.

After the rewrite, the code approximately becomes
  &({type _val; bpf_probe_read(&_val, sizeof(_val), &(p->m)); _val)

Firstly the rewriting is really unnecessary, and secondly
the compilation will fail since the addressOf cannot take address
of the rvalue _val.

C standard, however, allows the addressOf operand array subscript
expression, e.g.,
  &({type _val; bpf_probe_read(&_val, sizeof(_val), &(p->m)); _val)[0]

This patch intends to fix the problem by avoiding
the rewriting in the first place of addrressOf simple member expression.
It still permits addressOf the array subscript expression.
Signed-off-by: Yonghong Song <yhs@fb.com>

src/cc/frontends/clang/b_frontend_action.cc

@@ -229,7 +229,8 @@ bool MapVisitor::VisitCallExpr(CallExpr *Call) {
 
 ProbeVisitor::ProbeVisitor(ASTContext &C, Rewriter &rewriter,
                            set<Decl *> &m, bool track_helpers) :
-  C(C), rewriter_(rewriter), m_(m), track_helpers_(track_helpers) {}
+  C(C), rewriter_(rewriter), m_(m), track_helpers_(track_helpers),
+  addrof_stmt_(nullptr), is_addrof_(false) {}
 
 bool ProbeVisitor::assignsExtPtr(Expr *E, int *nbAddrOf) {
   if (IsContextMemberExpr(E)) {
@@ -296,7 +297,12 @@ bool ProbeVisitor::VisitVarDecl(VarDecl *D) {
 bool ProbeVisitor::TraverseStmt(Stmt *S) {
   if (whitelist_.find(S) != whitelist_.end())
     return true;
-  return RecursiveASTVisitor<ProbeVisitor>::TraverseStmt(S);
+  auto ret = RecursiveASTVisitor<ProbeVisitor>::TraverseStmt(S);
+  if (addrof_stmt_ == S) {
+    addrof_stmt_ = nullptr;
+    is_addrof_ = false;
+  }
+  return ret;
 }
 
 bool ProbeVisitor::VisitCallExpr(CallExpr *Call) {
@@ -380,6 +386,10 @@ bool ProbeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   return true;
 }
 bool ProbeVisitor::VisitUnaryOperator(UnaryOperator *E) {
+  if (E->getOpcode() == UO_AddrOf) {
+    addrof_stmt_ = E;
+    is_addrof_ = true;
+  }
   if (E->getOpcode() != UO_Deref)
     return true;
   if (memb_visited_.find(E) != memb_visited_.end())
@@ -396,6 +406,13 @@ bool ProbeVisitor::VisitUnaryOperator(UnaryOperator *E) {
   rewriter_.InsertTextAfterToken(expansionLoc(sub->getLocEnd()), post);
   return true;
 }
+bool ProbeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *E) {
+  // &({..._val; bpf_probe_read(&_val, ...); _val;})[0] is permitted
+  // by C standard.
+  if (is_addrof_)
+    is_addrof_ = false;
+  return true;
+}
 bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   if (memb_visited_.find(E) != memb_visited_.end()) return true;
 
@@ -422,6 +439,13 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   if (!rewriter_.isRewritable(E->getLocStart()))
     return true;
 
+  // parent expr has addrof, skip the rewrite, set is_addrof_ to flase so
+  // it won't affect next level of indirect address
+  if (is_addrof_) {
+    is_addrof_ = false;
+    return true;
+  }
+
   /* If the base of the dereference is a call to another function, we need to
    * visit that function first to know if a rewrite is necessary (i.e., if the
    * function returns an external pointer). */

src/cc/frontends/clang/b_frontend_action.h

@@ -102,6 +102,7 @@ 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_; }
@@ -123,6 +124,8 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   clang::Decl *ctx_;
   bool track_helpers_;
   std::list<int> ptregs_returned_;
+  const clang::Stmt *addrof_stmt_;
+  bool is_addrof_;
 };
 
 // A helper class to the frontend action, walks the decls

tests/python/test_clang.py

@@ -951,7 +951,7 @@ TRACEPOINT_PROBE(skb, kfree_skb) {
 #include <net/inet_sock.h>
 int test(struct pt_regs *ctx) {
     struct sock *sk;
-    sk = (struct sock *)ctx->di;
+    sk = (struct sock *)PT_REGS_PARM1(ctx);
     return sk->sk_dport;
 }
 """
@@ -1053,6 +1053,24 @@ static inline struct tcphdr *my_skb_transport_header(struct sk_buff *skb) {
 int test(struct pt_regs *ctx, struct sock *sk, struct sk_buff *skb) {
     return my_skb_transport_header(skb)->seq;
 }
+"""
+        b = BPF(text=text)
+        fn = b.load_func("test", BPF.KPROBE)
+
+    def test_no_probe_read_addrof(self):
+        text = """
+#include <linux/sched.h>
+#include <net/inet_sock.h>
+static inline int test_help(__be16 *addr) {
+    __be16 val = 0;
+    bpf_probe_read(&val, sizeof(val), addr);
+    return val;
+}
+int test(struct pt_regs *ctx) {
+    struct sock *sk;
+    sk = (struct sock *)PT_REGS_PARM1(ctx);
+    return test_help(&sk->sk_dport);
+}
 """
         b = BPF(text=text)
         fn = b.load_func("test", BPF.KPROBE)