Add update_table API: accepts sscanf style strings and populates a map

This is the culmination of the previous patches. It adds an api that can
input map data in a string format, but validating the numbers and
locations of data fields. The use case is for fuse file input/output. A
printf api may follow.

Take the table with key/leaf of:
  struct Key { int a; int b; };
  struct Leaf { int a; int b; int c; struct SubLeaf { int x; int y; } s; };

One would input to this table using:
  update_table(table_name, "{1 2}", "{1 2 -3 {9 0xa}}");

The implementation uses a JITed function for each unique type, that is
invoked to run sscanf on behalf of the caller. The input must have the
exact right number of arguments. Bit fields are supported, but the
caller must be aware of the collapse of those bitfields into an aligned
field, as well as endianness.
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/cc/bpf_common.cc

@@ -170,4 +170,16 @@ size_t bpf_table_leaf_size_id(void *program, size_t id) {
   return mod->table_leaf_size(id);
 }
 
+int bpf_table_update(void *program, const char *table_name, const char *key, const char *leaf) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->table_update(table_name, key, leaf);
+}
+
+int bpf_table_update_id(void *program, size_t id, const char *key, const char *leaf) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->table_update(id, key, leaf);
+}
+
 }

src/cc/bpf_common.h

@@ -48,6 +48,8 @@ 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_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);
 
 #ifdef __cplusplus
 }

src/cc/bpf_module.cc

@@ -26,6 +26,7 @@
 #include <linux/bpf.h>
 
 #include <llvm/ADT/STLExtras.h>
+#include <llvm/ExecutionEngine/GenericValue.h>
 #include <llvm/ExecutionEngine/MCJIT.h>
 #include <llvm/ExecutionEngine/SectionMemoryManager.h>
 #include <llvm/IRReader/IRReader.h>
@@ -121,14 +122,25 @@ void parse_type(IRBuilder<> &B, vector<Value *> *args, string *fmt, Type *type,
       *fmt += " ";
     }
     *fmt += "}";
-  } else if (dyn_cast<IntegerType>(type)) {
-    *fmt += "%lli";
+  } else if (IntegerType *it = dyn_cast<IntegerType>(type)) {
+    if (it->getBitWidth() <= 8)
+      *fmt += "%hhi";
+    else if (it->getBitWidth() <= 16)
+      *fmt += "%hi";
+    else if (it->getBitWidth() <= 32)
+      *fmt += "%i";
+    else if (it->getBitWidth() <= 64)
+      *fmt += "%li";
+    else
+      *fmt += "%lli";
     args->push_back(out);
   }
 }
 
-int BPFModule::make_reader(Module *mod, Type *type) {
-  if (readers_.find(type) != readers_.end()) return 0;
+Function * BPFModule::make_reader(Module *mod, Type *type) {
+  auto fn_it = readers_.find(type);
+  if (fn_it != readers_.end())
+    return fn_it->second;
 
   // int read(const char *in, Type *out) {
   //   int n = sscanf(in, "{ %i ... }", &out->field1, ...);
@@ -138,11 +150,15 @@ int BPFModule::make_reader(Module *mod, Type *type) {
 
   IRBuilder<> B(*ctx_);
 
-  vector<Type *> fn_args({B.getInt8PtrTy(), PointerType::getUnqual(type)});
+  // 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,
                                   "reader" + std::to_string(readers_.size()), mod);
   auto arg_it = fn->arg_begin();
+  Argument *arg_argc = arg_it++;
+  arg_argc->setName("argc");
   Argument *arg_in = arg_it++;
   arg_in->setName("in");
   Argument *arg_out = arg_it++;
@@ -163,11 +179,15 @@ int BPFModule::make_reader(Module *mod, Type *type) {
 
   vector<Type *> sscanf_fn_args({B.getInt8PtrTy(), B.getInt8PtrTy()});
   FunctionType *sscanf_fn_type = FunctionType::get(B.getInt32Ty(), sscanf_fn_args, /*isVarArg=*/true);
-  Function *sscanf_fn = mod->getFunction("__isoc99_sscanf");
+  Function *sscanf_fn = mod->getFunction("sscanf");
   if (!sscanf_fn)
-    sscanf_fn = Function::Create(sscanf_fn_type, GlobalValue::ExternalLinkage, "__isoc99_sscanf", mod);
+    sscanf_fn = Function::Create(sscanf_fn_type, GlobalValue::ExternalLinkage, "sscanf", mod);
+  sscanf_fn->setCallingConv(CallingConv::C);
+  sscanf_fn->addFnAttr(Attribute::NoUnwind);
 
   CallInst *call = B.CreateCall(sscanf_fn, args);
+  call->setTailCall(true);
+
   BasicBlock *label_then = BasicBlock::Create(*ctx_, "then", fn);
 
   Value *is_neq = B.CreateICmpNE(call, B.getInt32(args.size() - 2));
@@ -180,7 +200,7 @@ int BPFModule::make_reader(Module *mod, Type *type) {
   B.CreateRet(B.getInt32(0));
 
   readers_[type] = fn;
-  return 0;
+  return fn;
 }
 
 unique_ptr<ExecutionEngine> BPFModule::finalize_reader(unique_ptr<Module> m) {
@@ -226,7 +246,7 @@ int BPFModule::annotate() {
   auto m = make_unique<Module>("sscanf", *ctx_);
 
   size_t id = 0;
-  for (auto table : *tables_) {
+  for (auto &table : *tables_) {
     table_names_[table.name] = id++;
     GlobalValue *gvar = mod_->getNamedValue(table.name);
     if (!gvar) continue;
@@ -235,10 +255,16 @@ int BPFModule::annotate() {
         if (st->getNumElements() < 2) continue;
         Type *key_type = st->elements()[0];
         Type *leaf_type = st->elements()[1];
-        if (int rc = make_reader(&*m, key_type))
-          return rc;
-        if (int rc = make_reader(&*m, leaf_type))
-          return rc;
+        table.key_reader = make_reader(&*m, key_type);
+        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) {
+          errs() << "Failed to compile reader for " << *leaf_type << "\n";
+          continue;
+        }
       }
     }
   }
@@ -429,6 +455,39 @@ size_t BPFModule::table_leaf_size(const string &name) const {
   return table_leaf_size(it->second);
 }
 
+int BPFModule::table_update(const string &name, const char *key_str, const char *leaf_str) {
+  auto it = table_names_.find(name);
+  if (it == table_names_.end()) return 0;
+  return table_update(it->second, key_str, leaf_str);
+}
+
+int BPFModule::table_update(size_t id, const char *key_str, const char *leaf_str) {
+  if (id >= tables_->size()) return -1;
+
+  const TableDesc &desc = (*tables_)[id];
+  if (desc.fd < 0) return -1;
+
+  if (!reader_engine_ || !desc.key_reader || !desc.leaf_reader) {
+    fprintf(stderr, "Table sscanf not available\n");
+    return -1;
+  }
+
+  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())}));
+  if (rc.IntVal != 0)
+    return -1;
+  rc = reader_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);
+}
+
 // 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

