Merge pull request #125 from iovisor/bblanco_dev

Add printf for key/leaf and fix trace_printk bug

src/cc/bpf_common.cc

@@ -182,4 +182,15 @@ int bpf_table_update_id(void *program, size_t id, const char *key, const char *l
   return mod->table_update(id, key, leaf);
 }
 
+int bpf_table_key_snprintf(void *program, size_t id, char *buf, size_t buflen, const void *key) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->table_key_printf(id, buf, buflen, key);
+}
+int bpf_table_leaf_snprintf(void *program, size_t id, char *buf, size_t buflen, const void *leaf) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->table_key_printf(id, buf, buflen, leaf);
+}
+
 }

src/cc/bpf_common.h

@@ -48,6 +48,10 @@ size_t bpf_table_key_size(void *program, const char *table_name);
 size_t bpf_table_key_size_id(void *program, size_t id);
 size_t bpf_table_leaf_size(void *program, const char *table_name);
 size_t bpf_table_leaf_size_id(void *program, size_t id);
+int bpf_table_key_snprintf(void *program, size_t id, char *buf, size_t buflen, const void *key);
+int bpf_table_leaf_snprintf(void *program, size_t id, char *buf, size_t buflen, const void *leaf);
+//int bpf_table_key_sscanf(void *program, size_t id, const char *buf, void *key);
+//int bpf_table_leaf_sscanf(void *program, size_t id, const char *buf, void *leaf);
 int bpf_table_update(void *program, const char *table_name, const char *key, const char *leaf);
 int bpf_table_update_id(void *program, size_t id, const char *key, const char *leaf);
 

src/cc/bpf_module.cc

@@ -109,31 +109,52 @@ BPFModule::BPFModule(unsigned flags)
 
 BPFModule::~BPFModule() {
   engine_.reset();
+  rw_engine_.reset();
   ctx_.reset();
 }
 
+static void debug_printf(Module *mod, IRBuilder<> &B, const string &fmt, vector<Value *> args) {
+  GlobalVariable *fmt_gvar = B.CreateGlobalString(fmt, "fmt");
+  args.insert(args.begin(), B.CreateInBoundsGEP(fmt_gvar, vector<Value *>({B.getInt64(0), B.getInt64(0)})));
+  args.insert(args.begin(), B.getInt64((uintptr_t)stderr));
+  Function *fprintf_fn = mod->getFunction("fprintf");
+  if (!fprintf_fn) {
+    vector<Type *> fprintf_fn_args({B.getInt64Ty(), B.getInt8PtrTy()});
+    FunctionType *fprintf_fn_type = FunctionType::get(B.getInt32Ty(), fprintf_fn_args, /*isvarArg=*/true);
+    fprintf_fn = Function::Create(fprintf_fn_type, GlobalValue::ExternalLinkage, "fprintf", mod);
+    fprintf_fn->setCallingConv(CallingConv::C);
+    fprintf_fn->addFnAttr(Attribute::NoUnwind);
+  }
+  B.CreateCall(fprintf_fn, args);
+}
+
 // recursive helper to capture the arguments
