add map support

This commit is contained in:
2025-09-21 23:55:10 +05:30
parent 0a27d5a520
commit e5a946a767
9 changed files with 349 additions and 31 deletions

211
src/core/bpf_map.cpp Normal file
View File

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