@@ -43,7 +43,7 @@ class BPFModule {
   int finalize();
   int annotate();
   std::unique_ptr<llvm::ExecutionEngine> finalize_reader(std::unique_ptr<llvm::Module> mod);
-  int make_reader(llvm::Module *mod, llvm::Type *type);
+  llvm::Function * make_reader(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);
@@ -74,6 +74,8 @@ class BPFModule {
   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_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;
   unsigned kern_version() const;
  private:

src/cc/table_desc.h

@@ -17,6 +17,10 @@
 #include <cstdint>
 #include <string>
 
+namespace llvm {
+class Function;
+}
+
 namespace ebpf {
 
 struct TableDesc {
@@ -27,6 +31,8 @@ struct TableDesc {
   size_t max_entries;
   std::string key_desc;
   std::string leaf_desc;
+  llvm::Function *key_reader;
+  llvm::Function *leaf_reader;
 };
 
 }  // namespace ebpf

src/python/bpf/__init__.py

@@ -45,6 +45,8 @@ lib.bpf_table_key_desc.restype = ct.c_char_p
 lib.bpf_table_key_desc.argtypes = [ct.c_void_p, ct.c_char_p]
 lib.bpf_table_leaf_desc.restype = ct.c_char_p
 lib.bpf_table_leaf_desc.argtypes = [ct.c_void_p, ct.c_char_p]
+lib.bpf_table_update.restype = ct.c_int
+lib.bpf_table_update.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_char_p, ct.c_char_p]
 
 # keep in sync with libbpf.h
 lib.bpf_get_next_key.restype = ct.c_int
@@ -258,6 +260,12 @@ class BPF(object):
             leaftype = BPF._decode_table_type(json.loads(leaf_desc.decode()))
         return BPF.Table(self, map_fd, keytype, leaftype)
 
+    def update_table(self, name, key, leaf):
+        res = lib.bpf_table_update(self.module, name.encode("ascii"), key.encode("ascii"),
+                leaf.encode("ascii"))
+        if res < 0:
+            raise Exception("update_table failed")
+
     @staticmethod
     def attach_raw_socket(fn, dev):
         if not isinstance(fn, BPF.Function):

tests/cc/test_clang.py

@@ -48,14 +48,22 @@ int count_foo(struct pt_regs *ctx, unsigned long a, unsigned long b) {
 
     def test_sscanf(self):
         text = """
-BPF_TABLE("hash", int, struct { u64 a; u64 b; u64 c:31; u64 d:33; struct { u32 a; u32 b; } s; }, stats, 10);
+BPF_TABLE("hash", int, struct { u64 a; u64 b; u64 c:36; u64 d:28; struct { u32 a; u32 b; } s; }, stats, 10);
 int foo(void *ctx) {
     return 0;
 }
 """
         b = BPF(text=text, debug=0)
         fn = b.load_func("foo", BPF.KPROBE)
-        # todo: the actual test
+        b.update_table("stats", "2", "{ 2 3 0x1000000004 { 5 6 }}")
+        t = b.get_table("stats")
+        l = t[t.Key(2)]
+        self.assertEqual(l.a, 2)
+        self.assertEqual(l.b, 3)
+        self.assertEqual(l.c, 4)
+        self.assertEqual(l.d, 1)
+        self.assertEqual(l.s.a, 5)
+        self.assertEqual(l.s.b, 6)
 
 if __name__ == "__main__":
     main()