Automatically detect key/leaf type in clang

* Parse the key and leaf types used for table definition, and pass the
  type info back into python using a json transport
* Remove the subscript operator support
* Migrate the tests that are able to drop they Key/Leaf definition. Keep
  it around for the tests that are B/C compatible
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

examples/vlan_learning.py

@@ -38,18 +38,13 @@ num_workers = 3
 num_clients = 9
 num_vlans = 16
 
-class ifindex_leaf_t(Structure):
-    _fields_ = [("out_ifindex", c_int),
-                ("tx_pkts", c_ulonglong),
-                ("tx_bytes", c_ulonglong)]
-
 # load the bpf program
 b = BPF(src_file="examples/vlan_learning.c", debug=0)
 phys_fn = b.load_func("handle_phys2virt", BPF.SCHED_CLS)
 virt_fn = b.load_func("handle_virt2phys", BPF.SCHED_CLS)
 
-ingress = b.get_table("ingress", c_ulonglong, ifindex_leaf_t)
-egress = b.get_table("egress", c_ulonglong, ifindex_leaf_t)
+ingress = b.get_table("ingress")
+egress = b.get_table("egress")
 
 ipdb_workers = []
 ipdb_clients = []

src/cc/b_frontend_action.cc

@@ -34,6 +34,44 @@ using std::to_string;
 using std::unique_ptr;
 using namespace clang;
 
+// Encode the struct layout as a json description
+BMapDeclVisitor::BMapDeclVisitor(ASTContext &C, string &result)
+    : C(C), result_(result) {}
+bool BMapDeclVisitor::VisitFieldDecl(FieldDecl *D) {
+  result_ += "\"";
+  result_ += D->getName();
+  result_ += "\",";
+  return true;
+}
+bool BMapDeclVisitor::VisitRecordDecl(RecordDecl *D) {
+  result_ += "[\"";
+  result_ += D->getName();
+  result_ += "\", [";
+  for (auto F : D->getDefinition()->fields()) {
+    result_ += "[";
+    TraverseDecl(F);
+    if (F->isBitField())
+      result_ += ", " + to_string(F->getBitWidthValue(C));
+    result_ += "], ";
+  }
+  if (!D->getDefinition()->field_empty())
+    result_.erase(result_.end() - 2);
+  result_ += "]]";
+  return false;
+}
+bool BMapDeclVisitor::VisitTagType(const TagType *T) {
+  return TraverseDecl(T->getDecl()->getDefinition());
+}
+bool BMapDeclVisitor::VisitTypedefType(const TypedefType *T) {
+  return TraverseDecl(T->getDecl());
+}
+bool BMapDeclVisitor::VisitBuiltinType(const BuiltinType *T) {
+  result_ += "\"";
+  result_ += T->getName(C.getPrintingPolicy());
+  result_ += "\"";
+  return true;
+}
+
 BTypeVisitor::BTypeVisitor(ASTContext &C, Rewriter &rewriter, map<string, BPFTable> &tables)
     : C(C), rewriter_(rewriter), out_(llvm::errs()), tables_(tables) {
 }
@@ -122,53 +160,6 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
   return true;
 }
 
