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 <[email protected]>

examples/vlan_learning.py

@@ -38,18 +38,13 @@
 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.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 Table(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 @@ def update(self, key, leaf, flags=0):
 
  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 @@ def next(self):
  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 @@ def load_func(self, func_name, prog_type):
 
  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 @@
 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 @@ def test_ping(self):
  # 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 @@
 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):