Recognize context member dereferences despite array accesses (#1828)

* Skip instead of bailing out if MemberExpr is not rewritable

* Recognize context member dereferences despite array accesses

For example, the rewriter should recognize, in the following, that
prev is an external pointer retrieved from the context pointer,
despite the access to the second element of the args array.

struct task_struct *prev = (struct task_struct *)ctx->args[1];

The same could be done for the translation of member dereferences to
bpf_probe_read calls, but that would be a little bit more complex (to
retrieve the correct base) and there's currently no tool that would
benefit from it.

* Test for the recognition of ext ptrs from context array

* tools: remove unnecessary bpf_probe_read calls

5d656bc7 made this calls unnecessary.

src/cc/frontends/clang/b_frontend_action.cc

@@ -420,7 +420,7 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   }
 
   if (!rewriter_.isRewritable(E->getLocStart()))
-    return 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
@@ -451,19 +451,25 @@ bool ProbeVisitor::IsContextMemberExpr(Expr *E) {
   if (!E->getType()->isPointerType())
     return false;
 
-  MemberExpr *Memb = dyn_cast<MemberExpr>(E->IgnoreParenCasts());
   Expr *base;
-  SourceLocation rhs_start, member;
+  SourceLocation member;
   bool found = false;
   MemberExpr *M;
-  for (M = Memb; M; M = dyn_cast<MemberExpr>(M->getBase())) {
-    rhs_start = M->getLocEnd();
-    base = M->getBase();
+  Expr *Ex = E->IgnoreParenCasts();
+  while (Ex->getStmtClass() == Stmt::ArraySubscriptExprClass
+         || Ex->getStmtClass() == Stmt::MemberExprClass) {
+    if (Ex->getStmtClass() == Stmt::ArraySubscriptExprClass) {
+      Ex = dyn_cast<ArraySubscriptExpr>(Ex)->getBase()->IgnoreParenCasts();
+    } else if (Ex->getStmtClass() == Stmt::MemberExprClass) {
+      M = dyn_cast<MemberExpr>(Ex);
+      base = M->getBase()->IgnoreParenCasts();
       member = M->getMemberLoc();
       if (M->isArrow()) {
         found = true;
         break;
       }
+      Ex = base;
+    }
   }
   if (!found) {
     return false;
@@ -472,7 +478,7 @@ bool ProbeVisitor::IsContextMemberExpr(Expr *E) {
     return false;
   }
 
-  if (DeclRefExpr *base_expr = dyn_cast<DeclRefExpr>(base->IgnoreImplicit())) {
+  if (DeclRefExpr *base_expr = dyn_cast<DeclRefExpr>(base)) {
     if (base_expr->getDecl() == ctx_) {
       return true;
     }

tests/python/test_clang.py

@@ -88,6 +88,18 @@ int count_tcp(struct pt_regs *ctx, struct sk_buff *skb) {
         b = BPF(text=text)
         fn = b.load_func("count_tcp", BPF.KPROBE)
 
+    def test_probe_read4(self):
+        text = """
+#define KBUILD_MODNAME "foo"
+#include <net/tcp.h>
+#define _(P) ({typeof(P) val = 0; bpf_probe_read(&val, sizeof(val), &P); val;})
+int test(struct pt_regs *ctx, struct sk_buff *skb) {
+    return _(TCP_SKB_CB(skb)->tcp_gso_size) + skb->protocol;
+}
+"""
+        b = BPF(text=text)
+        fn = b.load_func("test", BPF.KPROBE)
+
     def test_probe_read_whitelist1(self):
         text = """
 #define KBUILD_MODNAME "foo"
@@ -929,6 +941,18 @@ int test(struct pt_regs *ctx) {
     sk = (struct sock *)ctx->di;
     return sk->sk_dport;
 }
+"""
+        b = BPF(text=text)
+        fn = b.load_func("test", BPF.KPROBE)
+
+    def test_probe_read_ctx_array(self):
+        text = """
+#include <linux/sched.h>
+#include <net/inet_sock.h>
+int test(struct pt_regs *ctx) {
+    struct sock *newsk = (struct sock *)PT_REGS_RC(ctx);
+    return newsk->__sk_common.skc_rcv_saddr;
+}
 """
         b = BPF(text=text)
         fn = b.load_func("test", BPF.KPROBE)

tools/runqlat.py

@@ -152,45 +152,37 @@ RAW_TRACEPOINT_PROBE(sched_wakeup)
 {
     // TP_PROTO(struct task_struct *p)
     struct task_struct *p = (struct task_struct *)ctx->args[0];
-    u32 tgid, pid;
-
-    bpf_probe_read(&tgid, sizeof(tgid), &p->tgid);
-    bpf_probe_read(&pid, sizeof(pid), &p->pid);
-    return trace_enqueue(tgid, pid);
+    return trace_enqueue(p->tgid, p->pid);
 }
 
 RAW_TRACEPOINT_PROBE(sched_wakeup_new)
 {
     // TP_PROTO(struct task_struct *p)
     struct task_struct *p = (struct task_struct *)ctx->args[0];
-    u32 tgid, pid;
-
-    bpf_probe_read(&tgid, sizeof(tgid), &p->tgid);
-    bpf_probe_read(&pid, sizeof(pid), &p->pid);
-    return trace_enqueue(tgid, pid);
+    return trace_enqueue(p->tgid, p->pid);
 }
 
 RAW_TRACEPOINT_PROBE(sched_switch)
 {
     // TP_PROTO(bool preempt, struct task_struct *prev, struct task_struct *next)
     struct task_struct *prev = (struct task_struct *)ctx->args[1];
-    struct task_struct *next= (struct task_struct *)ctx->args[2];
+    struct task_struct *next = (struct task_struct *)ctx->args[2];
     u32 pid, tgid;
     long state;
 
     // ivcsw: treat like an enqueue event and store timestamp
     bpf_probe_read(&state, sizeof(long), &prev->state);
     if (state == TASK_RUNNING) {
-        bpf_probe_read(&tgid, sizeof(prev->tgid), &prev->tgid);
-        bpf_probe_read(&pid, sizeof(prev->pid), &prev->pid);
+        tgid = prev->tgid;
+        pid = prev->pid;
         if (!(FILTER || pid == 0)) {
             u64 ts = bpf_ktime_get_ns();
             start.update(&pid, &ts);
         }
     }
 
-    bpf_probe_read(&tgid, sizeof(next->tgid), &next->tgid);
-    bpf_probe_read(&pid, sizeof(next->pid), &next->pid);
+    tgid = next->tgid;
+    pid = next->pid;
     if (FILTER || pid == 0)
         return 0;
     u64 *tsp, delta;

tools/tcpaccept.py

@@ -88,16 +88,14 @@ int kretprobe__inet_csk_accept(struct pt_regs *ctx)
 
     // pull in details
     u16 family = 0, lport = 0;
-    bpf_probe_read(&family, sizeof(family), &newsk->__sk_common.skc_family);
-    bpf_probe_read(&lport, sizeof(lport), &newsk->__sk_common.skc_num);
+    family = newsk->__sk_common.skc_family;
+    lport = newsk->__sk_common.skc_num;
 
     if (family == AF_INET) {
         struct ipv4_data_t data4 = {.pid = pid, .ip = 4};
         data4.ts_us = bpf_ktime_get_ns() / 1000;
-        bpf_probe_read(&data4.saddr, sizeof(u32),
-            &newsk->__sk_common.skc_rcv_saddr);
-        bpf_probe_read(&data4.daddr, sizeof(u32),
-            &newsk->__sk_common.skc_daddr);
+        data4.saddr = newsk->__sk_common.skc_rcv_saddr;
+        data4.daddr = newsk->__sk_common.skc_daddr;
         data4.lport = lport;
         bpf_get_current_comm(&data4.task, sizeof(data4.task));
         ipv4_events.perf_submit(ctx, &data4, sizeof(data4));