Automate bpf_probe_read

* Rewrite accesses to args beyond the first to use bpf_probe_read
 - Support struct access
 - Support POD type accesses
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

examples/task_switch.c

@@ -7,13 +7,12 @@ struct key_t {
 };
 // map_type, key_type, leaf_type, table_name, num_entry
 BPF_TABLE("hash", struct key_t, u64, stats, 1024);
-int count_sched(struct pt_regs *ctx) {
+int count_sched(struct pt_regs *ctx, struct task_struct *prev) {
   struct key_t key = {};
   u64 zero = 0, *val;
 
   key.curr_pid = bpf_get_current_pid_tgid();
-  bpf_probe_read(&key.prev_pid, 4,
-		 (void *)ctx->di + offsetof(struct task_struct, pid));
+  key.prev_pid = prev->pid;
 
   val = stats.lookup_or_init(&key, &zero);
   (*val)++;

src/cc/b_frontend_action.cc

@@ -28,6 +28,13 @@ int bpf_create_map(enum bpf_map_type map_type, int key_size, int value_size,
                    int max_entries);
 
 namespace ebpf {
+
+const char *calling_conv_regs_x86[] = {
+  "di", "si", "dx", "cx", "r8", "r9"
+};
+// todo: support more archs
+const char **calling_conv_regs = calling_conv_regs_x86;
+
 using std::map;
 using std::string;
 using std::to_string;
@@ -91,7 +98,7 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
             << "named arguments in BPF program definition";
         return false;
       }
-      fn_args_.push_back(arg->getName());
+      fn_args_.push_back(arg);
     }
   }
   return true;
@@ -205,10 +212,10 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
 
         string text;
         if (Decl->getName() == "incr_cksum_l3") {
-          text = "bpf_l3_csum_replace_(" + fn_args_[0] + ", (u64)";
+          text = "bpf_l3_csum_replace_(" + fn_args_[0]->getName().str() + ", (u64)";
           text += args[0] + ", " + args[1] + ", " + args[2] + ", sizeof(" + args[2] + "))";
         } else if (Decl->getName() == "incr_cksum_l4") {
-          text = "bpf_l4_csum_replace_(" + fn_args_[0] + ", (u64)";
+          text = "bpf_l4_csum_replace_(" + fn_args_[0]->getName().str() + ", (u64)";
           text += args[0] + ", " + args[1] + ", " + args[2];
           text += ", ((" + args[3] + " & 0x1) << 4) | sizeof(" + args[2] + "))";
         } else if (Decl->getName() == "bpf_trace_printk") {
@@ -229,6 +236,50 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
   return true;
 }
 
+bool BTypeVisitor::TraverseMemberExpr(MemberExpr *E) {
+  for (auto child : E->children())
+    if (!TraverseStmt(child))
+      return false;
+  if (!WalkUpFromMemberExpr(E))
+    return false;
+  return true;
+}
+
+bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
+  if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(E->getBase()->IgnoreImplicit())) {
+    auto it = std::find(fn_args_.begin() + 1, fn_args_.end(), Ref->getDecl());
+    if (it != fn_args_.end()) {
+      FieldDecl *F = dyn_cast<FieldDecl>(E->getMemberDecl());
+      string base_type = Ref->getType()->getPointeeType().getAsString();
+      string pre, post;
+      pre = "({ " + E->getType().getAsString() + " _val; memset(&_val, 0, sizeof(_val));";
+      pre += " bpf_probe_read(&_val, sizeof(_val), ";
+      post = " + offsetof(" + base_type + ", " + F->getName().str() + ")";
+      post += "); _val; })";
+      rewriter_.InsertText(E->getLocStart(), pre);
+      rewriter_.ReplaceText(SourceRange(E->getOperatorLoc(), E->getLocEnd()), post);
+    }
+  }
+  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;
@@ -248,7 +299,7 @@ bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
             uint64_t sz = F->isBitField() ? F->getBitWidthValue(C) : C.getTypeSize(F->getType());
             string base = rewriter_.getRewrittenText(SourceRange(Base->getLocStart(), Base->getLocEnd()));
             string rhs = rewriter_.getRewrittenText(SourceRange(RHS->getLocStart(), RHS->getLocEnd()));
-            string text = "bpf_dins_pkt(" + fn_args_[0] + ", (u64)" + base + "+" + to_string(ofs >> 3)
+            string text = "bpf_dins_pkt(" + fn_args_[0]->getName().str() + ", (u64)" + base + "+" + to_string(ofs >> 3)
                 + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ", " + rhs + ")";
             rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
           }
@@ -277,7 +328,7 @@ bool BTypeVisitor::VisitImplicitCastExpr(ImplicitCastExpr *E) {
           }
           uint64_t ofs = C.getFieldOffset(F);
           uint64_t sz = F->isBitField() ? F->getBitWidthValue(C) : C.getTypeSize(F->getType());
-          string text = "bpf_dext_pkt(" + fn_args_[0] + ", (u64)" + Ref->getDecl()->getName().str() + "+"
+          string text = "bpf_dext_pkt(" + fn_args_[0]->getName().str() + ", (u64)" + Ref->getDecl()->getName().str() + "+"
               + to_string(ofs >> 3) + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ")";
           rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
         }

src/cc/b_frontend_action.h

@@ -69,9 +69,12 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   explicit BTypeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter,
                         std::map<std::string, BPFTable> &tables);
   bool TraverseCallExpr(clang::CallExpr *Call);
+  bool TraverseMemberExpr(clang::MemberExpr *E);
   bool VisitFunctionDecl(clang::FunctionDecl *D);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitVarDecl(clang::VarDecl *Decl);
+  bool VisitMemberExpr(clang::MemberExpr *E);
+  bool VisitDeclRefExpr(clang::DeclRefExpr *E);
   bool VisitBinaryOperator(clang::BinaryOperator *E);
   bool VisitImplicitCastExpr(clang::ImplicitCastExpr *E);
 
@@ -80,7 +83,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   clang::Rewriter &rewriter_;  /// modifications to the source go into this class
   llvm::raw_ostream &out_;  /// for debugging
   std::map<std::string, BPFTable> &tables_;  /// store the open FDs
-  std::vector<std::string> fn_args_;
+  std::vector<clang::ParmVarDecl *> fn_args_;
 };
 
 // A helper class to the frontend action, walks the decls

tests/cc/test_clang.py

@@ -25,5 +25,28 @@ EOP:
         b = BPF(text=text, debug=0)
         fn = b.load_func("handle_packet", BPF.SCHED_CLS)
 
+    def test_probe_read1(self):
+        text = """
+#include <linux/sched.h>
+#include <uapi/linux/ptrace.h>
+int count_sched(struct pt_regs *ctx, struct task_struct *prev) {
+    pid_t p = prev->pid;
+    return (p != -1);
+}
+"""
+        b = BPF(text=text, debug=0)
+        fn = b.load_func("count_sched", BPF.KPROBE)
+
+    def test_probe_read2(self):
+        text = """
+#include <linux/sched.h>
+#include <uapi/linux/ptrace.h>
+int count_foo(struct pt_regs *ctx, unsigned long a, unsigned long b) {
+    return (a != b);
+}
+"""
+        b = BPF(text=text, debug=0)
+        fn = b.load_func("count_foo", BPF.KPROBE)
+
 if __name__ == "__main__":
     main()