Merge pull request #18 from plumgrid/issue16_cleanups

Address comments and add tail_call example

src/bpf.py

@@ -109,9 +109,8 @@ class BPF(object):
                 raise StopIteration()
             return next_key
 
-    def __init__(self, name, dp_file="", dph_file="", text=None, debug=0):
+    def __init__(self, dp_file="", dph_file="", text=None, debug=0):
         self.debug = debug
-        self.name = name
         self.funcs = {}
         if text:
             self.module = lib.bpf_module_create_from_string(text.encode("ascii"), self.debug)
@@ -124,7 +123,7 @@ class BPF(object):
 
     def load_func(self, func_name, prog_type):
         if lib.bpf_function_start(self.module, func_name.encode("ascii")) == None:
-            raise Exception("Unknown program %s" % self.name)
+            raise Exception("Unknown program %s" % func_name)
 
         fd = lib.bpf_prog_load(prog_type,
                 lib.bpf_function_start(self.module, func_name.encode("ascii")),
@@ -142,7 +141,7 @@ class BPF(object):
 
         return fn
 
-    def load_table(self, name, keytype, leaftype):
+    def get_table(self, name, keytype, leaftype):
         map_fd = lib.bpf_table_fd(self.module,
                 ct.c_char_p(name.encode("ascii")))
         if map_fd < 0:
@@ -165,7 +164,9 @@ class BPF(object):
         fn.sock = sock
 
     @staticmethod
-    def attach_filter(fn, ifindex, prio, classid):
+    def attach_classifier(fn, ifname, prio=10, classid=1):
+        with open("/sys/class/net/%s/ifindex" % ifname) as f:
+            ifindex = int(f.read())
         if not isinstance(fn, BPF.Function):
             raise Exception("arg 1 must be of type BPF.Function")
         res = lib.bpf_attach_filter(fn.fd, fn.name, ifindex, prio, classid)

src/cc/b_frontend_action.cc

@@ -50,19 +50,22 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
         string prefix, suffix;
         string map_update_policy = "BPF_ANY";
         if (memb_name == "get") {
-          prefix = "bpf_map_lookup_elem_";
+          prefix = "bpf_map_lookup_elem";
           suffix = ")";
         } else if (memb_name == "put") {
-          prefix = "bpf_map_update_elem_";
+          prefix = "bpf_map_update_elem";
           suffix = ", " + map_update_policy + ")";
         } else if (memb_name == "delete") {
-          prefix = "bpf_map_delete_elem_";
+          prefix = "bpf_map_delete_elem";
+          suffix = ")";
+        } else if (memb_name == "call") {
+          prefix = "bpf_tail_call_";
           suffix = ")";
         } else {
           llvm::errs() << "error: unknown bpf_table operation " << memb_name << "\n";
           return false;
         }
-        prefix += "(bpf_pseudo_fd(1, " + fd + "), ";
+        prefix += "((void *)bpf_pseudo_fd(1, " + fd + "), ";
 
         SourceRange argRange(Call->getArg(0)->getLocStart(),
                              Call->getArg(Call->getNumArgs()-1)->getLocEnd());
@@ -153,6 +156,8 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       map_type = BPF_MAP_TYPE_HASH;
     else if (A->getName() == "maps/array")
       map_type = BPF_MAP_TYPE_ARRAY;
+    else if (A->getName() == "maps/prog")
+      map_type = BPF_MAP_TYPE_PROG_ARRAY;
     table.fd = bpf_create_map(map_type, table.key_size, table.leaf_size, table.max_entries);
     if (table.fd < 0) {
       llvm::errs() << "error: could not open bpf fd\n";

src/cc/bpf_helpers.h

@@ -18,6 +18,7 @@ struct _name##_table_t { \
   _leaf_type * (*get) (_key_type *); \
   int (*put) (_key_type *, _leaf_type *); \
   int (*delete) (_key_type *); \
+  void (*call) (void *, int index); \
   _leaf_type data[_max_entries]; \
 }; \
 __attribute__((section("maps/" _table_type))) \
@@ -34,17 +35,19 @@ unsigned _version SEC("version") = LINUX_VERSION_CODE;
 /* helper functions called from eBPF programs written in C */
 static void *(*bpf_map_lookup_elem)(void *map, void *key) =
 	(void *) BPF_FUNC_map_lookup_elem;
-static int (*bpf_map_update_elem)(void *map, void *key, void *value,
-				  unsigned long long flags) =
+static int (*bpf_map_update_elem)(void *map, void *key, void *value, u64 flags) =
 	(void *) BPF_FUNC_map_update_elem;
 static int (*bpf_map_delete_elem)(void *map, void *key) =
 	(void *) BPF_FUNC_map_delete_elem;
-static int (*bpf_probe_read)(void *dst, unsigned long long size, void *unsafe_ptr) =
+static int (*bpf_probe_read)(void *dst, u64 size, void *unsafe_ptr) =
 	(void *) BPF_FUNC_probe_read;
-static unsigned long long (*bpf_ktime_get_ns)(void) =
+static u64 (*bpf_ktime_get_ns)(void) =
 	(void *) BPF_FUNC_ktime_get_ns;
-static int (*bpf_trace_printk)(const char *fmt, unsigned long long fmt_size, ...) =
+static int (*bpf_trace_printk)(const char *fmt, u64 fmt_size, ...) =
 	(void *) BPF_FUNC_trace_printk;
+static void bpf_tail_call_(u64 map_fd, void *ctx, int index) {
+  ((void (*)(void *, u64, int))BPF_FUNC_tail_call)(ctx, map_fd, index);
+}
 
 /* llvm builtin functions that eBPF C program may use to
  * emit BPF_LD_ABS and BPF_LD_IND instructions
@@ -207,11 +210,6 @@ int bpf_map_delete_elem_(uintptr_t map, void *key) {
   return bpf_map_delete_elem((void *)map, key);
 }
 
-SEC("helpers")
-int bpf_skb_store_bytes_(void *ctx, u64 off, void *from, u64 len, u64 flags) {
-  return bpf_skb_store_bytes(ctx, off, from, len, flags);
-}
-
 SEC("helpers")
 int bpf_l3_csum_replace_(void *ctx, u64 off, u64 from, u64 to, u64 flags) {
   switch (flags & 0xf) {

tests/jit/CMakeLists.txt

@@ -4,6 +4,8 @@ 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_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}
   COMMAND ${TEST_WRAPPER} py_trace1 sudo ${CMAKE_CURRENT_SOURCE_DIR}/trace1.py trace1.b kprobe.b)
 add_test(NAME py_trace2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}

tests/jit/test1.py

@@ -24,10 +24,10 @@ class Leaf(Structure):
 
 class TestBPFSocket(TestCase):
     def setUp(self):
-        b = BPF("test1", arg1, arg2, debug=0)
+        b = BPF(arg1, arg2, debug=0)
         fn = b.load_func("main", BPF.SOCKET_FILTER)
         BPF.attach_socket(fn, "eth0")
-        self.stats = b.load_table("stats", Key, Leaf)
+        self.stats = b.get_table("stats", Key, Leaf)
 
     def test_ping(self):
         cmd = ["ping", "-f", "-c", "100", "172.16.1.1"]

tests/jit/test2.py

@@ -23,12 +23,10 @@ class Leaf(Structure):
 
 class TestBPFSocket(TestCase):
     def setUp(self):
-        b = BPF("test2", arg1, arg2, debug=0)
-        with open("/sys/class/net/eth0/ifindex") as f:
-            ifindex = int(f.read())
+        b = BPF(arg1, arg2, debug=0)
         fn = b.load_func("main", BPF.SCHED_CLS)
-        BPF.attach_filter(fn, ifindex, 10, 1)
-        self.xlate = b.load_table("xlate", Key, Leaf)
+        BPF.attach_classifier(fn, "eth0")
+        self.xlate = b.get_table("xlate", Key, Leaf)
 
     def test_xlate(self):
         key = Key(IPAddress("172.16.1.1").value, IPAddress("172.16.1.2").value)

tests/jit/test3.c

+#include "../../src/cc/bpf_helpers.h"
+
+BPF_TABLE("prog", int, int, jump, 64);
+BPF_TABLE("array", int, u64, stats, 64);
+
+enum states {
+  S_EOP = 1,
+  S_ETHER,
+  S_ARP,
+  S_IP
+};
+
+BPF_EXPORT(parse_ether)
+int parse_ether(struct __sk_buff *skb) {
+  size_t cur = 0;
+  size_t next = cur + 14;
+
+  int key = S_ETHER;
+  u64 *leaf = stats.get(&key);
+  if (leaf) (*leaf)++;
+
+  switch (bpf_dext_pkt(skb, cur + 12, 0, 16)) {
+    case 0x0800: jump.call(skb, S_IP);
+    case 0x0806: jump.call(skb, S_ARP);
+  }
+  jump.call(skb, S_EOP);
+  return 0;
+}
+
+BPF_EXPORT(parse_arp)
+int parse_arp(struct __sk_buff *skb) {
+  size_t cur = 14;  // TODO: get from ctx
+  size_t next = cur + 28;
+
+  int key = S_ARP;
+  u64 *leaf = stats.get(&key);
+  if (leaf) (*leaf)++;
+
+  jump.call(skb, S_EOP);
+  return 0;
+}
+
+BPF_EXPORT(parse_ip)
+int parse_ip(struct __sk_buff *skb) {
+  size_t cur = 14;  // TODO: get from ctx
+  size_t next = cur + 20;
+
+  int key = S_IP;
+  u64 *leaf = stats.get(&key);
+  if (leaf) (*leaf)++;
+
+  jump.call(skb, S_EOP);
+  return 0;
+}
+
+BPF_EXPORT(eop)
+int eop(struct __sk_buff *skb) {
+  int key = S_EOP;
+  u64 *leaf = stats.get(&key);
+  if (leaf) (*leaf)++;
+  return 0;
+}

tests/jit/test3.py

+#!/usr/bin/env python
+
+from ctypes import c_ushort, c_int, c_ulonglong
+from netaddr import IPAddress
+from bpf import BPF
+from socket import socket, AF_INET, SOCK_DGRAM
+import sys
+from time import sleep
+from unittest import main, TestCase
+
+arg1 = sys.argv.pop(1)
+
+S_EOP = 1
+S_ETHER = 2
+S_ARP = 3
+S_IP = 4
+
+class TestBPFSocket(TestCase):
+    def setUp(self):
+        b = BPF(dp_file=arg1, debug=0)
+        ether_fn = b.load_func("parse_ether", BPF.SCHED_CLS)
+        arp_fn = b.load_func("parse_arp", BPF.SCHED_CLS)
+        ip_fn = b.load_func("parse_ip", BPF.SCHED_CLS)
+        eop_fn = b.load_func("eop", BPF.SCHED_CLS)
+        BPF.attach_classifier(ether_fn, "eth0")
+        self.jump = b.get_table("jump", c_int, c_int)
+        self.jump.put(c_int(S_ARP), c_int(arp_fn.fd))
+        self.jump.put(c_int(S_IP), c_int(ip_fn.fd))
+        self.jump.put(c_int(S_EOP), c_int(eop_fn.fd))
+        self.stats = b.get_table("stats", c_int, c_ulonglong)
+
+    def test_jumps(self):
+        udp = socket(AF_INET, SOCK_DGRAM)
+        udp.sendto(b"a" * 10, ("172.16.1.1", 5000))
+        self.assertGreater(self.stats.get(c_int(S_IP)).value, 0)
+        self.assertGreater(self.stats.get(c_int(S_ARP)).value, 0)
+        self.assertGreater(self.stats.get(c_int(S_EOP)).value, 1)
+
+if __name__ == "__main__":
+    main()

tests/jit/trace1.py

@@ -20,11 +20,11 @@ class Leaf(Structure):
 
 class TestKprobe(TestCase):
     def setUp(self):
-        b = BPF("trace1", arg1, arg2, debug=0)
+        b = BPF(arg1, arg2, debug=0)
         fn1 = b.load_func("sys_wr", BPF.KPROBE)
         fn2 = b.load_func("sys_rd", BPF.KPROBE)
         fn3 = b.load_func("sys_bpf", BPF.KPROBE)
-        self.stats = b.load_table("stats", Key, Leaf)
+        self.stats = b.get_table("stats", Key, Leaf)
         BPF.attach_kprobe(fn1, "sys_write", 0, -1)
         BPF.attach_kprobe(fn2, "sys_read", 0, -1)
         BPF.attach_kprobe(fn2, "htab_map_get_next_key", 0, -1)

tests/jit/trace2.py

@@ -28,9 +28,9 @@ class Counters(Structure):
 
 class TestTracingEvent(TestCase):
     def setUp(self):
-        b = BPF("trace2", text=text, debug=0)
+        b = BPF(text=text, debug=0)
         fn = b.load_func("count_sched", BPF.KPROBE)
-        self.stats = b.load_table("stats", Ptr, Counters)
+        self.stats = b.get_table("stats", Ptr, Counters)
         BPF.attach_kprobe(fn, "schedule+50", 0, -1)
 
     def test_sched1(self):

tests/jit/trace3.py

@@ -14,10 +14,10 @@ if len(sys.argv) > 1:
 
 class TestBlkRequest(TestCase):
     def setUp(self):
-        b = BPF("trace3", arg1, arg2, debug=0)
+        b = BPF(arg1, arg2, debug=0)
         fn1 = b.load_func("probe_blk_start_request", BPF.KPROBE)
         fn2 = b.load_func("probe_blk_update_request", BPF.KPROBE)
-        self.latency = b.load_table("latency", c_uint, c_ulong)
+        self.latency = b.get_table("latency", c_uint, c_ulong)
         BPF.attach_kprobe(fn1, "blk_start_request", -1, 0)
         BPF.attach_kprobe(fn2, "blk_update_request", -1, 0)