Address some of Issue #16

src/bpf.py

@@ -3,21 +3,25 @@
 
 lib = ct.cdll.LoadLibrary("libbpfprog.so")
 
-lib.bpf_program_create.restype = ct.c_void_p
-lib.bpf_program_create.argtypes = [ct.c_char_p, ct.c_char_p, ct.c_uint]
-lib.bpf_program_destroy.restype = None
-lib.bpf_program_destroy.argtypes = [ct.c_void_p]
-lib.bpf_program_start.restype = ct.c_void_p
-lib.bpf_program_start.argtypes = [ct.c_void_p, ct.c_char_p]
-lib.bpf_program_size.restype = ct.c_size_t
-lib.bpf_program_size.argtypes = [ct.c_void_p, ct.c_char_p]
-lib.bpf_program_license.restype = ct.c_char_p
-lib.bpf_program_license.argtypes = [ct.c_void_p]
-lib.bpf_program_kern_version.restype = ct.c_uint
-lib.bpf_program_kern_version.argtypes = [ct.c_void_p]
-lib.bpf_program_table_fd.restype = ct.c_int
-lib.bpf_program_table_fd.argtypes = [ct.c_void_p, ct.c_char_p]
-
+# keep in sync with bpf_common.h
+lib.bpf_module_create.restype = ct.c_void_p
+lib.bpf_module_create.argtypes = [ct.c_char_p, ct.c_char_p, ct.c_uint]
+lib.bpf_module_create_from_string.restype = ct.c_void_p
+lib.bpf_module_create_from_string.argtypes = [ct.c_char_p, ct.c_uint]
+lib.bpf_module_destroy.restype = None
+lib.bpf_module_destroy.argtypes = [ct.c_void_p]
+lib.bpf_module_license.restype = ct.c_char_p
+lib.bpf_module_license.argtypes = [ct.c_void_p]
+lib.bpf_module_kern_version.restype = ct.c_uint
+lib.bpf_module_kern_version.argtypes = [ct.c_void_p]
+lib.bpf_function_start.restype = ct.c_void_p
+lib.bpf_function_start.argtypes = [ct.c_void_p, ct.c_char_p]
+lib.bpf_function_size.restype = ct.c_size_t
+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]
+
+# keep in sync with libbpf.h
 lib.bpf_get_next_key.restype = ct.c_int
 lib.bpf_get_next_key.argtypes = [ct.c_int, ct.c_void_p, ct.c_void_p]
 lib.bpf_lookup_elem.restype = ct.c_int
@@ -40,42 +44,16 @@
 lib.bpf_attach_kprobe.argtypes = [ct.c_int, ct.c_char_p, ct.c_char_p, ct.c_int, ct.c_int, ct.c_int]
 
 class BPF(object):
-    BPF_PROG_TYPE_SOCKET_FILTER = 1
-    BPF_PROG_TYPE_KPROBE = 2
-    BPF_PROG_TYPE_SCHED_CLS = 3
-    BPF_PROG_TYPE_SCHED_ACT = 4
-    def __init__(self, name, dp_file, dph_file,
-            prog_type=BPF_PROG_TYPE_SOCKET_FILTER,
-            debug=0):
-        self.debug = debug
-        self.name = name
-        self.prog_type = prog_type
-        self.fd = {}
-        self.prog = lib.bpf_program_create(dp_file.encode("ascii"),
-                dph_file.encode("ascii"), self.debug)
-
-        if self.prog == None:
-            raise Exception("Failed to compile BPF program %s" % dp_file)
-
-        if prog_type == BPF.BPF_PROG_TYPE_KPROBE:
-            return
-
-        self.load(self.name)
+    SOCKET_FILTER = 1
+    KPROBE = 2
+    SCHED_CLS = 3
+    SCHED_ACT = 4
 
-    def load(self, prog_name):
-        if lib.bpf_program_start(self.prog, prog_name.encode("ascii")) == None:
-            raise Exception("Unknown program %s" % self.name)
-
-        self.fd[prog_name] = lib.bpf_prog_load(self.prog_type,
-                lib.bpf_program_start(self.prog, prog_name.encode("ascii")),
-                lib.bpf_program_size(self.prog, prog_name.encode("ascii")),
-                lib.bpf_program_license(self.prog),
-                lib.bpf_program_kern_version(self.prog))
-
-        if self.fd[prog_name] < 0:
-            print((ct.c_char * 65536).in_dll(lib, "bpf_log_buf").value)
-            #print(ct.c_char_p.in_dll(lib, "bpf_log_buf").value)
-            raise Exception("Failed to load BPF program %s" % self.name)
+    class Function(object):
+        def __init__(self, bpf, name, fd):
+            self.bpf = bpf
+            self.name = name
+            self.fd = fd
 
     class Table(object):
         def __init__(self, bpf, map_fd, keytype, leaftype):
