Commit ac56b0ef authored by 4ast's avatar 4ast

Merge pull request #17 from plumgrid/issue16

Address some of Issue #16
parents c848e557 414d6145
......@@ -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")
......
......@@ -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})
......
#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);
}
}
......@@ -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
}
......
......@@ -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,9 +118,16 @@ 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] == '/')
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()});
......@@ -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
......@@ -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;
......
......@@ -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)
......@@ -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"]
......
......@@ -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__":
......
......@@ -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:
......
......@@ -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):
......
......@@ -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)
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment