Address some of Issue #16

* Introduce Function object in bpf.py
 - prog load types (kprobe, socket, etc.) are independent
 - fd is a member of each function, to be used later (e.g. tail call table)
* Allow functions to be defined inline in the .py as a text argument
* Rename BPFProgram to BPFModule, which should make more sense
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/bpf.py

@@ -3,21 +3,25 @@ import os
 
 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.restype = ct.c_int
 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 @@ class BPF(object):
                 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

-#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
 }

src/cc/bpf_program.cc → src/cc/bpf_module.cc

@@ -46,7 +46,7 @@
 #include "type_check.h"
 #include "codegen_llvm.h"
 #include "b_frontend_action.h"
-#include "bpf_program.h"
+#include "bpf_module.h"
 #include "kbuild_helper.h"
 
 namespace ebpf {
@@ -91,7 +91,7 @@ class MyMemoryManager : public SectionMemoryManager {
   map<string, tuple<uint8_t *, uintptr_t>> *sections_;
 };
 
-BPFProgram::BPFProgram(unsigned flags)
+BPFModule::BPFModule(unsigned flags)
     : flags_(flags), ctx_(new LLVMContext) {
   LLVMInitializeBPFTarget();
   LLVMInitializeBPFTargetMC();
@@ -100,12 +100,12 @@ BPFProgram::BPFProgram(unsigned flags)
   LLVMLinkInMCJIT(); /* call empty function to force linking of MCJIT */
 }
 
-BPFProgram::~BPFProgram() {
+BPFModule::~BPFModule() {
   engine_.reset();
   ctx_.reset();
 }
 
-int BPFProgram::load_file_module(unique_ptr<llvm::Module> *mod, const string &file) {
+int BPFModule::load_file_module(unique_ptr<llvm::Module> *mod, const string &file, bool in_memory) {
   using namespace clang;
 
   struct utsname un;
@@ -118,11 +118,18 @@ int BPFProgram::load_file_module(unique_ptr<llvm::Module> *mod, const string &fi
   if (!dstack.ok())
     return -1;
 
-  string abs_file = string(dstack.cwd()) + "/" + file;
-  if (file[0] == '/')
-    abs_file = file;
+  string abs_file;
+  if (in_memory) {
+    abs_file = "<bcc-memory-buffer>";
+  } else {
+    if (file.substr(0, 1) == "/")
+      abs_file = file;
+    else
+      abs_file = string(dstack.cwd()) + "/" + file;
+  }
+
   vector<const char *> flags_cstr({"-O0", "-emit-llvm", "-I", dstack.cwd(),
-                                   "-x", "c", "-c", abs_file.c_str()});
+                                  "-x", "c", "-c", abs_file.c_str()});
 
   KBuildHelper kbuild_helper;
   vector<string> kflags;
@@ -172,6 +179,13 @@ int BPFProgram::load_file_module(unique_ptr<llvm::Module> *mod, const string &fi
                                           const_cast<const char **>(ccargs.data()) + ccargs.size(), diags))
     return -1;
 
+  if (in_memory) {
+    invocation1->getPreprocessorOpts().addRemappedFile("<bcc-memory-buffer>",
+                                                       llvm::MemoryBuffer::getMemBuffer(file).release());
+    invocation1->getFrontendOpts().Inputs.clear();
+    invocation1->getFrontendOpts().Inputs.push_back(FrontendInputFile("<bcc-memory-buffer>", IK_C));
+  }
+
   CompilerInstance compiler1;
   compiler1.setInvocation(invocation1.release());
   compiler1.createDiagnostics();
@@ -208,11 +222,41 @@ int BPFProgram::load_file_module(unique_ptr<llvm::Module> *mod, const string &fi
   return 0;
 }
 
+// load an entire c file as a module
+int BPFModule::load_cfile(const string &file, bool in_memory) {
+  unique_ptr<Module> mod;
+  if (load_file_module(&mod, file, in_memory))
+    return -1;
+  mod_ = &*mod;
+
+  mod_->setDataLayout("e-m:e-i64:64-f80:128-n8:16:32:64-S128");
+  mod_->setTargetTriple("bpf-pc-linux");
+
+  for (auto fn = mod_->getFunctionList().begin(); fn != mod_->getFunctionList().end(); ++fn)
+    fn->addFnAttr(Attribute::AlwaysInline);
+
+  string err;
+  engine_ = unique_ptr<ExecutionEngine>(EngineBuilder(move(mod))
+      .setErrorStr(&err)
+      .setMCJITMemoryManager(make_unique<MyMemoryManager>(&sections_))
+      .setMArch("bpf")
+      .create());
+  if (!engine_) {
+    fprintf(stderr, "Could not create ExecutionEngine: %s\n", err.c_str());
+    return -1;
+  }
+
+  return 0;
+}
+
+// NOTE: this is a duplicate of the above, but planning to deprecate if we
+// settle on clang as the frontend
+
 // Load in a pre-built list of functions into the initial Module object, then
 // build an ExecutionEngine.
-int BPFProgram::load_includes(const string &tmpfile) {
+int BPFModule::load_includes(const string &tmpfile) {
   unique_ptr<Module> mod;
-  if (load_file_module(&mod, tmpfile))
+  if (load_file_module(&mod, tmpfile, false))
     return -1;
   mod_ = &*mod;
 
@@ -236,13 +280,13 @@ int BPFProgram::load_includes(const string &tmpfile) {
   return 0;
 }
 
-void BPFProgram::dump_ir() {
+void BPFModule::dump_ir() {
   legacy::PassManager PM;
   PM.add(createPrintModulePass(outs()));
   PM.run(*mod_);
 }
 
-int BPFProgram::parse() {
+int BPFModule::parse() {
   int rc;
 
   proto_parser_ = make_unique<ebpf::cc::Parser>(proto_filename_);
@@ -283,7 +327,7 @@ int BPFProgram::parse() {
   return 0;
 }
 
-int BPFProgram::finalize() {
+int BPFModule::finalize() {
   if (verifyModule(*mod_, &errs())) {
     if (flags_ & 1)
       dump_ir();
@@ -305,7 +349,7 @@ int BPFProgram::finalize() {
   return 0;
 }
 
-uint8_t * BPFProgram::start(const string &name) const {
+uint8_t * BPFModule::start(const string &name) const {
   auto section = sections_.find("." + name);
   if (section == sections_.end())
     return nullptr;
@@ -313,7 +357,7 @@ uint8_t * BPFProgram::start(const string &name) const {
   return get<0>(section->second);
 }
 
-size_t BPFProgram::size(const string &name) const {
+size_t BPFModule::size(const string &name) const {
   auto section = sections_.find("." + name);
   if (section == sections_.end())
     return 0;
@@ -321,7 +365,7 @@ size_t BPFProgram::size(const string &name) const {
   return get<1>(section->second);
 }
 
-char * BPFProgram::license() const {
+char * BPFModule::license() const {
   auto section = sections_.find("license");
   if (section == sections_.end())
     return nullptr;
@@ -329,7 +373,7 @@ char * BPFProgram::license() const {
   return (char *)get<0>(section->second);
 }
 
-unsigned BPFProgram::kern_version() const {
+unsigned BPFModule::kern_version() const {
   auto section = sections_.find("version");
   if (section == sections_.end())
     return 0;
@@ -337,7 +381,7 @@ unsigned BPFProgram::kern_version() const {
   return *(unsigned *)get<0>(section->second);
 }
 
-int BPFProgram::table_fd(const string &name) const {
+int BPFModule::table_fd(const string &name) const {
   int fd = codegen_ ? codegen_->get_table_fd(name) : -1;
   if (fd >= 0) return fd;
   auto table_it = tables_->find(name);
@@ -345,7 +389,7 @@ int BPFProgram::table_fd(const string &name) const {
   return table_it->second.fd;
 }
 
-int BPFProgram::load(const string &filename, const string &proto_filename) {
+int BPFModule::load(const string &filename, const string &proto_filename) {
   if (!sections_.empty()) {
     fprintf(stderr, "Program already initialized\n");
     return -1;
@@ -354,7 +398,7 @@ int BPFProgram::load(const string &filename, const string &proto_filename) {
   proto_filename_ = proto_filename;
   if (proto_filename_.empty()) {
     // direct load of .b file
-    if (int rc = load_includes(filename_))
+    if (int rc = load_cfile(filename_, false))
       return rc;
   } else {
     // old lex .b file
@@ -366,4 +410,18 @@ int BPFProgram::load(const string &filename, const string &proto_filename) {
   return 0;
 }
 
+int BPFModule::load_string(const string &text) {
+  if (!sections_.empty()) {
+    fprintf(stderr, "Program already initialized\n");
+    return -1;
+  }
+  filename_ = "<memory>";
+  if (int rc = load_cfile(text, true))
+    return rc;
+
+  if (int rc = finalize())
+    return rc;
+  return 0;
+}
+
 } // namespace ebpf

src/cc/bpf_program.h → src/cc/bpf_module.h

@@ -38,19 +38,21 @@ class CodegenLLVM;
 class Parser;
 }
 
-class BPFProgram {
+class BPFModule {
  private:
   int init_engine();
   int parse();
   int finalize();
   void dump_ir();
-  int load_file_module(std::unique_ptr<llvm::Module> *mod, const std::string &file);
+  int load_file_module(std::unique_ptr<llvm::Module> *mod, const std::string &file, bool in_memory);
   int load_includes(const std::string &tmpfile);
+  int load_cfile(const std::string &file, bool in_memory);
   int kbuild_flags(const char *uname_release, std::vector<std::string> *cflags);
  public:
-  BPFProgram(unsigned flags);
-  ~BPFProgram();
+  BPFModule(unsigned flags);
+  ~BPFModule();
   int load(const std::string &filename, const std::string &proto_filename);
+  int load_string(const std::string &text);
   uint8_t * start(const std::string &name) const;
   size_t size(const std::string &name) const;
   int table_fd(const std::string &name) const;

tests/jit/CMakeLists.txt

@@ -6,9 +6,7 @@ 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_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_b WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-  COMMAND ${TEST_WRAPPER} py_trace2_b sudo ${CMAKE_CURRENT_SOURCE_DIR}/trace2.py trace2.b kprobe.b)
-add_test(NAME py_trace2_c WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-  COMMAND ${TEST_WRAPPER} py_trace2_c sudo ${CMAKE_CURRENT_SOURCE_DIR}/trace2.py trace2.c)
+add_test(NAME py_trace2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_trace2 sudo ${CMAKE_CURRENT_SOURCE_DIR}/trace2.py)
 add_test(NAME py_trace3_c WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_trace3_c sudo ${CMAKE_CURRENT_SOURCE_DIR}/trace3.py trace3.c)

tests/jit/test1.py

@@ -24,9 +24,10 @@ class Leaf(Structure):
 
 class TestBPFSocket(TestCase):
     def setUp(self):
-        self.prog = BPF("main", arg1, arg2, debug=0)
-        self.prog.attach("eth0")
-        self.stats = self.prog.table("stats", Key, Leaf)
+        b = BPF("test1", 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)
 
     def test_ping(self):
         cmd = ["ping", "-f", "-c", "100", "172.16.1.1"]

tests/jit/test2.py

@@ -23,20 +23,20 @@ class Leaf(Structure):
 
 class TestBPFSocket(TestCase):
     def setUp(self):
-        self.prog = BPF("main", arg1, arg2,
-                BPF.BPF_PROG_TYPE_SCHED_CLS, debug=0)
+        b = BPF("test2", arg1, arg2, debug=0)
         with open("/sys/class/net/eth0/ifindex") as f:
             ifindex = int(f.read())
-        self.prog.attach_filter(ifindex, 10, 1)
+        fn = b.load_func("main", BPF.SCHED_CLS)
+        BPF.attach_filter(fn, ifindex, 10, 1)
+        self.xlate = b.load_table("xlate", Key, Leaf)
 
     def test_xlate(self):
-        xlate = self.prog.table("xlate", Key, Leaf)
         key = Key(IPAddress("172.16.1.1").value, IPAddress("172.16.1.2").value)
         leaf = Leaf(IPAddress("192.168.1.1").value, IPAddress("192.168.1.2").value, 0)
-        xlate.put(key, leaf)
+        self.xlate.put(key, leaf)
         udp = socket(AF_INET, SOCK_DGRAM)
         udp.sendto(b"a" * 10, ("172.16.1.1", 5000))
-        leaf = xlate.get(key)
+        leaf = self.xlate.get(key)
         self.assertGreater(leaf.xlated_pkts, 0)
 
 if __name__ == "__main__":

tests/jit/trace1.py

@@ -20,15 +20,14 @@ class Leaf(Structure):
 
 class TestKprobe(TestCase):
     def setUp(self):
-        self.prog = BPF("trace1", arg1, arg2,
-                prog_type=BPF.BPF_PROG_TYPE_KPROBE, debug=0)
-        self.prog.load("sys_wr")
-        self.prog.load("sys_rd")
-        self.prog.load("sys_bpf")
-        self.stats = self.prog.table("stats", Key, Leaf)
-        self.prog.attach_kprobe("sys_write", "sys_wr", 0, -1)
-        self.prog.attach_kprobe("sys_read", "sys_rd", 0, -1)
-        self.prog.attach_kprobe("htab_map_get_next_key", "sys_bpf", 0, -1)
+        b = BPF("trace1", 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)
+        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)
 
     def test_trace1(self):
         with open("/dev/null", "a") as f:

tests/jit/trace2.py

@@ -6,10 +6,20 @@ from time import sleep
 import sys
 from unittest import main, TestCase
 
-arg1 = sys.argv.pop(1)
-arg2 = ""
-if len(sys.argv) > 1:
-  arg2 = sys.argv.pop(1)
+text = """
+#include <linux/ptrace.h>
+#include "../../src/cc/bpf_helpers.h"
+struct Ptr { u64 ptr; };
+struct Counters { u64 stat1; };
+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};
+  stats.data[(u64)&key].stat1++;
+  return 0;
+}
+"""
 
 class Ptr(Structure):
     _fields_ = [("ptr", c_ulong)]
@@ -18,11 +28,10 @@ class Counters(Structure):
 
 class TestTracingEvent(TestCase):
     def setUp(self):
-        self.prog = BPF("trace2", arg1, arg2,
-                prog_type=BPF.BPF_PROG_TYPE_KPROBE, debug=0)
-        self.prog.load("count_sched")
-        self.stats = self.prog.table("stats", Ptr, Counters)
-        self.prog.attach_kprobe("schedule+50", "count_sched", 0, -1)
+        b = BPF("trace2", text=text, debug=0)
+        fn = b.load_func("count_sched", BPF.KPROBE)
+        self.stats = b.load_table("stats", Ptr, Counters)
+        BPF.attach_kprobe(fn, "schedule+50", 0, -1)
 
     def test_sched1(self):
         for i in range(0, 100):

tests/jit/trace3.py

@@ -14,22 +14,22 @@ if len(sys.argv) > 1:
 
 class TestBlkRequest(TestCase):
     def setUp(self):
-        self.prog = BPF("trace3", arg1, arg2,
-                prog_type=BPF.BPF_PROG_TYPE_KPROBE, debug=0)
-        self.prog.load("probe_blk_start_request")
-        self.prog.load("probe_blk_update_request")
-        self.latency = self.prog.table("latency", c_uint, c_ulong)
-        self.prog.attach_kprobe("blk_start_request", "probe_blk_start_request", 0, -1)
-        self.prog.attach_kprobe("blk_update_request", "probe_blk_update_request", 0, -1)
+        b = BPF("trace3", 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)
+        BPF.attach_kprobe(fn1, "blk_start_request", -1, 0)
+        BPF.attach_kprobe(fn2, "blk_update_request", -1, 0)
 
     def test_blk1(self):
         import subprocess
         import os
+        # use /opt instead of /tmp so that it hits a real disk
         for i in range(0, 2):
-            with open("/srv/trace3.txt", "w") as f:
-                f.write("a" * 4096 * 4096)
+            subprocess.call(["dd", "if=/dev/zero", "of=/opt/trace3.txt",
+                             "count=1024", "bs=4096"])
             subprocess.call(["sync"])
-        os.unlink("/srv/trace3.txt")
+        os.unlink("/opt/trace3.txt")
         for key in self.latency.iter():
             leaf = self.latency.get(key)
             print("latency %u:" % key.value, "count %u" % leaf.value)