Fix nested dereference rewrites (#1835)

* Fix nested rewrites dereferences

When the rewriter meets a dereference of a member dereference it
fails to properly rewrite them into calls to bpf_probe_read.  The
reason is that Clang is unable to track the position of rewritten
text, but we can accommodate this by inserting text around the
dereference instead of completely rewriting it.  We are already
doing that for member dereference, but not for simple dereference.

* Test for the rewrite of nested dereferences

src/cc/frontends/clang/b_frontend_action.cc

@@ -388,12 +388,12 @@ bool ProbeVisitor::VisitUnaryOperator(UnaryOperator *E) {
   if (!ProbeChecker(sub, ptregs_, track_helpers_).needs_probe())
     return true;
   memb_visited_.insert(E);
-  string rhs = rewriter_.getRewrittenText(expansionRange(sub->getSourceRange()));
-  string text;
-  text = "({ typeof(" + E->getType().getAsString() + ") _val; __builtin_memset(&_val, 0, sizeof(_val));";
-  text += " bpf_probe_read(&_val, sizeof(_val), (u64)";
-  text += rhs + "); _val; })";
-  rewriter_.ReplaceText(expansionRange(E->getSourceRange()), text);
+  string pre, post;
+  pre = "({ typeof(" + E->getType().getAsString() + ") _val; __builtin_memset(&_val, 0, sizeof(_val));";
+  pre += " bpf_probe_read(&_val, sizeof(_val), (u64)";
+  post = "); _val; })";
+  rewriter_.ReplaceText(expansionLoc(E->getOperatorLoc()), 1, pre);
+  rewriter_.InsertTextAfterToken(expansionLoc(sub->getLocEnd()), post);
   return true;
 }
 bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
@@ -442,7 +442,7 @@ bool ProbeVisitor::VisitMemberExpr(MemberExpr *E) {
   pre = "({ typeof(" + E->getType().getAsString() + ") _val; __builtin_memset(&_val, 0, sizeof(_val));";
   pre += " bpf_probe_read(&_val, sizeof(_val), (u64)&";
   post = rhs + "); _val; })";
-  rewriter_.InsertText(E->getLocStart(), pre);
+  rewriter_.InsertText(expansionLoc(E->getLocStart()), pre);
   rewriter_.ReplaceText(expansionRange(SourceRange(member, E->getLocEnd())), post);
   return true;
 }
@@ -495,6 +495,11 @@ ProbeVisitor::expansionRange(SourceRange range) {
 #endif
 }
 
+SourceLocation
+ProbeVisitor::expansionLoc(SourceLocation loc) {
+  return rewriter_.getSourceMgr().getExpansionLoc(loc);
+}
+
 template <unsigned N>
 DiagnosticBuilder ProbeVisitor::error(SourceLocation loc, const char (&fmt)[N]) {
   unsigned int diag_id = C.getDiagnostics().getCustomDiagID(DiagnosticsEngine::Error, fmt);

src/cc/frontends/clang/b_frontend_action.h

@@ -109,6 +109,7 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
   bool assignsExtPtr(clang::Expr *E, int *nbAddrOf);
   bool IsContextMemberExpr(clang::Expr *E);
   clang::SourceRange expansionRange(clang::SourceRange range);
+  clang::SourceLocation expansionLoc(clang::SourceLocation loc);
   template <unsigned N>
   clang::DiagnosticBuilder error(clang::SourceLocation loc, const char (&fmt)[N]);
 

tests/python/test_clang.py

@@ -303,6 +303,19 @@ int trace_entry(struct pt_regs *ctx, struct file *file) {
         b = BPF(text=text, debug=0)
         fn = b.load_func("trace_entry", BPF.KPROBE)
 
+    def test_nested_probe_read_deref(self):
+        text = """
+#include <uapi/linux/ptrace.h>
+struct sock {
+    u32 *sk_daddr;
+};
+int test(struct pt_regs *ctx, struct sock *skp) {
+    return *(skp->sk_daddr);
+}
+"""
+        b = BPF(text=text)
+        fn = b.load_func("test", BPF.KPROBE)
+
     def test_char_array_probe(self):
         BPF(text="""#include <linux/blkdev.h>
 int kprobe__blk_update_request(struct pt_regs *ctx, struct request *req) {