Merge pull request #22 from plumgrid/bblanco_dev

Add packet read syntax to c parser

src/cc/b_frontend_action.cc

@@ -15,6 +15,7 @@ int bpf_create_map(enum bpf_map_type map_type, int key_size, int value_size,
 namespace ebpf {
 using std::map;
 using std::string;
+using std::to_string;
 using std::unique_ptr;
 using namespace clang;
 
@@ -39,6 +40,11 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
       if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
         if (!A->getName().startswith("maps"))
           return true;
+
+        SourceRange argRange(Call->getArg(0)->getLocStart(),
+                             Call->getArg(Call->getNumArgs()-1)->getLocEnd());
+        string args = rewriter_.getRewrittenText(argRange);
+
         // find the table fd, which was opened at declaration time
         auto table_it = tables_.find(Ref->getDecl()->getName());
         if (table_it == tables_.end()) {
@@ -46,13 +52,31 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
               << "initialized handle for bpf_table";
           return false;
         }
-        string fd = std::to_string(table_it->second.fd);
+        string fd = to_string(table_it->second.fd);
         string prefix, suffix;
         string map_update_policy = "BPF_ANY";
-        if (memb_name == "get") {
+        string txt;
+        if (memb_name == "lookup_or_init") {
+          string map_update_policy = "BPF_NOEXIST";
+          string name = Ref->getDecl()->getName();
+          string arg0 = rewriter_.getRewrittenText(SourceRange(Call->getArg(0)->getLocStart(),
+                                                               Call->getArg(0)->getLocEnd()));
+          string arg1 = rewriter_.getRewrittenText(SourceRange(Call->getArg(1)->getLocStart(),
+                                                               Call->getArg(1)->getLocEnd()));
+          string lookup = "bpf_map_lookup_elem_(bpf_pseudo_fd(1, " + fd + ")";
+          string update = "bpf_map_update_elem_(bpf_pseudo_fd(1, " + fd + ")";
+          txt  = "({typeof(" + name + ".leaf) *leaf = " + lookup + ", " + arg0 + "); ";
+          txt += "if (!leaf) {";
+          txt += " " + update + ", " + arg0 + ", " + arg1 + ", " + map_update_policy + ");";
+          txt += " leaf = " + lookup + ", " + arg0 + ");";
+          txt += " if (!leaf) return -1;";
+          txt += "}";
+          txt += "leaf;})";
+        } else {
+          if (memb_name == "lookup") {
             prefix = "bpf_map_lookup_elem";
             suffix = ")";
-        } else if (memb_name == "put") {
+          } else if (memb_name == "update") {
             prefix = "bpf_map_update_elem";
             suffix = ", " + map_update_policy + ")";
           } else if (memb_name == "delete") {
@@ -67,10 +91,9 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
           }
           prefix += "((void *)bpf_pseudo_fd(1, " + fd + "), ";
 
-        SourceRange argRange(Call->getArg(0)->getLocStart(),
-                             Call->getArg(Call->getNumArgs()-1)->getLocEnd());
-        string args = rewriter_.getRewrittenText(argRange);
-        rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), prefix + args + suffix);
+          txt = prefix + args + suffix;
+        }
+        rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), txt);
         return true;
       }
     }
@@ -102,7 +125,7 @@ bool BTypeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *Arr) {
                 << "initialized handle for bpf_table";
             return false;
           }
-          string fd = std::to_string(table_it->second.fd);
+          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());
@@ -125,6 +148,25 @@ bool BTypeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *Arr) {
   return true;
 }
 
+bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
+  Expr *Base = E->getBase()->IgnoreImplicit();
+  if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Base)) {
+    if (DeprecatedAttr *A = Ref->getDecl()->getAttr<DeprecatedAttr>()) {
+      if (A->getMessage() == "packet") {
+        if (FieldDecl *F = dyn_cast<FieldDecl>(E->getMemberDecl())) {
+          uint64_t ofs = C.getFieldOffset(F);
+          uint64_t sz = F->isBitField() ? F->getBitWidthValue(C) : C.getTypeSize(F->getType());
+          string base = rewriter_.getRewrittenText(SourceRange(Base->getLocStart(), Base->getLocEnd()));
+          string text = "bpf_dext_pkt(skb, (u64)" + base + "+" + to_string(ofs >> 3)
+              + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ")";
+          rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
+        }
+      }
+    }
+  }
+  return true;
+}
+
 // Open table FDs when bpf tables (as denoted by section("maps*") attribute)
 // are declared.
 bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
