Merge pull request #24 from plumgrid/bblanco_dev

Add support for packet field write

src/bpf.py

@@ -149,7 +149,7 @@ class BPF(object):
         return BPF.Table(self, map_fd, keytype, leaftype)
 
     @staticmethod
-    def attach_socket(fn, dev):
+    def attach_raw_socket(fn, dev):
         if not isinstance(fn, BPF.Function):
             raise Exception("arg 1 must be of type BPF.Function")
         sock = lib.bpf_open_raw_sock(dev.encode("ascii"))

src/cc/b_frontend_action.cc

@@ -148,12 +148,42 @@ bool BTypeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *Arr) {
   return true;
 }
 
-bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
-  Expr *Base = E->getBase()->IgnoreImplicit();
+bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
+  if (!E->isAssignmentOp())
+    return true;
+  Expr *LHS = E->getLHS()->IgnoreImplicit();
+  Expr *RHS = E->getRHS()->IgnoreImplicit();
+  if (MemberExpr *Memb = dyn_cast<MemberExpr>(LHS)) {
+    if (DeclRefExpr *Base = dyn_cast<DeclRefExpr>(Memb->getBase()->IgnoreImplicit())) {
+      if (DeprecatedAttr *A = Base->getDecl()->getAttr<DeprecatedAttr>()) {
+        if (A->getMessage() == "packet") {
+          if (FieldDecl *F = dyn_cast<FieldDecl>(Memb->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 rhs = rewriter_.getRewrittenText(SourceRange(RHS->getLocStart(), RHS->getLocEnd()));
+            string text = "bpf_dins_pkt(skb, (u64)" + base + "+" + to_string(ofs >> 3)
+                + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ", " + rhs + ")";
+            rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
+          }
+        }
+      }
+    }
+  }
+  return true;
+}
+bool BTypeVisitor::VisitImplicitCastExpr(ImplicitCastExpr *E) {
+  // use dext only for RValues
+  if (E->getCastKind() != CK_LValueToRValue)
+    return true;
+  MemberExpr *Memb = dyn_cast<MemberExpr>(E->IgnoreImplicit());
+  if (!Memb)
+    return true;
+  Expr *Base = Memb->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())) {
+        if (FieldDecl *F = dyn_cast<FieldDecl>(Memb->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()));

src/cc/b_frontend_action.h

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

src/cc/bpf_helpers.h

@@ -254,6 +254,11 @@ 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 incr_cksum_l3(expr, oldval, newval) \
+  bpf_l3_csum_replace_(skb, (u64)expr, oldval, newval, sizeof(newval))
+#define incr_cksum_l4(expr, oldval, newval, is_pseudo) \
+  bpf_l4_csum_replace_(skb, (u64)expr, oldval, newval, ((is_pseudo & 0x1) << 4) | sizeof(newval))
+
 #define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
 
 #endif

src/cc/proto.h

@@ -7,42 +7,42 @@ struct ethernet_t {
 } __attribute__((packed));
 
 struct dot1q_t {
-  u32 pri:3;
-  u32 cfi:1;
-  u32 vlanid:12;
-  u32 type:16;
+  u16 pri:3;
+  u16 cfi:1;
+  u16 vlanid:12;
+  u16 type;
 } __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
+  u8 ver:4;           // byte 0
+  u8 hlen:4;
+  u8 tos;
+  u16 tlen;
+  u16 identification; // byte 4
+  u16 ffo_unused:1;
+  u16 df:1;
+  u16 mf:1;
+  u16 foffset:13;
+  u8 ttl;             // byte 8
+  u8 nextp;
+  u16 hchecksum;
+  u32 src;            // byte 12
+  u32 dst;            // byte 16
 } __attribute__((packed));
 
 struct udp_t {
-  u32 sport:16;
-  u32 dport:16;
-  u32 length:16;
-  u32 crc:16;
+  u16 sport;
+  u16 dport;
+  u16 length;
+  u16 crc;
 } __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
+  u16 src_port;   // byte 0
+  u16 dst_port;
+  u32 seq_num;    // byte 4
+  u32 ack_num;    // byte 8
+  u8 offset:4;    // byte 12
   u8 reserved:4;
   u8 flag_cwr:1;
   u8 flag_ece:1;
@@ -52,7 +52,7 @@ struct tcp_t {
   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;
+  u16 rcv_wnd;
+  u16 cksum;      // byte 16
+  u16 urg_ptr;
 } __attribute__((packed));

tests/jit/CMakeLists.txt

@@ -2,8 +2,10 @@ add_test(NAME py_test1_b WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test1_b namespace ${CMAKE_CURRENT_SOURCE_DIR}/test1.py test1.b proto.b)
 add_test(NAME py_test1_c WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test1_c namespace ${CMAKE_CURRENT_SOURCE_DIR}/test1.py test1.c)
-add_test(NAME py_test2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-  COMMAND ${TEST_WRAPPER} py_test2 namespace ${CMAKE_CURRENT_SOURCE_DIR}/test2.py test2.b proto.b)
+add_test(NAME py_test2_b WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test2_b namespace ${CMAKE_CURRENT_SOURCE_DIR}/test2.py test2.b proto.b)
+add_test(NAME py_test2_c WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test2_c namespace ${CMAKE_CURRENT_SOURCE_DIR}/test2.py test2.c)
 add_test(NAME py_test3 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test3 namespace ${CMAKE_CURRENT_SOURCE_DIR}/test3.py test3.c)
 add_test(NAME py_trace1 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}

tests/jit/test1.b

@@ -12,7 +12,7 @@ struct skbuff {
   u32 type:32;
 };
 
-u32 main(struct skbuff *skb) {
+u32 on_packet(struct skbuff *skb) {
   u32 ret:32 = 0;
 
   goto proto::ethernet;

tests/jit/test1.c

@@ -12,8 +12,8 @@ struct IPLeaf {
 
 BPF_TABLE("hash", struct IPKey, struct IPLeaf, stats, 256);
 
-BPF_EXPORT(main)
-int _main(struct __sk_buff *skb) {
+BPF_EXPORT(on_packet)
+int on_packet(struct __sk_buff *skb) {
   BEGIN(ethernet);
 
   PROTO(ethernet) {

tests/jit/test1.py

@@ -25,8 +25,8 @@ class Leaf(Structure):
 class TestBPFSocket(TestCase):
     def setUp(self):
         b = BPF(arg1, arg2, debug=0)
-        fn = b.load_func("main", BPF.SOCKET_FILTER)
-        BPF.attach_socket(fn, "eth0")
+        fn = b.load_func("on_packet", BPF.SOCKET_FILTER)
+        BPF.attach_raw_socket(fn, "eth0")
         self.stats = b.get_table("stats", Key, Leaf)
 
     def test_ping(self):

tests/jit/test2.b

@@ -17,7 +17,7 @@ struct skbuff {
   u32 type:32;
 };
 
-u32 main (struct skbuff *skb) {
+u32 on_packet (struct skbuff *skb) {
   u32 ret:32 = 1;
 
   u32 orig_dip:32 = 0;

tests/jit/test2.c

+#include "../../src/cc/bpf_helpers.h"
+#include "../../src/cc/proto.h"
+struct IPKey {
+  u32 dip;
+  u32 sip;
+};
+struct IPLeaf {
+  u32 xdip;
+  u32 xsip;
+  u64 xlated_pkts;
+};
+BPF_TABLE("hash", struct IPKey, struct IPLeaf, xlate, 1024);
+
+BPF_EXPORT(on_packet)
+int on_packet(struct __sk_buff *skb) {
+  BEGIN(ethernet);
+
+  u32 orig_dip = 0;
+  u32 orig_sip = 0;
+  struct IPLeaf *xleaf;
+
+  PROTO(ethernet) {
+    switch (ethernet->type) {
+      case 0x0800: goto ip;
+      case 0x8100: goto dot1q;
+    }
+    goto EOP;
+  }
+
+  PROTO(dot1q) {
+    switch (dot1q->type) {
+      case 0x0800: goto ip;
+    }
+    goto EOP;
+  }
+
+  PROTO(ip) {
+    orig_dip = ip->dst;
+    orig_sip = ip->src;
+    struct IPKey key = {.dip=orig_dip, .sip=orig_sip};
+    xleaf = xlate.lookup(&key);
+    if (xleaf) {
+      ip->dst = xleaf->xdip;
+      incr_cksum_l3(&ip->hchecksum, orig_dip, xleaf->xdip);
+      ip->src = xleaf->xsip;
+      incr_cksum_l3(&ip->hchecksum, orig_sip, xleaf->xsip);
+      lock_xadd(&xleaf->xlated_pkts, 1);
+    }
+    switch (ip->nextp) {
+      case 6: goto tcp;
+      case 17: goto udp;
+    }
+    goto EOP;
+  }
+
+  PROTO(udp) {
+    if (xleaf) {
+      incr_cksum_l4(&udp->crc, orig_dip, xleaf->xdip, 1);
+      incr_cksum_l4(&udp->crc, orig_sip, xleaf->xsip, 1);
+    }
+    goto EOP;
+  }
+
+  PROTO(tcp) {
+    if (xleaf) {
+      incr_cksum_l4(&tcp->cksum, orig_dip, xleaf->xdip, 1);
+      incr_cksum_l4(&tcp->cksum, orig_sip, xleaf->xsip, 1);
+    }
+    goto EOP;
+  }
+
+EOP:
+  return 0;
+}

tests/jit/test2.py

@@ -23,8 +23,8 @@ class Leaf(Structure):
 
 class TestBPFSocket(TestCase):
     def setUp(self):
-        b = BPF(arg1, arg2, debug=0)
-        fn = b.load_func("main", BPF.SCHED_CLS)
+        b = BPF(arg1, arg2, debug=1)
+        fn = b.load_func("on_packet", BPF.SCHED_CLS)
         BPF.attach_classifier(fn, "eth0")
         self.xlate = b.get_table("xlate", Key, Leaf)
 

tests/wrapper.sh.in

@@ -28,6 +28,7 @@ function ns_run() {
   sudo ip netns exec $ns tc qdisc add dev eth0 root prio
   sudo ip netns exec $ns ip addr add dev eth0 172.16.1.2/24
   sudo ip netns exec $ns ip link set eth0 up
+  sudo ip netns exec $ns ethtool -K eth0 tx off
   sudo ip addr add dev $ns.out 172.16.1.1/24
   sudo ip link set $ns.out up
   sudo bash -c "PYTHONPATH=$PYTHONPATH LD_LIBRARY_PATH=$LD_LIBRARY_PATH ip netns exec $ns $cmd $1 $2"