-// look for table subscript references, and turn them into auto table entries:
-// table.data[key]
-//  becomes:
-// struct Key key = {123};
-// struct Leaf *leaf = table.get(&key);
-// if (!leaf) {
-//   struct Leaf zleaf = {0};
-//   table.put(&key, &zleaf, BPF_NOEXIST);
-//   leaf = table.get(&key);
-//   if (!leaf) return -1;
-// }
-bool BTypeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *Arr) {
-  Expr *LHS = Arr->getLHS()->IgnoreImplicit();
-  Expr *RHS = Arr->getRHS()->IgnoreImplicit();
-  if (MemberExpr *Memb = dyn_cast<MemberExpr>(LHS)) {
-    if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
-      if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
-        if (A->getName().startswith("maps")) {
-          auto table_it = tables_.find(Ref->getDecl()->getName());
-          if (table_it == tables_.end()) {
-            C.getDiagnostics().Report(Ref->getLocEnd(), diag::err_expected)
-                << "initialized handle for bpf_table";
-            return false;
-          }
-          string fd = to_string(table_it->second.fd);
-          string map_update_policy = "BPF_NOEXIST";
-          string name = Ref->getDecl()->getName();
-          SourceRange argRange(RHS->getLocStart(), RHS->getLocEnd());
-          string args = rewriter_.getRewrittenText(argRange);
-          string lookup = "bpf_map_lookup_elem_(bpf_pseudo_fd(1, " + fd + ")";
-          string update = "bpf_map_update_elem_(bpf_pseudo_fd(1, " + fd + ")";
-          string txt = "(*({typeof(" + name + ".leaf) *leaf = " + lookup + ", " + args + "); ";
-          txt         += "if (!leaf) {";
-          txt         += " typeof(" + name + ".leaf) zleaf = {0};";
-          txt         += " " + update + ", " + args + ", &zleaf, " + map_update_policy + ");";
-          txt         += " leaf = " + lookup + ", " + args + ");";
-          txt         += " if (!leaf) return 0;";
-          txt         += "}";
-          txt         += "leaf;}))";
-          rewriter_.ReplaceText(SourceRange(Arr->getLocStart(), Arr->getLocEnd()), txt);
-        }
-      }
-    }
-  }
-  return true;
-}
-
 bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
   if (!E->isAssignmentOp())
     return true;
@@ -237,8 +228,12 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       size_t sz = C.getTypeSize(F->getType()) >> 3;
       if (F->getName() == "key") {
         table.key_size = sz;
+        BMapDeclVisitor visitor(C, table.key_desc);
+        visitor.TraverseType(F->getType());
       } else if (F->getName() == "leaf") {
         table.leaf_size = sz;
+        BMapDeclVisitor visitor(C, table.leaf_desc);
+        visitor.TraverseType(F->getType());
       } else if (F->getName() == "data") {
         table.max_entries = sz / table.leaf_size;
       }
@@ -256,7 +251,7 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       llvm::errs() << "error: could not open bpf fd\n";
       return false;
     }
-    tables_[Decl->getName()] = table;
+    tables_[Decl->getName()] = std::move(table);
   }
   return true;
 }

src/cc/b_frontend_action.h

@@ -41,6 +41,23 @@ struct BPFTable {
   size_t key_size;
   size_t leaf_size;
   size_t max_entries;
+  std::string key_desc;
+  std::string leaf_desc;
+};
+
+// Helper visitor for constructing a string representation of a key/leaf decl
+class BMapDeclVisitor : public clang::RecursiveASTVisitor<BMapDeclVisitor> {
+ public:
+  explicit BMapDeclVisitor(clang::ASTContext &C, std::string &result);
+  bool VisitRecordDecl(clang::RecordDecl *Decl);
+  bool VisitFieldDecl(clang::FieldDecl *Decl);
+  bool VisitBuiltinType(const clang::BuiltinType *T);
+  bool VisitTypedefType(const clang::TypedefType *T);
+  bool VisitTagType(const clang::TagType *T);
+  const std::string & str() const { return result_; }
+ private:
+  clang::ASTContext &C;
+  std::string &result_;
 };
 
 // Type visitor and rewriter for B programs.
@@ -54,7 +71,6 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   bool VisitFunctionDecl(clang::FunctionDecl *D);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitVarDecl(clang::VarDecl *Decl);
-  bool VisitArraySubscriptExpr(clang::ArraySubscriptExpr *E);
   bool VisitBinaryOperator(clang::BinaryOperator *E);
   bool VisitImplicitCastExpr(clang::ImplicitCastExpr *E);
 

src/cc/bpf_common.cc

@@ -71,4 +71,16 @@ int bpf_table_fd(void *program, const char *table_name) {
   return mod->table_fd(table_name);
 }
 