-void parse_type(IRBuilder<> &B, vector<Value *> *args, string *fmt, Type *type, Value *out) {
+static void parse_type(IRBuilder<> &B, vector<Value *> *args, string *fmt,
+                       Type *type, Value *out, bool is_writer) {
   if (StructType *st = dyn_cast<StructType>(type)) {
     *fmt += "{ ";
     unsigned idx = 0;
     for (auto field : st->elements()) {
-      parse_type(B, args, fmt, field, B.CreateStructGEP(type, out, idx++));
+      parse_type(B, args, fmt, field, B.CreateStructGEP(type, out, idx++), is_writer);
       *fmt += " ";
     }
     *fmt += "}";
   } else if (IntegerType *it = dyn_cast<IntegerType>(type)) {
+    if (is_writer)
+      *fmt += "0x";
     if (it->getBitWidth() <= 8)
-      *fmt += "%hhi";
+      *fmt += "%hh";
     else if (it->getBitWidth() <= 16)
-      *fmt += "%hi";
+      *fmt += "%h";
     else if (it->getBitWidth() <= 32)
-      *fmt += "%i";
-    else if (it->getBitWidth() <= 64)
-      *fmt += "%li";
+      *fmt += "%l";
     else
-      *fmt += "%lli";
-    args->push_back(out);
+      *fmt += "%ll";
+    if (is_writer)
+      *fmt += "x";
+    else
+      *fmt += "i";
+    args->push_back(is_writer ? B.CreateLoad(out) : out);
   }
 }
 
@@ -168,14 +189,13 @@ Function * BPFModule::make_reader(Module *mod, Type *type) {
   BasicBlock *label_exit = BasicBlock::Create(*ctx_, "exit", fn);
   B.SetInsertPoint(label_entry);
 
-  vector<Value *> args;
+  vector<Value *> args({arg_in, nullptr});
   string fmt;
-  parse_type(B, &args, &fmt, type, arg_out);
+  parse_type(B, &args, &fmt, type, arg_out, false);
 
   GlobalVariable *fmt_gvar = B.CreateGlobalString(fmt, "fmt");
 
-  args.insert(args.begin(), B.CreateInBoundsGEP(fmt_gvar, vector<Value *>({B.getInt64(0), B.getInt64(0)})));
-  args.insert(args.begin(), arg_in);
+  args[1] = B.CreateInBoundsGEP(fmt_gvar, vector<Value *>({B.getInt64(0), B.getInt64(0)}));
 
   vector<Type *> sscanf_fn_args({B.getInt8PtrTy(), B.getInt8PtrTy()});
   FunctionType *sscanf_fn_type = FunctionType::get(B.getInt32Ty(), sscanf_fn_args, /*isVarArg=*/true);
@@ -203,7 +223,63 @@ Function * BPFModule::make_reader(Module *mod, Type *type) {
   return fn;
 }
 
-unique_ptr<ExecutionEngine> BPFModule::finalize_reader(unique_ptr<Module> m) {
+Function * BPFModule::make_writer(Module *mod, Type *type) {
+  auto fn_it = writers_.find(type);
+  if (fn_it != writers_.end())
+    return fn_it->second;
+
+  // int write(int len, char *out, Type *in) {
+  //   return snprintf(out, len, "{ %i ... }", out->field1, ...);
+  // }
+
+  IRBuilder<> B(*ctx_);
+
+  // The JIT currently supports a limited number of function prototypes, use the
+  // int (*) (int, char **, const char **) version
+  vector<Type *> fn_args({B.getInt32Ty(), B.getInt8PtrTy(), PointerType::getUnqual(type)});
+  FunctionType *fn_type = FunctionType::get(B.getInt32Ty(), fn_args, /*isVarArg=*/false);
+  Function *fn = Function::Create(fn_type, GlobalValue::ExternalLinkage,
+                                  "writer" + std::to_string(writers_.size()), mod);
+  auto arg_it = fn->arg_begin();
+  Argument *arg_len = arg_it++;
+  arg_len->setName("len");
+  Argument *arg_out = arg_it++;
+  arg_out->setName("out");
+  Argument *arg_in = arg_it++;
+  arg_in->setName("in");
+
+  BasicBlock *label_entry = BasicBlock::Create(*ctx_, "entry", fn);
+  B.SetInsertPoint(label_entry);
+
+  vector<Value *> args({arg_out, B.CreateZExt(arg_len, B.getInt64Ty()), nullptr});
+  string fmt;
+  parse_type(B, &args, &fmt, type, arg_in, true);
+
+  GlobalVariable *fmt_gvar = B.CreateGlobalString(fmt, "fmt");
+
+  args[2] = B.CreateInBoundsGEP(fmt_gvar, vector<Value *>({B.getInt64(0), B.getInt64(0)}));
+
+  if (0)
+    debug_printf(mod, B, "%d %p %p\n", vector<Value *>({arg_len, arg_out, arg_in}));
+
+  vector<Type *> snprintf_fn_args({B.getInt8PtrTy(), B.getInt64Ty(), B.getInt8PtrTy()});
+  FunctionType *snprintf_fn_type = FunctionType::get(B.getInt32Ty(), snprintf_fn_args, /*isVarArg=*/true);
+  Function *snprintf_fn = mod->getFunction("snprintf");
+  if (!snprintf_fn)
+    snprintf_fn = Function::Create(snprintf_fn_type, GlobalValue::ExternalLinkage, "snprintf", mod);
+  snprintf_fn->setCallingConv(CallingConv::C);
+  snprintf_fn->addFnAttr(Attribute::NoUnwind);
+
+  CallInst *call = B.CreateCall(snprintf_fn, args);
+  call->setTailCall(true);
+
+  B.CreateRet(call);
+
+  writers_[type] = fn;
+  return fn;
+}
+
+unique_ptr<ExecutionEngine> BPFModule::finalize_rw(unique_ptr<Module> m) {
   Module *mod = &*m;
 
   run_pass_manager(*mod);
@@ -256,22 +332,24 @@ int BPFModule::annotate() {
         Type *key_type = st->elements()[0];
         Type *leaf_type = st->elements()[1];
         table.key_reader = make_reader(&*m, key_type);
-        if (!table.key_reader) {
+        if (!table.key_reader)
           errs() << "Failed to compile reader for " << *key_type << "\n";
-          continue;
-        }
         table.leaf_reader = make_reader(&*m, leaf_type);
-        if (!table.leaf_reader) {
+        if (!table.leaf_reader)
           errs() << "Failed to compile reader for " << *leaf_type << "\n";
-          continue;
-        }
+        table.key_writer = make_writer(&*m, key_type);
+        if (!table.key_writer)
+          errs() << "Failed to compile writer for " << *key_type << "\n";
+        table.leaf_writer = make_writer(&*m, leaf_type);
+        if (!table.leaf_writer)
+          errs() << "Failed to compile writer for " << *leaf_type << "\n";
       }
     }
   }
 
-  reader_engine_ = finalize_reader(move(m));
-  if (reader_engine_)
-    reader_engine_->finalizeObject();
+  rw_engine_ = finalize_rw(move(m));
+  if (rw_engine_)
+    rw_engine_->finalizeObject();
 
   return 0;
 }
@@ -467,7 +545,7 @@ int BPFModule::table_update(size_t id, const char *key_str, const char *leaf_str
   const TableDesc &desc = (*tables_)[id];
   if (desc.fd < 0) return -1;
 
-  if (!reader_engine_ || !desc.key_reader || !desc.leaf_reader) {
+  if (!rw_engine_ || !desc.key_reader || !desc.leaf_reader) {
     fprintf(stderr, "Table sscanf not available\n");
     return -1;
   }
@@ -475,19 +553,80 @@ int BPFModule::table_update(size_t id, const char *key_str, const char *leaf_str
   unique_ptr<uint8_t[]> key(new uint8_t[desc.key_size]);
   unique_ptr<uint8_t[]> leaf(new uint8_t[desc.leaf_size]);
   GenericValue rc;
-  rc = reader_engine_->runFunction(desc.key_reader, vector<GenericValue>({GenericValue(),
-                                                                          GenericValue((void *)key_str),
-                                                                          GenericValue((void *)key.get())}));
+  rc = rw_engine_->runFunction(desc.key_reader, vector<GenericValue>({GenericValue(),
+                                                                      GenericValue((void *)key_str),
+                                                                      GenericValue((void *)key.get())}));
   if (rc.IntVal != 0)
     return -1;
-  rc = reader_engine_->runFunction(desc.leaf_reader, vector<GenericValue>({GenericValue(),
-                                                                           GenericValue((void *)leaf_str),
-                                                                           GenericValue((void *)leaf.get())}));
+  rc = rw_engine_->runFunction(desc.leaf_reader, vector<GenericValue>({GenericValue(),
+                                                                       GenericValue((void *)leaf_str),
+                                                                       GenericValue((void *)leaf.get())}));
   if (rc.IntVal != 0)
     return -1;
   return bpf_update_elem(desc.fd, key.get(), leaf.get(), 0);
 }
 
+struct TableIterator {
+  TableIterator(size_t key_size, size_t leaf_size)
+      : key(new uint8_t[key_size]), leaf(new uint8_t[leaf_size]) {
+  }
+  unique_ptr<uint8_t[]> key;
+  unique_ptr<uint8_t[]> leaf;
+  uint8_t keyb[512];
+};
+
+int BPFModule::table_key_printf(size_t id, char *buf, size_t buflen, const void *key) {
+  if (id >= tables_->size()) {
+    fprintf(stderr, "table id %zu out of range\n", id);
+    return -1;
+  }
+
+  const TableDesc &desc = (*tables_)[id];
+  if (!desc.key_writer) {
+    fprintf(stderr, "table snprintf not implemented for %s key\n", desc.name.c_str());
+    return -1;
+  }
+  GenericValue gv_buflen;
+  gv_buflen.IntVal = APInt(32, buflen, true);
+  vector<GenericValue> args({gv_buflen, GenericValue((void *)buf), GenericValue((void *)key)});
+  GenericValue rc = rw_engine_->runFunction(desc.key_writer, args);
+  if (rc.IntVal.isNegative()) {
+    perror("snprintf");
+    return -1;
+  }
+  if (rc.IntVal.sge(buflen)) {
+    fprintf(stderr, "snprintf ran out of buffer space\n");
+    return -1;
+  }
+  return 0;
+}
+
+int BPFModule::table_leaf_printf(size_t id, char *buf, size_t buflen, const void *leaf) {
+  if (id >= tables_->size()) {
+    fprintf(stderr, "table id %zu out of range\n", id);
+    return -1;
+  }
+
+  const TableDesc &desc = (*tables_)[id];
+  if (!desc.leaf_writer) {
+    fprintf(stderr, "table snprintf not implemented for %s leaf\n", desc.name.c_str());
+    return -1;
+  }
+  GenericValue gv_buflen;
+  gv_buflen.IntVal = buflen;
+  vector<GenericValue> args({gv_buflen, GenericValue((void *)buf), GenericValue((void *)leaf)});
+  GenericValue rc = rw_engine_->runFunction(desc.leaf_writer, args);
+  if (rc.IntVal.isNegative()) {
+    perror("snprintf");
+    return -1;
+  }
+  if (rc.IntVal.sge(buflen)) {
+    fprintf(stderr, "snprintf ran out of buffer space\n");
+    return -1;
+  }
+  return 0;
+}
+
 // load a B file, which comes in two parts
 int BPFModule::load_b(const string &filename, const string &proto_filename) {
   if (!sections_.empty()) {

src/cc/bpf_module.h

@@ -42,8 +42,9 @@ class BPFModule {
   int parse(llvm::Module *mod);
   int finalize();
   int annotate();
-  std::unique_ptr<llvm::ExecutionEngine> finalize_reader(std::unique_ptr<llvm::Module> mod);
+  std::unique_ptr<llvm::ExecutionEngine> finalize_rw(std::unique_ptr<llvm::Module> mod);
   llvm::Function * make_reader(llvm::Module *mod, llvm::Type *type);
+  llvm::Function * make_writer(llvm::Module *mod, llvm::Type *type);
   void dump_ir(llvm::Module &mod);
   int load_file_module(std::unique_ptr<llvm::Module> *mod, const std::string &file, bool in_memory);
   int load_includes(const std::string &tmpfile);
@@ -70,10 +71,12 @@ class BPFModule {
   const char * table_key_desc(const std::string &name) const;
   size_t table_key_size(size_t id) const;
   size_t table_key_size(const std::string &name) const;
+  int table_key_printf(size_t id, char *buf, size_t buflen, const void *key);
   const char * table_leaf_desc(size_t id) const;
   const char * table_leaf_desc(const std::string &name) const;
   size_t table_leaf_size(size_t id) const;
   size_t table_leaf_size(const std::string &name) const;
+  int table_leaf_printf(size_t id, char *buf, size_t buflen, const void *leaf);
   int table_update(size_t id, const char *key, const char *leaf);
   int table_update(const std::string &name, const char *key, const char *leaf);
   char * license() const;
@@ -84,7 +87,7 @@ class BPFModule {
   std::string proto_filename_;
   std::unique_ptr<llvm::LLVMContext> ctx_;
   std::unique_ptr<llvm::ExecutionEngine> engine_;
-  std::unique_ptr<llvm::ExecutionEngine> reader_engine_;
+  std::unique_ptr<llvm::ExecutionEngine> rw_engine_;
   std::unique_ptr<llvm::Module> mod_;
   std::unique_ptr<BLoader> b_loader_;
   std::unique_ptr<ClangLoader> clang_loader_;
@@ -93,6 +96,7 @@ class BPFModule {
   std::map<std::string, size_t> table_names_;
   std::vector<std::string> function_names_;
   std::map<llvm::Type *, llvm::Function *> readers_;
+  std::map<llvm::Type *, llvm::Function *> writers_;
 };
 
 }  // namespace ebpf

src/cc/frontends/clang/b_frontend_action.cc

@@ -227,24 +227,24 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
         if (Decl->getName() == "incr_cksum_l3") {
           text = "bpf_l3_csum_replace_(" + fn_args_[0]->getName().str() + ", (u64)";
           text += args[0] + ", " + args[1] + ", " + args[2] + ", sizeof(" + args[2] + "))";
+          rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
         } else if (Decl->getName() == "incr_cksum_l4") {
           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] + "))";
+          rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
         } else if (Decl->getName() == "bpf_trace_printk") {
           //  #define bpf_trace_printk(fmt, args...)
           //    ({ char _fmt[] = fmt; bpf_trace_printk_(_fmt, sizeof(_fmt), args...); })
           text = "({ char _fmt[] = " + args[0] + "; bpf_trace_printk_(_fmt, sizeof(_fmt)";
-          if (args.size() > 1)
-            text += ", ";
-          for (auto arg = args.begin() + 1; arg != args.end(); ++arg) {
-            text += *arg;
-            if (arg + 1 != args.end())
-              text += ", ";
+          if (args.size() <= 1) {
+            text += "); })";
+            rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
+          } else {
+            rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getArg(0)->getLocEnd()), text);
+            rewriter_.InsertTextAfter(Call->getLocEnd(), "); }");
           }
-          text += "); })";
         }
-        rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
       }
     }
   }

src/cc/table_desc.h

@@ -33,6 +33,8 @@ struct TableDesc {
   std::string leaf_desc;
   llvm::Function *key_reader;
   llvm::Function *leaf_reader;
+  llvm::Function *key_writer;
+  llvm::Function *leaf_writer;
 };
 
 }  // namespace ebpf

tests/cc/test_clang.py

@@ -82,6 +82,18 @@ int do_request(struct pt_regs *ctx, struct request *req) {
 
     return 0;
 }
+"""
+        b = BPF(text=text, debug=0)
+        fn = b.load_func("do_request", BPF.KPROBE)
+
+    def test_blk_start_request(self):
+        text = """
+#include <linux/blkdev.h>
+#include <uapi/linux/ptrace.h>
+int do_request(struct pt_regs *ctx, int req) {
+    bpf_trace_printk("req ptr: 0x%x\\n", req);
+    return 0;
+}
 """
         b = BPF(text=text, debug=0)
         fn = b.load_func("do_request", BPF.KPROBE)