diff --git a/CMakeLists.txt b/CMakeLists.txt index 61d575e..8bc7371 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -4,8 +4,14 @@ project(pylibbpf) # pybind11 include_directories(${CMAKE_SOURCE_DIR}/src) add_subdirectory(pybind11) -pybind11_add_module(pylibbpf src/core/bpf_program.h src/core/bpf_exception.h - src/bindings/main.cpp src/core/bpf_program.cpp) +pybind11_add_module( + pylibbpf + src/core/bpf_program.h + src/core/bpf_exception.h + src/core/bpf_map.h + src/bindings/main.cpp + src/core/bpf_program.cpp + src/core/bpf_map.cpp) # --- libbpf build rules --- set(LIBBPF_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/libbpf/src) diff --git a/examples/execve.py b/examples/execve.py index 2db23f2..c6a3f9f 100644 --- a/examples/execve.py +++ b/examples/execve.py @@ -1,6 +1,7 @@ import time from ctypes import c_int32, c_int64, c_uint64, c_void_p +from pylibbpf import BpfMap from pythonbpf import BPF, bpf, bpfglobal, map, section from pythonbpf.maps import HashMap @@ -8,7 +9,7 @@ from pythonbpf.maps import HashMap @bpf @map def last() -> HashMap: - return HashMap(key=c_uint64, value=c_uint64, max_entries=1) + return HashMap(key=c_uint64, value=c_uint64, max_entries=3) @bpf @@ -37,6 +38,17 @@ def LICENSE() -> str: b = BPF() b.load_and_attach() +mp = BpfMap(b, last) +mp[42] = 100 # Update entry +value = mp[42] # Lookup entry +del mp[42] # Delete entry +mp[69] = 420 +mp[31] = 42 +mp[21] = 34 +print(mp.items()) +# Iterate through map +for key in mp.keys(): + print(mp[key]) while True: print("running") time.sleep(1) diff --git a/pyproject.toml b/pyproject.toml index 662c6a0..a819ec6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,9 +69,9 @@ line-length = 88 [tool.ruff.lint] extend-select = [ - "B", # flake8-bugbear - "I", # isort - "PGH", # pygrep-hooks - "RUF", # Ruff-specific - "UP", # pyupgrade + "B", # flake8-bugbear + "I", # isort + "PGH", # pygrep-hooks + "RUF", # Ruff-specific + "UP", # pyupgrade ] diff --git a/src/bindings/main.cpp b/src/bindings/main.cpp index ffbd296..ea3a6e3 100644 --- a/src/bindings/main.cpp +++ b/src/bindings/main.cpp @@ -3,10 +3,12 @@ #define MACRO_STRINGIFY(x) STRINGIFY(x) extern "C" { -#include "libbpf.h" +#include } + #include "core/bpf_program.h" #include "core/bpf_exception.h" +#include "core/bpf_map.h" namespace py = pybind11; @@ -21,6 +23,7 @@ PYBIND11_MODULE(pylibbpf, m) { :toctree: _generate BpfProgram + BpfMap BpfException )pbdoc"; @@ -28,14 +31,33 @@ PYBIND11_MODULE(pylibbpf, m) { py::register_exception(m, "BpfException"); py::class_(m, "BpfProgram") - .def(py::init()) - .def(py::init()) - .def("load", &BpfProgram::load) - .def("attach", &BpfProgram::attach) - // .def("detach", &BpfProgram::detach) - .def("load_and_attach", &BpfProgram::load_and_attach) - .def("is_loaded", &BpfProgram::is_loaded) - .def("is_attached", &BpfProgram::is_attached); + .def(py::init()) + .def(py::init()) + .def("load", &BpfProgram::load) + .def("attach", &BpfProgram::attach) + .def("destroy", &BpfProgram::destroy) + .def("load_and_attach", &BpfProgram::load_and_attach) + .def("is_loaded", &BpfProgram::is_loaded) + .def("is_attached", &BpfProgram::is_attached); + + py::class_(m, "BpfMap") + .def(py::init()) + .def("lookup", &BpfMap::lookup) + .def("update", &BpfMap::update) + .def("delete", &BpfMap::delete_elem) + .def("get_next_key", &BpfMap::get_next_key, py::arg("key") = py::none()) + .def("items", &BpfMap::items) + .def("keys", &BpfMap::keys) + .def("values", &BpfMap::values) + .def("get_name", &BpfMap::get_name) + .def("get_type", &BpfMap::get_type) + .def("get_key_size", &BpfMap::get_key_size) + .def("get_value_size", &BpfMap::get_value_size) + .def("get_max_entries", &BpfMap::get_max_entries) + .def("__getitem__", &BpfMap::lookup) + .def("__setitem__", &BpfMap::update) + .def("__delitem__", &BpfMap::delete_elem); + #ifdef VERSION_INFO m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO); diff --git a/src/core/bpf_exception.h b/src/core/bpf_exception.h index efd95d3..4616c19 100644 --- a/src/core/bpf_exception.h +++ b/src/core/bpf_exception.h @@ -6,11 +6,13 @@ class BpfException final : public std::runtime_error { public: - explicit BpfException(const std::string& message) - : std::runtime_error(message) {} + explicit BpfException(const std::string &message) + : std::runtime_error(message) { + } - explicit BpfException(const char* message) - : std::runtime_error(message) {} + explicit BpfException(const char *message) + : std::runtime_error(message) { + } }; #endif // PYLIBBPF_BPF_EXCEPTION_H diff --git a/src/core/bpf_map.cpp b/src/core/bpf_map.cpp new file mode 100644 index 0000000..3e19aa0 --- /dev/null +++ b/src/core/bpf_map.cpp @@ -0,0 +1,211 @@ +#include "bpf_map.h" + +#include "bpf_exception.h" + +BpfMap::BpfMap(BpfProgram *program_, const py::object &map_from_python) { + if (py::isinstance(map_from_python)) { + const auto name = map_from_python.attr("__name__").cast(); + bpf_program = program_; + map_ = bpf_object__find_map_by_name(bpf_program->get_obj(), name.c_str()); + if (!map_) { + throw BpfException("Failed to find map by name"); + } + map_fd = bpf_map__fd(map_); + if (map_fd == -1) { + throw BpfException("Failed to open map File Descriptor"); + } + } else { + throw BpfException("Invalid map object passed to function."); + } +} + +std::vector BpfMap::python_to_bytes(const py::object &obj, size_t size) { + std::vector result(size, 0); + + if (py::isinstance(obj)) { + const auto value = obj.cast(); + std::memcpy(result.data(), &value, std::min(size, sizeof(uint64_t))); + } else if (py::isinstance(obj)) { + const auto bytes_str = obj.cast(); + std::memcpy(result.data(), bytes_str.data(), std::min(size, bytes_str.size())); + } else if (py::isinstance(obj)) { + const auto str_val = obj.cast(); + std::memcpy(result.data(), str_val.data(), std::min(size, str_val.size())); + } + + return result; +} + +py::object BpfMap::bytes_to_python(const std::vector &data) { + // Try to interpret as integer if it's a common integer size + if (data.size() == 4) { + uint32_t value; + std::memcpy(&value, data.data(), 4); + return py::cast(value); + } else if (data.size() == 8) { + uint64_t value; + std::memcpy(&value, data.data(), 8); + return py::cast(value); + } else { + // Return as bytes + return py::bytes(reinterpret_cast(data.data()), data.size()); + } +} + +void BpfMap::update(const py::object &key, const py::object &value) const { + const size_t key_size = bpf_map__key_size(map_); + const size_t value_size = bpf_map__value_size(map_); + + const auto key_bytes = python_to_bytes(key, key_size); + const auto value_bytes = python_to_bytes(value, value_size); + + const int ret = bpf_map__update_elem( + map_, + key_bytes.data(), + key_size, + value_bytes.data(), + value_size, + BPF_ANY); + if (ret != 0) { + throw BpfException("Failed to update map element"); + } +} + +void BpfMap::delete_elem(const py::object &key) const { + const size_t key_size = bpf_map__key_size(map_); + std::vector key_bytes; + key_bytes = python_to_bytes(key, key_size); + + if (const int ret = bpf_map__delete_elem(map_, key_bytes.data(), key_size, BPF_ANY); ret != 0) { + throw BpfException("Failed to delete map element"); + } +} + +py::list BpfMap::get_next_key(const py::object &key) const { + const size_t key_size = bpf_map__key_size(map_); + std::vector next_key(key_size); + + int ret; + if (key.is_none()) { + ret = bpf_map__get_next_key(map_, nullptr, next_key.data(), key_size); + } else { + const auto key_bytes = python_to_bytes(key, key_size); + ret = bpf_map__get_next_key(map_, key_bytes.data(), next_key.data(), key_size); + } + + py::list result; + if (ret == 0) { + result.append(bytes_to_python(next_key)); + } + return result; +} + +py::list BpfMap::keys() const { + py::list result; + const size_t key_size = bpf_map__key_size(map_); + + std::vector key(key_size); + std::vector next_key(key_size); + + int ret = bpf_map__get_next_key(map_, nullptr, key.data(), key_size); + + while (ret == 0) { + result.append(bytes_to_python(key)); + ret = bpf_map__get_next_key(map_, key.data(), next_key.data(), key_size); + key = next_key; + } + + return result; +} + +py::list BpfMap::values() const { + py::list result; + const size_t key_size = bpf_map__key_size(map_); + const size_t value_size = bpf_map__value_size(map_); + + std::vector key(key_size); + std::vector next_key(key_size); + std::vector value(value_size); + + int ret = bpf_map__get_next_key(map_, nullptr, key.data(), key_size); + + while (ret == 0) { + if (bpf_map__lookup_elem(map_, key.data(), key_size, value.data(), value_size, BPF_ANY) == 0) { + result.append(bytes_to_python(value)); + } + ret = bpf_map__get_next_key(map_, key.data(), next_key.data(), key_size); + key = next_key; + } + + return result; +} + +std::string BpfMap::get_name() const { + const char *name = bpf_map__name(map_); + return name ? std::string(name) : ""; +} + +int BpfMap::get_type() const { + return bpf_map__type(map_); +} + +int BpfMap::get_key_size() const { + return bpf_map__key_size(map_); +} + +int BpfMap::get_value_size() const { + return bpf_map__value_size(map_); +} + +int BpfMap::get_max_entries() const { + return bpf_map__max_entries(map_); +} + +py::dict BpfMap::items() const { + py::dict result; + const size_t key_size = bpf_map__key_size(map_); + const size_t value_size = bpf_map__value_size(map_); + + std::vector key(key_size); + std::vector next_key(key_size); + std::vector value(value_size); + + // Get first key + int ret = bpf_map__get_next_key(map_, nullptr, key.data(), key_size); + + while (ret == 0) { + // Lookup value for current key + if (bpf_map__lookup_elem(map_, key.data(), key_size, value.data(), value_size, BPF_ANY) == 0) { + result[bytes_to_python(key)] = bytes_to_python(value); + } + + // Get next key + ret = bpf_map__get_next_key(map_, key.data(), next_key.data(), key_size); + key = next_key; + } + + return result; +} + +py::object BpfMap::lookup(const py::object &key) const { + const __u32 key_size = bpf_map__key_size(map_); + const __u32 value_size = bpf_map__value_size(map_); + + const auto key_bytes = python_to_bytes(key, key_size); + std::vector value_bytes(value_size); + + // The flags field here matters only when spin locks are used which is close to fucking never, so fuck no, + // im not adding it + const int ret = bpf_map__lookup_elem( + map_, + key_bytes.data(), + key_size, + value_bytes.data(), + value_size, + BPF_ANY); + if (ret != 0) { + return py::none(); + } + + return bytes_to_python(value_bytes); +} diff --git a/src/core/bpf_map.h b/src/core/bpf_map.h new file mode 100644 index 0000000..e6ad13a --- /dev/null +++ b/src/core/bpf_map.h @@ -0,0 +1,55 @@ +#ifndef PYLIBBPF_MAPS_H +#define PYLIBBPF_MAPS_H + +#include +#include +#include +#include + +#include "bpf_program.h" + +namespace py = pybind11; + +class BpfMap { +private: + struct bpf_map *map_; + int map_fd = -1; + //TODO: turn below into a shared pointer and ref count it so that there is no resource leakage + BpfProgram *bpf_program; + +public: + BpfMap(BpfProgram *program_, const py::object &map_from_python); + + ~BpfMap() = default; + + [[nodiscard]] py::object lookup(const py::object &key) const; + + void update(const py::object &key, const py::object &value) const; + + void delete_elem(const py::object &key) const; + + py::list get_next_key(const py::object &key = py::none()) const; + + py::dict items() const; + + py::list keys() const; + + py::list values() const; + + [[nodiscard]] std::string get_name() const; + + int get_type() const; + + int get_key_size() const; + + int get_value_size() const; + + int get_max_entries() const; + +private: + static std::vector python_to_bytes(const py::object &obj, size_t size); + + static py::object bytes_to_python(const std::vector &data); +}; + +#endif //PYLIBBPF_MAPS_H diff --git a/src/core/bpf_program.cpp b/src/core/bpf_program.cpp index b8a8ec6..5b9fbb2 100644 --- a/src/core/bpf_program.cpp +++ b/src/core/bpf_program.cpp @@ -3,7 +3,7 @@ #include #include -BpfProgram::BpfProgram(std::string object_path, std::string program_name) +BpfProgram::BpfProgram(std::string object_path, std::string program_name) : obj_(nullptr), prog_(nullptr), link_(nullptr), object_path_(std::move(object_path)), program_name_(std::move(program_name)) { } @@ -15,6 +15,10 @@ BpfProgram::~BpfProgram() { } } +struct bpf_object * BpfProgram::get_obj() const { + return obj_; +} + bool BpfProgram::load() { // Open the eBPF object file obj_ = bpf_object__open_file(object_path_.c_str(), nullptr); @@ -48,8 +52,7 @@ bool BpfProgram::load() { } bool BpfProgram::attach() { - for (auto [prog, link] : programs) - { + for (auto [prog, link]: programs) { if (!prog) { throw BpfException("Program not loaded"); } @@ -66,7 +69,7 @@ bool BpfProgram::attach() { bool BpfProgram::destroy() { bool success = true; - for (auto [prog, link] : programs) { + for (auto [prog, link]: programs) { if (!prog) { throw BpfException("Program not loaded"); } diff --git a/src/core/bpf_program.h b/src/core/bpf_program.h index 7ec73d7..93532da 100644 --- a/src/core/bpf_program.h +++ b/src/core/bpf_program.h @@ -1,7 +1,7 @@ #ifndef PYLIBBPF_BPF_PROGRAM_H #define PYLIBBPF_BPF_PROGRAM_H -#include "libbpf.h" +#include #include #include @@ -9,19 +9,26 @@ namespace py = pybind11; class BpfProgram { private: - struct bpf_object* obj_; - struct bpf_program* prog_; - struct bpf_link* link_; + struct bpf_object *obj_; + struct bpf_program *prog_; + struct bpf_link *link_; std::string object_path_; std::string program_name_; - std::vector> programs; + std::vector > programs; + public: - explicit BpfProgram(std::string object_path, std::string program_name = ""); + explicit BpfProgram(std::string object_path, std::string program_name = ""); + ~BpfProgram(); + struct bpf_object *get_obj() const; + bool load(); + bool attach(); + bool destroy(); + void load_and_attach(); [[nodiscard]] bool is_loaded() const { return obj_ != nullptr; }