From c5a485b526ff4bba15d729208c6dca289e3c78df Mon Sep 17 00:00:00 2001 From: Pragyansh Chaturvedi Date: Sat, 18 Oct 2025 20:59:31 +0530 Subject: [PATCH] Reimplement BpfMap --- src/core/bpf_map.cpp | 415 +++++++++++++++++++++------------------ src/core/bpf_map.h | 38 +++- src/core/bpf_program.cpp | 2 + 3 files changed, 261 insertions(+), 194 deletions(-) diff --git a/src/core/bpf_map.cpp b/src/core/bpf_map.cpp index 3e19aa0..b6f9069 100644 --- a/src/core/bpf_map.cpp +++ b/src/core/bpf_map.cpp @@ -1,43 +1,241 @@ #include "bpf_map.h" - +#include "bpf_object.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."); +BpfMap::BpfMap(std::shared_ptr parent, struct bpf_map *raw_map, const std::string &map_name) + : parent_obj_(parent), + map_(raw_map), + map_fd_(-1), + map_name_(map_name), + key_size_(0), + value_size_(0) { + if (!parent) + throw BpfException("Parent BpfObject is null"); + if(!(parent->is_loaded())) + throw BpfException("Parent BpfObject is not loaded"); + if (!raw_map) + throw BpfException("bpf_map pointer is null"); + + map_fd_ = bpf_map__fd(map_); + if (map_fd_ < 0) + throw BpfException("Failed to get file descriptor for map '" + map_name_ + "'"); + + key_size_ = bpf_map__key_size(map_); + value_size_ = bpf_map__value_size(map_); +} + +py::object BpfMap::lookup(const py::object &key) const { + if (map_fd_ < 0) + throw BpfException("Map '" + map_name_ + "' is not initialized properly"); + + BufferManager<> key_buf, value_buf; + auto key_span = key_buf.get_span(key_size_); + auto value_span = value_buf.get_span(value_size_); + + // Convert Python → bytes + python_to_bytes_inplace(key, key_span); + + // The flags field here matters only when spin locks are used. + // Skipping it for now. + const int ret = bpf_map__lookup_elem( + map_, + key_span.data(), + key_size_, + value_span.data(), + value_size_, + BPF_ANY); + if (ret < 0) { + if (ret == -ENOENT) + throw py::key_error("Key not found in map '" + map_name_ + "'"); + throw BpfException( + "Failed to lookup key in map '" + map_name_ + "': " + + std::strerror(-ret) + ); + } + + return bytes_to_python(value_span); +} + +void BpfMap::update(const py::object &key, const py::object &value) const { + if (map_fd_ < 0) + throw BpfException("Map '" + map_name_ + "' is not initialized properly"); + + BufferManager<> key_buf, value_buf; + auto key_span = key_buf.get_span(key_size_); + auto value_span = value_buf.get_span(value_size_); + + python_to_bytes_inplace(key, key_span); + python_to_bytes_inplace(value, value_span); + + const int ret = bpf_map__update_elem( + map_, + key_span.data(), + key_size_, + value_span.data(), + value_size_, + BPF_ANY); + if (ret < 0) { + throw BpfException( + "Failed to update key in map '" + map_name_ + "': " + + std::strerror(-ret) + ); } } -std::vector BpfMap::python_to_bytes(const py::object &obj, size_t size) { - std::vector result(size, 0); +void BpfMap::delete_elem(const py::object &key) const { + if (map_fd_ < 0) + throw BpfException("Map '" + map_name_ + "' is not initialized properly"); - 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())); + BufferManager<> key_buf; + auto key_span = key_buf.get_span(key_size_); + + // Convert Python → bytes + python_to_bytes_inplace(key, key_span); + + const int ret = bpf_map__delete_elem(map_, key_span.data(), key_size_, BPF_ANY); + + if (ret != 0) { + if (ret == -ENOENT) + throw py::key_error("Key not found in map '" + map_name_ + "'"); + throw BpfException( + "Failed to delete key from map '" + map_name_ + "': " + + std::strerror(-ret) + ); + } +} + +py::object BpfMap::get_next_key(const py::object &key) const { + BufferManager<> next_key_buf; + auto next_key = next_key_buf.get_span(key_size_); + + int ret; + if (key.is_none()) { + ret = bpf_map__get_next_key(map_, nullptr, next_key.data(), key_size_); + } else { + BufferManager<> key_buf; + auto key_bytes = key_buf.get_span(key_size_); + python_to_bytes_inplace(key, key_bytes); + ret = bpf_map__get_next_key(map_, key_bytes.data(), next_key.data(), key_size_); } + if (ret < 0) { + if (ret == -ENOENT) { + // No more keys + return py::none(); + } + throw BpfException( + "Failed to get next key in map '" + map_name_ + "': " + + std::strerror(-ret) + ); + } + + return bytes_to_python(next_key); +} + +py::dict BpfMap::items() const { + py::dict result; + + py::object current_key = get_next_key(py::none()); + if (current_key.is_none()) { + return result; + } + + while (!current_key.is_none()) { + try { + py::object value = lookup(current_key); + result[current_key] = value; + current_key = get_next_key(current_key); + } catch (const py::key_error&) { + break; + } + } + return result; } -py::object BpfMap::bytes_to_python(const std::vector &data) { - // Try to interpret as integer if it's a common integer size +py::list BpfMap::keys() const { + py::list result; + + py::object current_key = get_next_key(py::none()); + if (current_key.is_none()) { + return result; + } + + while (!current_key.is_none()) { + result.append(current_key); + current_key = get_next_key(current_key); + } + + return result; +} + +py::list BpfMap::values() const { + py::list result; + + py::object current_key = get_next_key(py::none()); + if (current_key.is_none()) { + return result; + } + + while (!current_key.is_none()) { + try { + py::object value = lookup(current_key); + result.append(value); + current_key = get_next_key(current_key); + } catch (const py::key_error&) { + break; + } + } + + return result; +} + +int BpfMap::get_type() const { + return bpf_map__type(map_); +} + +int BpfMap::get_max_entries() const { + return bpf_map__max_entries(map_); +} + + +// Helper functions +void BpfMap::python_to_bytes_inplace(const py::object &obj, std::span buffer) { + std::fill(buffer.begin(), buffer.end(), 0); + + if (py::isinstance(obj)) { + if (buffer.size() <= sizeof(uint64_t)) { + uint64_t value = obj.cast(); + std::memcpy(buffer.data(), &value, buffer.size()); + } else { + throw BpfException("Integer key/value size exceeds maximum (8 bytes)"); + } + } else if (py::isinstance(obj)) { + std::string bytes_str = obj.cast(); + + if (bytes_str.size() > buffer.size()) { + throw BpfException( + "Bytes size " + std::to_string(bytes_str.size()) + + " exceeds expected size " + std::to_string(buffer.size()) + ); + } + + std::memcpy(buffer.data(), bytes_str.data(), bytes_str.size()); + } else if (py::isinstance(obj)) { + std::string str_val = obj.cast(); + + if (str_val.size() >= buffer.size()) { + throw BpfException("String size exceeds expected size"); + } + + std::memcpy(buffer.data(), str_val.data(), str_val.size()); + buffer[str_val.size()] = '\0'; + } else { + throw BpfException("Unsupported type for BPF map key/value"); + } +} + +py::object BpfMap::bytes_to_python(std::span data) { if (data.size() == 4) { uint32_t value; std::memcpy(&value, data.data(), 4); @@ -47,165 +245,6 @@ py::object BpfMap::bytes_to_python(const std::vector &data) { std::memcpy(&value, data.data(), 8); return py::cast(value); } else { - // Return as bytes - return py::bytes(reinterpret_cast(data.data()), data.size()); + 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 index 550e25d..03302b5 100644 --- a/src/core/bpf_map.h +++ b/src/core/bpf_map.h @@ -5,6 +5,11 @@ #include #include #include +#include +#include +#include +#include +#include class BpfObject; @@ -16,16 +21,37 @@ private: struct bpf_map *map_; int map_fd_; std::string map_name_; + __u32 key_size_, value_size_; + + template + struct BufferManager { + std::array stack_buf; + std::vector heap_buf; + + std::span get_span(size_t size) { + if (size <= StackSize) { + return std::span(stack_buf.data(), size); + } else { + heap_buf.resize(size); + return std::span(heap_buf); + } + } + }; public: - BpfMap(std::shared_ptr, struct bpf_map *raw_map, const std::string &map_name); + BpfMap(std::shared_ptr parent, struct bpf_map *raw_map, const std::string &map_name); ~BpfMap() = default; + BpfMap(const BpfMap&) = delete; + BpfMap& operator=(const BpfMap&) = delete; + BpfMap(BpfMap&&) noexcept = default; + BpfMap& operator=(BpfMap&&) noexcept = 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::object get_next_key(const py::object &key = py::none()) const; py::dict items() const; py::list keys() const; py::list values() const; @@ -33,13 +59,13 @@ public: [[nodiscard]] std::string get_name() const { return map_name_; } [[nodiscard]] int get_fd() const { return map_fd_; } [[nodiscard]] int get_type() const; - [[nodiscard]] int get_key_size() const; - [[nodiscard]] int get_value_size() const; + [[nodiscard]] int get_key_size() const { return key_size_; }; + [[nodiscard]] int get_value_size() const { return value_size_; }; [[nodiscard]] 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); + static void python_to_bytes_inplace(const py::object &obj, std::span buffer); + static py::object bytes_to_python(std::span data); }; #endif //PYLIBBPF_MAPS_H diff --git a/src/core/bpf_program.cpp b/src/core/bpf_program.cpp index c606b5d..1d4202a 100644 --- a/src/core/bpf_program.cpp +++ b/src/core/bpf_program.cpp @@ -10,6 +10,8 @@ BpfProgram::BpfProgram(std::shared_ptr parent, struct bpf_program *ra program_name_(program_name) { if (!parent) throw BpfException("Parent BpfObject is null"); + if(!(parent->is_loaded())) + throw BpfException("Parent BpfObject is not loaded"); if (!raw_prog) throw BpfException("bpf_program pointer is null"); }