skip probe rewriter for bpf_probe_read()

bpf_probe_read() is often used to access pointees in bpf programs.
Recent rewriter has become smarter so a lot of bpf_probe_read()
can be replaced with simple pointer/member access.

In certain cases, bpf_probe_read() is still preferred though.
For example, kernel net/tcp.h defined TCP_SKB_CB as below
  #define TCP_SKB_CB(__skb)	((struct tcp_skb_cb *)&((__skb)->cb[0]))
User can use below to access tcp_gso_size of a skb data structure.
  TCP_SKB_CB(skb)->tcp_gso_size
The rewriter will fail as it attempts to rewrite (__skb)->cb[0].

Instead of chasing down to prevent exactly the above pattern,
this patch detects function bpf_probe_read() in ProbeVisitor and
will skip it so bpf_probe_read()'s third parameter is a AddrOf.
This can also help other cases where rewriter is not
capable and user used bpf_probe_read() as the workaround.

Also fixed tcptop.py to use direct assignment instead of
bpf_probe_read. Otherwise, rewriter will actually rewrite
src address reference inside the bpf_probe_read().
Signed-off-by: Yonghong Song <yhs@fb.com>

src/cc/frontends/clang/b_frontend_action.cc

@@ -242,7 +242,7 @@ bool ProbeVisitor::assignsExtPtr(Expr *E, int *nbAddrOf) {
 
           if (memb_name == "lookup" || memb_name == "lookup_or_init") {
             if (m_.find(Ref->getDecl()) != m_.end()) {
-            // Retrieved an ext. pointer from a map, mark LHS as ext. pointer.
+              // Retrieved an ext. pointer from a map, mark LHS as ext. pointer.
               // Pointers from maps always need a single dereference to get the
               // actual value.  The value may be an external pointer but cannot
               // be a pointer to an external pointer as the verifier prohibits
@@ -269,7 +269,20 @@ bool ProbeVisitor::VisitVarDecl(VarDecl *D) {
   }
   return true;
 }
+
 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;
+      }
+      return true;
+    }
+  }
+
   if (FunctionDecl *F = dyn_cast<FunctionDecl>(Call->getCalleeDecl())) {
     if (F->hasBody()) {
       unsigned i = 0;

tests/python/test_clang.py

@@ -76,6 +76,24 @@ int count_foo(struct pt_regs *ctx, unsigned long a, unsigned long b) {
         b = BPF(text=text, debug=0)
         fn = b.load_func("count_foo", BPF.KPROBE)
 
+    def test_probe_read3(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;
+}
+"""
+        b = BPF(text=text)
+        fn = b.load_func("count_tcp", BPF.KPROBE)
+
     def test_probe_read_keys(self):
         text = """
 #include <uapi/linux/ptrace.h>

tools/tcptop.py

@@ -118,14 +118,10 @@ int kprobe__tcp_sendmsg(struct pt_regs *ctx, struct sock *sk,
     } else if (family == AF_INET6) {
         struct ipv6_key_t ipv6_key = {.pid = pid};
 
-        bpf_probe_read(&ipv6_key.saddr0, sizeof(ipv6_key.saddr0),
-            &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[0]);
-        bpf_probe_read(&ipv6_key.saddr1, sizeof(ipv6_key.saddr1),
-            &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[2]);
-        bpf_probe_read(&ipv6_key.daddr0, sizeof(ipv6_key.daddr0),
-            &sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[0]);
-        bpf_probe_read(&ipv6_key.daddr1, sizeof(ipv6_key.daddr1),
-            &sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[2]);
+        ipv6_key.saddr0 = *(u64 *)&sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[0];
+        ipv6_key.saddr1 = *(u64 *)&sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[2];
+        ipv6_key.daddr0 = *(u64 *)&sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[0];
+        ipv6_key.daddr1 = *(u64 *)&sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[2];
         ipv6_key.lport = sk->__sk_common.skc_num;
         dport = sk->__sk_common.skc_dport;
         ipv6_key.dport = ntohs(dport);
@@ -165,14 +161,10 @@ int kprobe__tcp_cleanup_rbuf(struct pt_regs *ctx, struct sock *sk, int copied)
 
     } else if (family == AF_INET6) {
         struct ipv6_key_t ipv6_key = {.pid = pid};
-        bpf_probe_read(&ipv6_key.saddr0, sizeof(ipv6_key.saddr0),
-            &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[0]);
-        bpf_probe_read(&ipv6_key.saddr1, sizeof(ipv6_key.saddr1),
-            &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[2]);
-        bpf_probe_read(&ipv6_key.daddr0, sizeof(ipv6_key.daddr0),
-            &sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[0]);
-        bpf_probe_read(&ipv6_key.daddr1, sizeof(ipv6_key.daddr1),
-            &sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[2]);
+        ipv6_key.saddr0 = *(u64 *)&sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[0];
+        ipv6_key.saddr1 = *(u64 *)&sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32[2];
+        ipv6_key.daddr0 = *(u64 *)&sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[0];
+        ipv6_key.daddr1 = *(u64 *)&sk->__sk_common.skc_v6_daddr.in6_u.u6_addr32[2];
         ipv6_key.lport = sk->__sk_common.skc_num;
         dport = sk->__sk_common.skc_dport;
         ipv6_key.dport = ntohs(dport);