Fix the inline replace of kprobe args

The way in which args 1+ were being replaced in the C file was
fragile. Instead, assign the registers from ptregs into the function
arguments as the first statement(s) in the body of the function.
e.g.:
int sys_clone(struct ptregs *ctx, struct request *req) {
  // do something with req
}
becomes:
int sys_clone(struct ptregs *ctx, struct request *req) {
  req = ctx->di;
  // do something with req

Fixes: #192
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/cc/frontends/clang/b_frontend_action.cc

@@ -102,6 +102,7 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
     rewriter_.InsertText(D->getLocStart(), attr);
     // remember the arg names of the current function...first one is the ctx
     fn_args_.clear();
+    string preamble = "{";
     for (auto arg : D->params()) {
       if (arg->getName() == "") {
         C.getDiagnostics().Report(arg->getLocEnd(), diag::err_expected)
@@ -109,7 +110,15 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
         return false;
       }
       fn_args_.push_back(arg);
+      if (fn_args_.size() > 1) {
+        size_t d = fn_args_.size() - 2;
+        const char *reg = calling_conv_regs[d];
+        preamble += arg->getName().str() + " = " + fn_args_[0]->getName().str() + "->" + string(reg) + ";";
+      }
     }
+    // for each trace argument, convert the variable from ptregs to something on stack
+    if (CompoundStmt *S = dyn_cast<CompoundStmt>(D->getBody()))
+      rewriter_.ReplaceText(S->getLBracLoc(), 1, preamble);
   }
   return true;
 }
@@ -278,23 +287,6 @@ bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
   return true;
 }
 
-bool BTypeVisitor::VisitDeclRefExpr(DeclRefExpr *E) {
-  auto it = std::find(fn_args_.begin() + 1, fn_args_.end(), E->getDecl());
-  if (it != fn_args_.end()) {
-    if (!rewriter_.isRewritable(E->getLocStart())) {
-      C.getDiagnostics().Report(E->getLocStart(), diag::err_expected)
-          << "use of probe argument not in a macro";
-      return false;
-    }
-    size_t d = std::distance(fn_args_.begin() + 1, it);
-    const char *reg = calling_conv_regs[d];
-    string text = "((u64)" + fn_args_[0]->getName().str() + "->" + string(reg) + ")";
-    rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
-    return true;
-  }
-  return true;
-}
-
 bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   if (!E->isAssignmentOp())
     return true;

tests/cc/test_clang.py

@@ -46,6 +46,28 @@ 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_read_keys(self):
+        text = """
+#include <uapi/linux/ptrace.h>
+#include <linux/blkdev.h>
+BPF_HASH(start, struct request *);
+int do_request(struct pt_regs *ctx, struct request *req) {
+    u64 ts = bpf_ktime_get_ns();
+    start.update(&req, &ts);
+    return 0;
+}
+
+int do_completion(struct pt_regs *ctx, struct request *req) {
+    u64 *tsp = start.lookup(&req);
+    if (tsp != 0) {
+        start.delete(&req);
+    }
+    return 0;
+}
+"""
+        b = BPF(text=text, debug=0)
+        fns = b.load_funcs(BPF.KPROBE)
+
     def test_sscanf(self):
         text = """
 BPF_TABLE("hash", int, struct { u64 a; u64 b; u64 c:36; u64 d:28; struct { u32 a; u32 b; } s; }, stats, 10);