Improve json type support for misc struct/union types

The ability of the clang rewriter to extract the type information for
some types of structs, unions, and pointers to the aforementioned was
somewhat buggy. This became exposed in a test_clang case after a user
upgraded to a newer kernel, wherein the struct definition changed. The
functionality in question is only used to pass json-ified representation
of the struct to python in order to program the Key/Leaf metaclass.

Improve support for this and other types, including unions.
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/cc/frontends/clang/b_frontend_action.cc

@@ -67,10 +67,7 @@ bool BMapDeclVisitor::VisitRecordDecl(RecordDecl *D) {
   result_ += "\", [";
   for (auto F : D->getDefinition()->fields()) {
     result_ += "[";
-    if (F->getType()->isPointerType())
-      result_ += "\"" + F->getName().str() + "\", \"unsigned long long\"";
-    else
-      TraverseDecl(F);
+    TraverseDecl(F);
     if (const ConstantArrayType *T = dyn_cast<ConstantArrayType>(F->getType()))
       result_ += ", [" + T->getSize().toString(10, false) + "]";
     if (F->isBitField())
@@ -79,9 +76,19 @@ bool BMapDeclVisitor::VisitRecordDecl(RecordDecl *D) {
   }
   if (!D->getDefinition()->field_empty())
     result_.erase(result_.end() - 2);
-  result_ += "]]";
+  result_ += "]";
+  if (D->isUnion())
+    result_ += ", \"union\"";
+  else if (D->isStruct())
+    result_ += ", \"struct\"";
+  result_ += "]";
   return true;
 }
+// pointer to anything should be treated as terminal, don't recurse further
+bool BMapDeclVisitor::VisitPointerType(const PointerType *T) {
+  result_ += "\"unsigned long long\"";
+  return false;
+}
 bool BMapDeclVisitor::VisitTagType(const TagType *T) {
   return TraverseDecl(T->getDecl()->getDefinition());
 }
@@ -510,22 +517,32 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       sscanf(un.release, "%d.%d.", &major, &minor);
     }
 
-    TableDesc table;
+    TableDesc table = {};
     table.name = Decl->getName();
 
     unsigned i = 0;
     for (auto F : RD->fields()) {
       size_t sz = C.getTypeSize(F->getType()) >> 3;
       if (F->getName() == "key") {
+        if (sz == 0) {
+          unsigned diag_id = C.getDiagnostics().getCustomDiagID(DiagnosticsEngine::Error,
+                                                                "invalid zero-sized leaf");
+          C.getDiagnostics().Report(F->getLocStart(), diag_id);
+          return false;
+        }
         table.key_size = sz;
         BMapDeclVisitor visitor(C, table.key_desc);
-        if (!visitor.TraverseType(F->getType()))
-          return false;
+        visitor.TraverseType(F->getType());
       } else if (F->getName() == "leaf") {
+        if (sz == 0) {
+          unsigned diag_id = C.getDiagnostics().getCustomDiagID(DiagnosticsEngine::Error,
+                                                                "invalid zero-sized leaf");
+          C.getDiagnostics().Report(F->getLocStart(), diag_id);
+          return false;
+        }
         table.leaf_size = sz;
         BMapDeclVisitor visitor(C, table.leaf_desc);
-        if (!visitor.TraverseType(F->getType()))
-          return false;
+        visitor.TraverseType(F->getType());
       } else if (F->getName() == "data") {
         table.max_entries = sz / table.leaf_size;
       }

src/cc/frontends/clang/b_frontend_action.h

@@ -49,6 +49,7 @@ class BMapDeclVisitor : public clang::RecursiveASTVisitor<BMapDeclVisitor> {
   bool VisitBuiltinType(const clang::BuiltinType *T);
   bool VisitTypedefType(const clang::TypedefType *T);
   bool VisitTagType(const clang::TagType *T);
+  bool VisitPointerType(const clang::PointerType *T);
  private:
   clang::ASTContext &C;
   std::string &result_;

src/python/bcc/__init__.py

@@ -518,7 +518,13 @@ class BPF(object):
                     fields.append((t[0], BPF._decode_table_type(t[1]), t[2]))
             else:
                 raise Exception("Failed to decode type %s" % str(t))
-        cls = type(str(desc[0]), (ct.Structure,), dict(_fields_=fields))
+        base = ct.Structure
+        if len(desc) > 2:
+            if desc[2] == u"union":
+                base = ct.Union
+            elif desc[2] == u"struct":
+                base = ct.Structure
+        cls = type(str(desc[0]), (base,), dict(_fields_=fields))
         return cls
 
     def get_table(self, name, keytype=None, leaftype=None):

tests/cc/test_clang.py

@@ -247,7 +247,7 @@ int kprobe____kmalloc(struct pt_regs *ctx, size_t size) {
     if (leaf)
         leaf->size += size;
     return 0;
-}""", debug=4)
+}""")
 
     def test_unop_probe_read(self):
         text = """
@@ -263,5 +263,29 @@ int trace_entry(struct pt_regs *ctx, struct request *req) {
         b = BPF(text=text)
         fn = b.load_func("trace_entry", BPF.KPROBE)
 
+    def test_complex_leaf_types(self):
+        text = """
+struct list;
+struct list {
+  struct list *selfp;
+  struct list *another_selfp;
+  struct list *selfp_array[2];
+};
+struct empty {
+};
+union emptyu {
+  struct empty *em1;
+  struct empty em2;
+  struct empty em3;
+  struct empty em4;
+};
+BPF_TABLE("array", int, struct list, t1, 1);
+BPF_TABLE("array", int, struct list *, t2, 1);
+BPF_TABLE("array", int, union emptyu, t3, 1);
+"""
+        b = BPF(text=text)
+        import ctypes
+        self.assertEqual(ctypes.sizeof(b["t3"].Leaf), 8)
+
 if __name__ == "__main__":
     main()