@@ -131,44 +109,88 @@ def next(self, key):
                 raise StopIteration()
             return next_key
 
-    def table(self, name, keytype, leaftype):
-        map_fd = lib.bpf_program_table_fd(self.prog,
+    def __init__(self, name, 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)
+        else:
+            self.module = lib.bpf_module_create(dp_file.encode("ascii"),
+                    dph_file.encode("ascii"), self.debug)
+
+        if self.module == None:
+            raise Exception("Failed to compile BPF module %s" % dp_file)
+
+    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)
+
+        fd = lib.bpf_prog_load(prog_type,
+                lib.bpf_function_start(self.module, func_name.encode("ascii")),
+                lib.bpf_function_size(self.module, func_name.encode("ascii")),
+                lib.bpf_module_license(self.module),
+                lib.bpf_module_kern_version(self.module))
+
+        if fd < 0:
+            print((ct.c_char * 65536).in_dll(lib, "bpf_log_buf").value)
+            #print(ct.c_char_p.in_dll(lib, "bpf_log_buf").value)
+            raise Exception("Failed to load BPF program %s" % func_name)
+
+        fn = BPF.Function(self, func_name, fd)
+        self.funcs[func_name] = fn
+
+        return fn
+
+    def load_table(self, name, keytype, leaftype):
+        map_fd = lib.bpf_table_fd(self.module,
                 ct.c_char_p(name.encode("ascii")))
         if map_fd < 0:
             raise Exception("Failed to find BPF Table %s" % name)
         return BPF.Table(self, map_fd, keytype, leaftype)
 
-    def attach(self, dev, prog_name=None):
-        prog_name = prog_name or self.name
-        self.sock = lib.bpf_open_raw_sock(dev.encode("ascii"))
-        if self.sock < 0:
+    @staticmethod
+    def attach_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"))
+        if sock < 0:
             errstr = os.strerror(ct.get_errno())
             raise Exception("Failed to open raw device %s: %s" % (dev, errstr))
-        res = lib.bpf_attach_socket(self.sock, self.fd[prog_name])
+        res = lib.bpf_attach_socket(sock, fn.fd)
         if res < 0:
             errstr = os.strerror(ct.get_errno())
             raise Exception("Failed to attach BPF to device %s: %s"
                     % (dev, errstr))
+        fn.sock = sock
 
-    def attach_filter(self, ifindex, prio, classid, prog_name=None):
-        prog_name = prog_name or self.name
-        res = lib.bpf_attach_filter(self.fd[prog_name], self.name.encode("ascii"), ifindex, prio, classid)
+    @staticmethod
+    def attach_filter(fn, ifindex, prio, classid):
+        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)
         if res < 0:
             raise Exception("Failed to filter with BPF")
 
-    def attach_kprobe(self, event, prog_name, pid=-1, cpu=0, group_fd=-1):
+    @staticmethod
+    def attach_kprobe(fn, event, pid=-1, cpu=0, group_fd=-1):
+        if not isinstance(fn, BPF.Function):
+            raise Exception("arg 1 must be of type BPF.Function")
         ev_name = "p_" + event.replace("+", "_")
         desc = "p:kprobes/%s %s" % (ev_name, event)
