From c48ab4b004b8515dc6d4973ce58295e059c94ea5 Mon Sep 17 00:00:00 2001
From: Brenden Blanco <bblanco@plumgrid.com>
Date: Thu, 1 Oct 2015 11:18:07 -0700
Subject: [PATCH] Don't treat fundamentally typed args as needing probe_read

The rewriter was aggressively parsing PoD types as requiring probe_read
and poisoning other decls when those arguments were used.

Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>
---
 src/cc/frontends/clang/b_frontend_action.cc | 16 +++++++++++-----
 tests/cc/test_clang.py                      | 13 +++++++++++++
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
index 5a04b67..3ec4efb 100644
--- a/src/cc/frontends/clang/b_frontend_action.cc
+++ b/src/cc/frontends/clang/b_frontend_action.cc
@@ -97,11 +97,15 @@ bool BMapDeclVisitor::VisitBuiltinType(const BuiltinType *T) {
 class ProbeChecker : public RecursiveASTVisitor<ProbeChecker> {
  public:
   explicit ProbeChecker(Expr *arg, const set<Decl *> &ptregs)
-      : needs_probe_(false), ptregs_(ptregs) {
-    if (arg)
+      : needs_probe_(false), is_transitive_(false), ptregs_(ptregs) {
+    if (arg) {
       TraverseStmt(arg);
+      if (arg->getType()->isPointerType())
+        is_transitive_ = needs_probe_;
+    }
   }
   bool VisitCallExpr(CallExpr *E) {
+    needs_probe_ = false;
     return false;
   }
   bool VisitDeclRefExpr(DeclRefExpr *E) {
@@ -110,8 +114,10 @@ class ProbeChecker : public RecursiveASTVisitor<ProbeChecker> {
     return true;
   }
   bool needs_probe() const { return needs_probe_; }
+  bool is_transitive() const { return is_transitive_; }
  private:
   bool needs_probe_;
+  bool is_transitive_;
   const set<Decl *> &ptregs_;
 };
 
@@ -131,7 +137,7 @@ ProbeVisitor::ProbeVisitor(Rewriter &rewriter) : rewriter_(rewriter) {}
 
 bool ProbeVisitor::VisitVarDecl(VarDecl *Decl) {
   if (Expr *E = Decl->getInit()) {
-    if (ProbeChecker(E, ptregs_).needs_probe())
+    if (ProbeChecker(E, ptregs_).is_transitive())
       set_ptreg(Decl);
   }
   return true;
@@ -157,7 +163,7 @@ bool ProbeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   if (!E->isAssignmentOp())
     return true;
   // copy probe attribute from RHS to LHS if present
-  if (ProbeChecker(E->getRHS(), ptregs_).needs_probe()) {
+  if (ProbeChecker(E->getRHS(), ptregs_).is_transitive()) {
     ProbeSetter setter(&ptregs_);
     setter.TraverseStmt(E->getLHS());
   }
@@ -570,7 +576,7 @@ bool ProbeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
     if (FunctionDecl *F = dyn_cast<FunctionDecl>(D)) {
       if (F->isExternallyVisible() && F->hasBody()) {
         for (auto arg : F->parameters()) {
-          if (arg != F->getParamDecl(0))
+          if (arg != F->getParamDecl(0) && !arg->getType()->isFundamentalType())
             visitor_.set_ptreg(arg);
         }
         visitor_.TraverseDecl(D);
diff --git a/tests/cc/test_clang.py b/tests/cc/test_clang.py
index ad102e1..276bd69 100755
--- a/tests/cc/test_clang.py
+++ b/tests/cc/test_clang.py
@@ -235,6 +235,19 @@ int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) {
 }
 """)
 
+    def test_probe_simple_assign(self):
+        b = BPF(text="""
+#include <uapi/linux/ptrace.h>
+#include <linux/gfp.h>
+struct leaf { size_t size; };
+BPF_HASH(simple_map, u32, struct leaf);
+int kprobe____kmalloc(struct pt_regs *ctx, size_t size) {
+    u32 pid = bpf_get_current_pid_tgid();
+    struct leaf* leaf = simple_map.lookup(&pid);
+    if (leaf)
+        leaf->size += size;
+    return 0;
+}""", debug=4)
 
 if __name__ == "__main__":
     main()
-- 
2.7.4