add map support

Author: varun-r-mallya

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)

examples/execve.py

@@ -1,14 +1,15 @@
 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
 
 
 @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)

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
 ]

src/bindings/main.cpp

@@ -3,10 +3,12 @@
 #define MACRO_STRINGIFY(x) STRINGIFY(x)
 
 extern "C" {
-#include "libbpf.h"
+#include <libbpf.h>
 }
+
 #include "core/bpf_program.h"
 #include "core/bpf_exception.h"
+#include "core/bpf_map.h"
 
 namespace py = pybind11;
 
@@ -21,21 +23,41 @@ PYBIND11_MODULE(pylibbpf, m) {
            :toctree: _generate
 
            BpfProgram
+           BpfMap
            BpfException
     )pbdoc";
 
     // Register the custom exception
     py::register_exception<BpfException>(m, "BpfException");
 
     py::class_<BpfProgram>(m, "BpfProgram")
-     .def(py::init<const std::string&>())
-     .def(py::init<const std::string&, const std::string&>())
-     .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<const std::string &>())
+            .def(py::init<const std::string &, const std::string &>())
+            .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_<BpfMap>(m, "BpfMap")
+            .def(py::init<BpfProgram *, py::object &>())
+            .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);

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

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<py::function>(map_from_python)) {
+        const auto name = map_from_python.attr("__name__").cast<std::string>();
+        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<uint8_t> BpfMap::python_to_bytes(const py::object &obj, size_t size) {
+    std::vector<uint8_t> result(size, 0);
+
+    if (py::isinstance<py::int_>(obj)) {
+        const auto value = obj.cast<uint64_t>();
+        std::memcpy(result.data(), &value, std::min(size, sizeof(uint64_t)));
+    } else if (py::isinstance<py::bytes>(obj)) {
+        const auto bytes_str = obj.cast<std::string>();
+        std::memcpy(result.data(), bytes_str.data(), std::min(size, bytes_str.size()));
+    } else if (py::isinstance<py::str>(obj)) {
+        const auto str_val = obj.cast<std::string>();
+        std::memcpy(result.data(), str_val.data(), std::min(size, str_val.size()));
+    }
+
+    return result;
+}
+
+py::object BpfMap::bytes_to_python(const std::vector<uint8_t> &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<const char *>(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<uint8_t> 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<uint8_t> 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<uint8_t> key(key_size);
+    std::vector<uint8_t> 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<uint8_t> key(key_size);
+    std::vector<uint8_t> next_key(key_size);
+    std::vector<uint8_t> 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<uint8_t> key(key_size);
+    std::vector<uint8_t> next_key(key_size);
+    std::vector<uint8_t> 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<uint8_t> 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);
+}