add map support

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);
+}

src/core/bpf_map.h

@@ -0,0 +1,55 @@
+#ifndef PYLIBBPF_MAPS_H
+#define PYLIBBPF_MAPS_H
+
+#include <libbpf.h>
+#include <pybind11/pybind11.h>
+#include <vector>
+#include <string>
+
+#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<uint8_t> python_to_bytes(const py::object &obj, size_t size);
+
+    static py::object bytes_to_python(const std::vector<uint8_t> &data);
+};
+
+#endif //PYLIBBPF_MAPS_H

src/core/bpf_program.cpp

@@ -3,7 +3,7 @@
 #include <filesystem>
 #include <utility>
 
-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");
         }

src/core/bpf_program.h

@@ -1,7 +1,7 @@
 #ifndef PYLIBBPF_BPF_PROGRAM_H
 #define PYLIBBPF_BPF_PROGRAM_H
 
-#include "libbpf.h"
+#include <libbpf.h>
 #include <pybind11/stl.h>
 #include <string>
 
@@ -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<std::pair<bpf_program*, bpf_link*>> programs;
+    std::vector<std::pair<bpf_program *, bpf_link *> > 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; }