From 74adb7349802e535f1f05b2d311e13df0a3cd9bc Mon Sep 17 00:00:00 2001
From: Brenden Blanco <bblanco@plumgrid.com>
Date: Tue, 29 Sep 2015 13:44:19 -0700
Subject: [PATCH] Mask function calls from influencing probe_read

Calls to functions were propagating the needs_probe state, causing map
values to require a probe_read improperly. Although there may be a use
case for this, generally it will not be needed and should be solved
differently.

Fixes: #256
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>
---
 src/cc/frontends/clang/b_frontend_action.cc | 17 ++++++++++-------
 tests/cc/test_clang.py                      | 21 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/src/cc/frontends/clang/b_frontend_action.cc b/src/cc/frontends/clang/b_frontend_action.cc
index 5c52f62..5a04b67 100644
--- a/src/cc/frontends/clang/b_frontend_action.cc
+++ b/src/cc/frontends/clang/b_frontend_action.cc
@@ -94,14 +94,17 @@ bool BMapDeclVisitor::VisitBuiltinType(const BuiltinType *T) {
   return true;
 }
 
-class ProbeChecker : public clang::RecursiveASTVisitor<ProbeChecker> {
+class ProbeChecker : public RecursiveASTVisitor<ProbeChecker> {
  public:
   explicit ProbeChecker(Expr *arg, const set<Decl *> &ptregs)
       : needs_probe_(false), ptregs_(ptregs) {
     if (arg)
       TraverseStmt(arg);
   }
-  bool VisitDeclRefExpr(clang::DeclRefExpr *E) {
+  bool VisitCallExpr(CallExpr *E) {
+    return false;
+  }
+  bool VisitDeclRefExpr(DeclRefExpr *E) {
     if (ptregs_.find(E->getDecl()) != ptregs_.end())
       needs_probe_ = true;
     return true;
@@ -113,10 +116,10 @@ class ProbeChecker : public clang::RecursiveASTVisitor<ProbeChecker> {
 };
 
 // Visit a piece of the AST and mark it as needing probe reads
-class ProbeSetter : public clang::RecursiveASTVisitor<ProbeSetter> {
+class ProbeSetter : public RecursiveASTVisitor<ProbeSetter> {
  public:
   explicit ProbeSetter(set<Decl *> *ptregs) : ptregs_(ptregs) {}
-  bool VisitDeclRefExpr(clang::DeclRefExpr *E) {
+  bool VisitDeclRefExpr(DeclRefExpr *E) {
     ptregs_->insert(E->getDecl());
     return true;
   }
@@ -161,7 +164,7 @@ bool ProbeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   return true;
 }
 bool ProbeVisitor::VisitUnaryOperator(UnaryOperator *E) {
-  if (E->getOpcode() != UO_Deref)
+  if (E->getOpcode() == UO_AddrOf)
     return true;
   if (memb_visited_.find(E) != memb_visited_.end())
     return true;
@@ -559,10 +562,10 @@ bool BTypeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
   return true;
 }
 
-ProbeConsumer::ProbeConsumer(clang::ASTContext &C, Rewriter &rewriter)
+ProbeConsumer::ProbeConsumer(ASTContext &C, Rewriter &rewriter)
     : visitor_(rewriter) {}
 
-bool ProbeConsumer::HandleTopLevelDecl(clang::DeclGroupRef Group) {
+bool ProbeConsumer::HandleTopLevelDecl(DeclGroupRef Group) {
   for (auto D : Group) {
     if (FunctionDecl *F = dyn_cast<FunctionDecl>(D)) {
       if (F->isExternallyVisible() && F->hasBody()) {
diff --git a/tests/cc/test_clang.py b/tests/cc/test_clang.py
index 312d4ee..ad102e1 100755
--- a/tests/cc/test_clang.py
+++ b/tests/cc/test_clang.py
@@ -214,6 +214,27 @@ int kprobe__sys_open(struct pt_regs *ctx, const char *filename,
 };
 """)
 
+    def test_task_switch(self):
+        b = BPF(text="""
+#include <uapi/linux/ptrace.h>
+#include <linux/sched.h>
+struct key_t {
+  u32 prev_pid;
+  u32 curr_pid;
+};
+BPF_TABLE("hash", struct key_t, u64, stats, 1024);
+int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) {
+  struct key_t key = {};
+  u64 zero = 0, *val;
+  key.curr_pid = bpf_get_current_pid_tgid();
+  key.prev_pid = prev->pid;
+
+  val = stats.lookup_or_init(&key, &zero);
+  (*val)++;
+  return 0;
+}
+""")
+
 
 if __name__ == "__main__":
     main()
-- 
2.7.4