+const char * bpf_table_key_desc(void *program, const char *table_name) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return nullptr;
+  return mod->table_key_desc(table_name);
+}
+
+const char * bpf_table_leaf_desc(void *program, const char *table_name) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return nullptr;
+  return mod->table_leaf_desc(table_name);
+}
+
 }

src/cc/bpf_common.h

@@ -28,6 +28,8 @@ unsigned bpf_module_kern_version(void *program);
 void * bpf_function_start(void *program, const char *name);
 size_t bpf_function_size(void *program, const char *name);
 int bpf_table_fd(void *program, const char *table_name);
+const char * bpf_table_key_desc(void *program, const char *table_name);
+const char * bpf_table_leaf_desc(void *program, const char *table_name);
 
 #ifdef __cplusplus
 }

src/cc/bpf_module.cc

@@ -408,6 +408,20 @@ int BPFModule::table_fd(const string &name) const {
   return table_it->second.fd;
 }
 
+const char * BPFModule::table_key_desc(const string &name) const {
+  if (codegen_) return nullptr;
+  auto table_it = tables_->find(name);
+  if (table_it == tables_->end()) return nullptr;
+  return table_it->second.key_desc.c_str();
+}
+
+const char * BPFModule::table_leaf_desc(const string &name) const {
+  if (codegen_) return nullptr;
+  auto table_it = tables_->find(name);
+  if (table_it == tables_->end()) return nullptr;
+  return table_it->second.leaf_desc.c_str();
+}
+
 int BPFModule::load(const string &filename, const string &proto_filename) {
   if (!sections_.empty()) {
     fprintf(stderr, "Program already initialized\n");

src/cc/bpf_module.h

@@ -54,6 +54,8 @@ class BPFModule {
   uint8_t * start(const std::string &name) const;
   size_t size(const std::string &name) const;
   int table_fd(const std::string &name) const;
+  const char * table_key_desc(const std::string &name) const;
+  const char * table_leaf_desc(const std::string &name) const;
   char * license() const;
   unsigned kern_version() const;
  private:

src/python/bpf/__init__.py

@@ -14,6 +14,7 @@
 
 import atexit
 import ctypes as ct
+import json
 import os
 
 lib = ct.CDLL("libbpfprog.so")
@@ -35,6 +36,10 @@ lib.bpf_function_size.restype = ct.c_size_t
 lib.bpf_function_size.argtypes = [ct.c_void_p, ct.c_char_p]
 lib.bpf_table_fd.restype = ct.c_int
 lib.bpf_table_fd.argtypes = [ct.c_void_p, ct.c_char_p]
+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]
 
 # keep in sync with libbpf.h
 lib.bpf_get_next_key.restype = ct.c_int
@@ -83,12 +88,12 @@ class BPF(object):
         def __init__(self, bpf, map_fd, keytype, leaftype):
             self.bpf = bpf
             self.map_fd = map_fd
-            self.keytype = keytype
-            self.leaftype = leaftype
+            self.Key = keytype
+            self.Leaf = leaftype
 
         def lookup(self, key):
             key_p = ct.pointer(key)
-            leaf = self.leaftype()
+            leaf = self.Leaf()
             leaf_p = ct.pointer(leaf)
             res = lib.bpf_lookup_elem(self.map_fd,
                     ct.cast(key_p, ct.c_void_p),
@@ -108,9 +113,9 @@ class BPF(object):
 
         class Iter(object):
             def __init__(self, table, keytype):
-                self.keytype = keytype
+                self.Key = keytype
                 self.table = table
-                self.key = keytype()
+                self.key = self.Key()
             def __iter__(self):
                 return self
             def __next__(self):
@@ -120,10 +125,10 @@ class BPF(object):
                 return self.key
 
         def iter(self):
-            return BPF.Table.Iter(self, self.keytype)
+            return BPF.Table.Iter(self, self.Key)
 
         def next(self, key):
-            next_key = self.keytype()
+            next_key = self.Key()
             next_key_p = ct.pointer(next_key)
             key_p = ct.pointer(key)
             res = lib.bpf_get_next_key(self.map_fd,
@@ -165,11 +170,51 @@ class BPF(object):
 
         return fn
 
-    def get_table(self, name, keytype, leaftype):
-        map_fd = lib.bpf_table_fd(self.module,
-                ct.c_char_p(name.encode("ascii")))
+    str2ctype = {
+        "_Bool": ct.c_bool,
+        "char": ct.c_char,
+        "wchar_t": ct.c_wchar,
+        "char": ct.c_byte,
+        "unsigned char": ct.c_ubyte,
+        "short": ct.c_short,
+        "unsigned short": ct.c_ushort,
+        "int": ct.c_int,
+        "unsigned int": ct.c_uint,
+        "long": ct.c_long,
+        "unsigned long": ct.c_ulong,
+        "long long": ct.c_longlong,
+        "unsigned long long": ct.c_ulonglong,
+        "float": ct.c_float,
+        "double": ct.c_double,
+        "long double": ct.c_longdouble
+    }
+    @staticmethod
+    def _decode_table_type(desc):
+        if isinstance(desc, str):
+            return BPF.str2ctype[desc]
+        fields = []
+        for t in desc[1]:
+            if len(t) == 2:
+                fields.append((t[0], BPF._decode_table_type(t[1])))
+            elif len(t) == 3:
+                fields.append((t[0], BPF._decode_table_type(t[1]), t[2]))
+        cls = type(desc[0], (ct.Structure,), dict(_fields_=fields))
+        return cls
+
+    def get_table(self, name, keytype=None, leaftype=None):
+        map_fd = lib.bpf_table_fd(self.module, name.encode("ascii"))
         if map_fd < 0:
             raise Exception("Failed to find BPF Table %s" % name)
+        if not keytype:
+            key_desc = lib.bpf_table_key_desc(self.module, name.encode("ascii"))
+            if not key_desc:
+                raise Exception("Failed to load BPF Table %s key desc" % name)
+            keytype = BPF._decode_table_type(json.loads(key_desc.decode()))
+        if not leaftype:
+            leaf_desc = lib.bpf_table_leaf_desc(self.module, name.encode("ascii"))
+            if not leaf_desc:
+                raise Exception("Failed to load BPF Table %s leaf desc" % name)
+            leaftype = BPF._decode_table_type(json.loads(leaf_desc.decode()))
         return BPF.Table(self, map_fd, keytype, leaftype)
 
     @staticmethod

tests/cc/test_stat1.py

@@ -17,12 +17,15 @@ arg2 = ""
 if len(sys.argv) > 1:
   arg2 = sys.argv.pop(1)
 
-class Key(Structure):
-    _fields_ = [("dip", c_uint),
-                ("sip", c_uint)]
-class Leaf(Structure):
-    _fields_ = [("rx_pkts", c_ulong),
-                ("tx_pkts", c_ulong)]
+Key = None
+Leaf = None
+if arg1.endswith(".b"):
+    class Key(Structure):
+        _fields_ = [("dip", c_uint),
+                    ("sip", c_uint)]
+    class Leaf(Structure):
+        _fields_ = [("rx_pkts", c_ulong),
+                    ("tx_pkts", c_ulong)]
 
 class TestBPFSocket(TestCase):
     def setUp(self):
@@ -38,7 +41,7 @@ class TestBPFSocket(TestCase):
         #    leaf = self.stats.lookup(key)
         #    print(IPAddress(key.sip), "=>", IPAddress(key.dip),
         #          "rx", leaf.rx_pkts, "tx", leaf.tx_pkts)
-        key = Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
+        key = self.stats.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
         leaf = self.stats.lookup(key)
         self.assertEqual(leaf.rx_pkts, 100)
         self.assertEqual(leaf.tx_pkts, 100)

tests/cc/test_trace1.py

@@ -14,11 +14,14 @@ arg2 = ""
 if len(sys.argv) > 1:
   arg2 = sys.argv.pop(1)
 
-class Key(Structure):
-    _fields_ = [("fd", c_ulong)]
-class Leaf(Structure):
-    _fields_ = [("stat1", c_ulong),
-                ("stat2", c_ulong)]
+Key = None
+Leaf = None
+if arg1.endswith(".b"):
+    class Key(Structure):
+        _fields_ = [("fd", c_ulong)]
+    class Leaf(Structure):
+        _fields_ = [("stat1", c_ulong),
+                    ("stat2", c_ulong)]
 
 class TestKprobe(TestCase):
     def setUp(self):

tests/cc/test_trace2.py

@@ -16,21 +16,17 @@ BPF_TABLE("hash", struct Ptr, struct Counters, stats, 1024);
 
 int count_sched(struct pt_regs *ctx) {
   struct Ptr key = {.ptr=ctx->bx};
-  stats.data[(u64)&key].stat1++;
+  struct Counters zleaf = {0};
+  stats.lookup_or_init(&key, &zleaf)->stat1++;
   return 0;
 }
 """
 
-class Ptr(Structure):
-    _fields_ = [("ptr", c_ulong)]
-class Counters(Structure):
-    _fields_ = [("stat1", c_ulong)]
-
 class TestTracingEvent(TestCase):
     def setUp(self):
         b = BPF(text=text, debug=0)
         fn = b.load_func("count_sched", BPF.KPROBE)
-        self.stats = b.get_table("stats", Ptr, Counters)
+        self.stats = b.get_table("stats")
         BPF.attach_kprobe(fn, "schedule+50", 0, -1)
 
     def test_sched1(self):

tests/cc/test_xlate1.py

@@ -2,7 +2,6 @@
 # Copyright (c) PLUMgrid, Inc.
 # Licensed under the Apache License, Version 2.0 (the "License")
 
-from ctypes import c_uint, c_ulonglong, Structure
 from netaddr import IPAddress
 from bpf import BPF
 from pyroute2 import IPRoute
@@ -17,15 +16,6 @@ arg2 = ""
 if len(sys.argv) > 1:
   arg2 = sys.argv.pop(1)
 
-class Key(Structure):
-    _fields_ = [("dip", c_uint),
-                ("sip", c_uint)]
-class Leaf(Structure):
-    _fields_ = [("xdip", c_uint),
-                ("xsip", c_uint),
-                ("ip_xlated_pkts", c_ulonglong),
-                ("arp_xlated_pkts", c_ulonglong)]
-
 class TestBPFFilter(TestCase):
     def setUp(self):
         b = BPF(arg1, arg2, debug=0)
@@ -43,14 +33,14 @@ class TestBPFFilter(TestCase):
         # add same program to both ingress/egress, so pkt is translated in both directions
         ip.tc("add-filter", "bpf", ifindex, ":1", fd=fn.fd, name=fn.name, parent="ffff:", action="ok", classid=1)
         ip.tc("add-filter", "bpf", ifindex, ":2", fd=fn.fd, name=fn.name, parent="1:", action="ok", classid=1)
-        self.xlate = b.get_table("xlate", Key, Leaf)
+        self.xlate = b.get_table("xlate")
 
     def test_xlate(self):
-        key1 = Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
-        leaf1 = Leaf(IPAddress("192.168.1.2").value, IPAddress("192.168.1.1").value, 0, 0)
+        key1 = self.xlate.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
+        leaf1 = self.xlate.Leaf(IPAddress("192.168.1.2").value, IPAddress("192.168.1.1").value, 0, 0)
         self.xlate.update(key1, leaf1)
-        key2 = Key(IPAddress("192.168.1.1").value, IPAddress("192.168.1.2").value)
-        leaf2 = Leaf(IPAddress("172.16.1.1").value, IPAddress("172.16.1.2").value, 0, 0)
+        key2 = self.xlate.Key(IPAddress("192.168.1.1").value, IPAddress("192.168.1.2").value)
+        leaf2 = self.xlate.Leaf(IPAddress("172.16.1.1").value, IPAddress("172.16.1.2").value, 0, 0)
         self.xlate.update(key2, leaf2)
         call(["ping", "-c1", "192.168.1.1"])
         leaf = self.xlate.lookup(key1)