-        res = lib.bpf_attach_kprobe(self.fd[prog_name], ev_name.encode("ascii"),
+        res = lib.bpf_attach_kprobe(fn.fd, ev_name.encode("ascii"),
                 desc.encode("ascii"), pid, cpu, group_fd)
         if res < 0:
             raise Exception("Failed to attach BPF to kprobe")
         return res
 
-    def attach_kretprobe(self, event, prog_name, pid=-1, cpu=0, group_fd=-1):
+    @staticmethod
+    def attach_kretprobe(fn, event, pid=-1, cpu=0, group_fd=-1):
+        if not isinstance(fn, BPF.Function):
+            raise Exception("arg 1 must be of type BPF.Function")
         ev_name = "r_" + event.replace("+", "_")
         desc = "r:kprobes/%s %s" % (ev_name, event)
-        res = lib.bpf_attach_kprobe(self.fd[prog_name], ev_name.encode("ascii"),
+        res = lib.bpf_attach_kprobe(fn.fd, ev_name.encode("ascii"),
                 desc.encode("ascii"), pid, cpu, group_fd)
         if res < 0:
             raise Exception("Failed to attach BPF to kprobe")

src/cc/CMakeLists.txt

@@ -9,7 +9,7 @@ ADD_FLEX_BISON_DEPENDENCY(Lexer Parser)
 
 set(CMAKE_SHARED_LINKER_FLAGS "-static-libstdc++ -Wl,--exclude-libs=ALL")
 
-add_library(bpfprog SHARED bpf_common.cc bpf_program.cc codegen_llvm.cc
+add_library(bpfprog SHARED bpf_common.cc bpf_module.cc codegen_llvm.cc
   node.cc parser.cc printer.cc type_check.cc libbpf.c b_frontend_action.cc
   kbuild_helper.cc
   ${BISON_Parser_OUTPUTS} ${FLEX_Lexer_OUTPUTS})

src/cc/bpf_common.cc

@@ -1,50 +1,59 @@
-#include "cc/bpf_program.h"
+#include "cc/bpf_module.h"
 #include "cc/bpf_common.h"
 
 extern "C" {
-void * bpf_program_create(const char *filename, const char *proto_filename, unsigned flags) {
-  auto prog = new ebpf::BPFProgram(flags);
-  if (prog->load(filename, proto_filename) != 0) {
-    delete prog;
+void * bpf_module_create(const char *filename, const char *proto_filename, unsigned flags) {
+  auto mod = new ebpf::BPFModule(flags);
+  if (mod->load(filename, proto_filename) != 0) {
+    delete mod;
     return nullptr;
   }
-  return prog;
+  return mod;
 }
 
-void bpf_program_destroy(void *program) {
-  auto prog = static_cast<ebpf::BPFProgram *>(program);
-  if (!prog) return;
-  delete prog;
+void * bpf_module_create_from_string(const char *text, unsigned flags) {
+  auto mod = new ebpf::BPFModule(flags);
+  if (mod->load_string(text) != 0) {
+    delete mod;
+    return nullptr;
+  }
+  return mod;
+}
+
+void bpf_module_destroy(void *program) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return;
+  delete mod;
 }
 
-void * bpf_program_start(void *program, const char *name) {
-  auto prog = static_cast<ebpf::BPFProgram *>(program);
-  if (!prog) return nullptr;
-  return prog->start(name);
+void * bpf_function_start(void *program, const char *name) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return nullptr;
+  return mod->start(name);
 }
 
-size_t bpf_program_size(void *program, const char *name) {
-  auto prog = static_cast<ebpf::BPFProgram *>(program);
-  if (!prog) return 0;
-  return prog->size(name);
+size_t bpf_function_size(void *program, const char *name) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->size(name);
 }
 
-char * bpf_program_license(void *program) {
-  auto prog = static_cast<ebpf::BPFProgram *>(program);
-  if (!prog) return nullptr;
-  return prog->license();
+char * bpf_module_license(void *program) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return nullptr;
+  return mod->license();
 }
 
-unsigned bpf_program_kern_version(void *program) {
-  auto prog = static_cast<ebpf::BPFProgram *>(program);
-  if (!prog) return 0;
-  return prog->kern_version();
+unsigned bpf_module_kern_version(void *program) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->kern_version();
 }
 
-int bpf_program_table_fd(void *program, const char *table_name) {
-  auto prog = static_cast<ebpf::BPFProgram *>(program);
-  if (!prog) return -1;
-  return prog->table_fd(table_name);
+int bpf_table_fd(void *program, const char *table_name) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return -1;
+  return mod->table_fd(table_name);
 }
 
 }

src/cc/bpf_common.h

@@ -6,13 +6,13 @@
 extern "C" {
 #endif
 
-void * bpf_program_create(const char *filename, const char *proto_filename, unsigned flags);
-void bpf_program_destroy(void *program);
-void * bpf_program_start(void *program, const char *name);
-size_t bpf_program_size(void *program, const char *name);
-char * bpf_program_license(void *program);
-unsigned bpf_program_kern_version(void *program);
-int bpf_program_table_fd(void *program, const char *table_name);
+void * bpf_module_create(const char *filename, const char *proto_filename, unsigned flags);
+void bpf_module_destroy(void *program);
+char * bpf_module_license(void *program);
+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);
 
 #ifdef __cplusplus
 }