@@ -167,11 +209,6 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
   }
   return true;
 }
-bool BTypeVisitor::VisitDeclRefExpr(DeclRefExpr *E) {
-  //ValueDecl *D = E->getDecl();
-  //BPFTableAttr *A = D->getAttr<BPFTableAttr>();
-  return true;
-}
 
 BTypeConsumer::BTypeConsumer(ASTContext &C, Rewriter &rewriter, map<string, BPFTable> &tables)
     : visitor_(C, rewriter, tables) {

src/cc/b_frontend_action.h

@@ -40,7 +40,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitVarDecl(clang::VarDecl *Decl);
   bool VisitArraySubscriptExpr(clang::ArraySubscriptExpr *E);
-  bool VisitDeclRefExpr(clang::DeclRefExpr *E);
+  bool VisitMemberExpr(clang::MemberExpr *E);
 
  private:
   clang::ASTContext &C;

src/cc/bpf_helpers.h

@@ -15,8 +15,9 @@
 struct _name##_table_t { \
   _key_type key; \
   _leaf_type leaf; \
-  _leaf_type * (*get) (_key_type *); \
-  int (*put) (_key_type *, _leaf_type *); \
+  _leaf_type * (*lookup) (_key_type *); \
+  _leaf_type * (*lookup_or_init) (_key_type *, _leaf_type *); \
+  int (*update) (_key_type *, _leaf_type *); \
   int (*delete) (_key_type *); \
   void (*call) (void *, int index); \
   _leaf_type data[_max_entries]; \
@@ -24,6 +25,19 @@ struct _name##_table_t { \
 __attribute__((section("maps/" _table_type))) \
 struct _name##_table_t _name
 
+// packet parsing state machine helpers
+#define STATE_MACHINE(name) \
+  BPF_EXPORT(name) int _##name(struct __sk_buff *skb)
+#define BEGIN(next) \
+  u64 _parse_cursor = 0; \
+  goto next
+
+#define PROTO(name) \
+  goto EOP; \
+name: ; \
+  struct name##_t *name __attribute__((deprecated("packet"))) = (void *)_parse_cursor; \
+  _parse_cursor += sizeof(*name);
+
 // export this function to llvm by putting it into a specially named section
 //#define BPF_EXPORT(_ret, _name, ...) SEC("." #_name) _ret _name(__VA_ARGS__)
 #define BPF_EXPORT(_name) __attribute__((section("." #_name)))
@@ -240,4 +254,6 @@ int bpf_l4_csum_replace_(void *ctx, u64 off, u64 from, u64 to, u64 flags) {
   return bpf_l4_csum_replace(ctx, off, from, to, flags);
 }
 
+#define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
+
 #endif

src/cc/bpf_module.cc

@@ -129,6 +129,7 @@ int BPFModule::load_file_module(unique_ptr<llvm::Module> *mod, const string &fil
   }
 
   vector<const char *> flags_cstr({"-O0", "-emit-llvm", "-I", dstack.cwd(),
+                                   "-Wno-deprecated-declarations",
                                    "-x", "c", "-c", abs_file.c_str()});
 
   KBuildHelper kbuild_helper;

src/cc/proto.h

+#include <linux/types.h>
+
+struct ethernet_t {
+  u64 dst:48;
+  u64 src:48;
+  u32 type:16;
+} __attribute__((packed));
+
+struct dot1q_t {
+  u32 pri:3;
+  u32 cfi:1;
+  u32 vlanid:12;
+  u32 type:16;
+} __attribute__((packed));
+
+struct ip_t {
+  u32 ver:4;              // byte 0
+  u32 hlen:4;
+  u32 tos:8;
+  u32 tlen:16;
+  u32 identification:16;  // byte 4
+  u32 ffo_unused:1;
+  u32 df:1;
+  u32 mf:1;
+  u32 foffset:13;
+  u32 ttl:8;              // byte 8
+  u32 nextp:8;
+  u32 hchecksum:16;
+  u32 src:32;             // byte 12
+  u32 dst:32;             // byte 16
+} __attribute__((packed));
+
+struct udp_t {
+  u32 sport:16;
+  u32 dport:16;
+  u32 length:16;
+  u32 crc:16;
+} __attribute__((packed));
+
+struct tcp_t {
+  u16 src_port:16;  // byte 0
+  u16 dst_port:16;
+  u32 seq_num:32;   // byte 4
+  u32 ack_num:32;   // byte 8
+  u8 offset:4;      // byte 12
+  u8 reserved:4;
+  u8 flag_cwr:1;
+  u8 flag_ece:1;
+  u8 flag_urg:1;
+  u8 flag_ack:1;
+  u8 flag_psh:1;
+  u8 flag_rst:1;
+  u8 flag_syn:1;
+  u8 flag_fin:1;
+  u16 rcv_wnd:16;
+  u16 cksum:16;     // byte 16
+  u16 urg_ptr:16;
+} __attribute__((packed));

tests/jit/test1.c

 #include "../../src/cc/bpf_helpers.h"
+#include "../../src/cc/proto.h"
 
 struct IPKey {
   u32 dip;
@@ -12,91 +13,37 @@ struct IPLeaf {
 BPF_TABLE("hash", struct IPKey, struct IPLeaf, stats, 256);
 
 BPF_EXPORT(main)
-int foo(struct __sk_buff *skb) {
-  size_t next = 0, cur = 0;
-ethernet:
-{
-  cur = next; next += 14;
+int _main(struct __sk_buff *skb) {
+  BEGIN(ethernet);
 
-  switch (bpf_dext_pkt(skb, cur + 12, 0, 16)) {
-    case 0x800: goto ip;
+  PROTO(ethernet) {
+    switch (ethernet->type) {
+      case 0x0800: goto ip;
       case 0x8100: goto dot1q;
-    default: goto EOP;
     }
-}
-dot1q:
-{
-  cur = next; next += 4;
-
-  switch (bpf_dext_pkt(skb, cur + 2, 0, 16)) {
+  }
+  PROTO(dot1q) {
+    switch (dot1q->type) {
       case 0x0800: goto ip;
-    default: goto EOP;
     }
-}
-
-ip:
-{
-  cur = next; next += 20;
-
-  int rx = 0;
-  int tx = 0;
-  struct IPKey key = {0};
-  if (bpf_dext_pkt(skb, cur + 16, 0, 32) > bpf_dext_pkt(skb, cur + 12, 0, 32)) {
-    key.sip = bpf_dext_pkt(skb, cur + 12, 0, 32);
-    key.dip = bpf_dext_pkt(skb, cur + 16, 0, 32);
+  }
+  PROTO(ip) {
+    int rx = 0, tx = 0;
+    struct IPKey key;
+    if (ip->dst > ip->src) {
+      key.dip = ip->dst;
+      key.sip = ip->src;
       rx = 1;
     } else {
-    key.dip = bpf_dext_pkt(skb, cur + 12, 0, 32);
-    key.sip = bpf_dext_pkt(skb, cur + 16, 0, 32);
+      key.dip = ip->src;
+      key.sip = ip->dst;
       tx = 1;
     }
-  // try to get here:
-  //stats[key].rx_pkts += rx;
-  //stats[key].tx_pkts += tx;
-  // or here:
-  //struct IPLeaf *leaf = stats[key];
-  //if (leaf) {
-  //  __sync_fetch_and_add(&leaf->rx_pkts, rx);
-  //  __sync_fetch_and_add(&leaf->tx_pkts, tx);
-  //}
-  struct IPLeaf *leaf;
-  leaf = stats.get(&key);
-  if (!leaf) {
     struct IPLeaf zleaf = {0};
-    stats.put(&key, &zleaf);
-    leaf = stats.get(&key);
-  }
-  if (leaf) {
-    __sync_fetch_and_add(&leaf->rx_pkts, rx);
-    __sync_fetch_and_add(&leaf->tx_pkts, tx);
-  }
-
-  switch (bpf_dext_pkt(skb, cur + 9, 0, 8)) {
-    case 6: goto tcp;
-    case 17: goto udp;
-    //case 47: goto gre;
-    default: goto EOP;
-  }
-}
-
-udp:
-{
-  cur = next; next += 8;
-
-  switch (bpf_dext_pkt(skb, cur + 2, 0, 16)) {
-    //case 8472: goto vxlan;
-    //case 4789: goto vxlan;
-    default: goto EOP;
+    struct IPLeaf *leaf = stats.lookup_or_init(&key, &zleaf);
+    lock_xadd(&leaf->rx_pkts, rx);
+    lock_xadd(&leaf->tx_pkts, tx);
   }
-}
-
-tcp:
-{
-  cur = next; next += 20;
-
-  goto EOP;
-}
-
 EOP:
   return 0;
 }

tests/jit/test3.c

@@ -16,7 +16,7 @@ int parse_ether(struct __sk_buff *skb) {
   size_t next = cur + 14;
 
   int key = S_ETHER;
-  u64 *leaf = stats.get(&key);
+  u64 *leaf = stats.lookup(&key);
   if (leaf) (*leaf)++;
 
   switch (bpf_dext_pkt(skb, cur + 12, 0, 16)) {
@@ -33,7 +33,7 @@ int parse_arp(struct __sk_buff *skb) {
   size_t next = cur + 28;
 
   int key = S_ARP;
-  u64 *leaf = stats.get(&key);
+  u64 *leaf = stats.lookup(&key);
   if (leaf) (*leaf)++;
 
   jump.call(skb, S_EOP);
@@ -46,7 +46,7 @@ int parse_ip(struct __sk_buff *skb) {
   size_t next = cur + 20;
 
   int key = S_IP;
-  u64 *leaf = stats.get(&key);
+  u64 *leaf = stats.lookup(&key);
   if (leaf) (*leaf)++;
 
   jump.call(skb, S_EOP);
@@ -56,7 +56,7 @@ int parse_ip(struct __sk_buff *skb) {
 BPF_EXPORT(eop)
 int eop(struct __sk_buff *skb) {
   int key = S_EOP;
-  u64 *leaf = stats.get(&key);
+  u64 *leaf = stats.lookup(&key);
   if (leaf) (*leaf)++;
   return 0;
 }

tests/jit/trace2.c

@@ -7,11 +7,7 @@ BPF_TABLE("hash", struct Ptr, struct Counters, stats, 1024);
 BPF_EXPORT(count_sched)
 int count_sched(struct pt_regs *ctx) {
   struct Ptr key = {.ptr=ctx->bx};
-#if 1
-  stats.data[(u64)&key].stat1++;
-#else
   struct Counters zleaf = {0};
-  stats.upsert(&key, &zleaf)->stat1++;
-#endif
+  stats.lookup_or_init(&key, &zleaf)->stat1++;
   return 0;
 }

tests/jit/trace3.c

@@ -30,14 +30,14 @@ BPF_EXPORT(probe_blk_start_request)
 int probe_blk_start_request(struct pt_regs *ctx) {
   struct Request rq = {.rq = ctx->di};
   struct Time tm = {.start = bpf_ktime_get_ns()};
-  requests.put(&rq, &tm);
+  requests.update(&rq, &tm);
   return 0;
 }
 
 BPF_EXPORT(probe_blk_update_request)
 int probe_blk_update_request(struct pt_regs *ctx) {
   struct Request rq = {.rq = ctx->di};
-  struct Time *tm = requests.get(&rq);
+  struct Time *tm = requests.lookup(&rq);
   if (!tm) return 0;
   u64 delta = bpf_ktime_get_ns() - tm->start;
   requests.delete(&rq);
@@ -46,6 +46,9 @@ int probe_blk_update_request(struct pt_regs *ctx) {
   u32 index = (lg * 64 + (delta - base) * 64 / base) * 3 / 64;
   if (index >= SLOTS)
     index = SLOTS - 1;
-  __sync_fetch_and_add(&latency.data[(u64)&index], 1);
+
+  u64 zero = 0;
+  u64 *val = latency.lookup_or_init(&index, &zero);
+  lock_xadd(val, 1);
   return 0;
 }