Merge pull request iovisor#77 from iovisor/bblanco_dev

Automate bpf_probe_read

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 @@ def test_printk(self):
  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()