Merge pull request #8 from pythonbpf/test-workflow

Add support for using struct as value_type in hashmap

Author: r41k0u

src/bindings/main.cpp

@@ -65,6 +65,10 @@ PYBIND11_MODULE(pylibbpf, m) {
       .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("set_value_struct", &BpfMap::set_value_struct,
+           py::arg("struct_name"))
+      .def("get_value_struct_name", &BpfMap::get_value_struct_name)
+      .def("has_struct_value", &BpfMap::has_struct_value)
       .def("__getitem__", &BpfMap::lookup, py::arg("key"))
       .def("__setitem__", &BpfMap::update, py::arg("key"), py::arg("value"))
       .def("__delitem__", &BpfMap::delete_elem, py::arg("key"));

src/core/bpf_map.cpp

@@ -1,6 +1,7 @@
 #include "core/bpf_map.h"
 #include "core/bpf_exception.h"
 #include "core/bpf_object.h"
+#include "utils/struct_parser.h"
 #include <algorithm>
 #include <cerrno>
 #include <cstring>
@@ -25,6 +26,24 @@ BpfMap::BpfMap(std::shared_ptr<BpfObject> parent, struct bpf_map *raw_map,
   value_size_ = bpf_map__value_size(map_);
 }
 
+void BpfMap::set_value_struct(const std::string &struct_name) {
+  auto parent = parent_obj_.lock();
+  if (!parent) {
+    throw BpfException("Parent BpfObject no longer exists");
+  }
+
+  struct_parser_ = parent->get_struct_parser();
+  if (!struct_parser_) {
+    throw BpfException("No struct definitions available in BpfObject");
+  }
+
+  if (!struct_parser_->has_struct(struct_name)) {
+    throw BpfException("Unknown struct: " + struct_name);
+  }
+
+  value_struct_name_ = struct_name;
+}
+
 py::object BpfMap::lookup(const py::object &key) const {
   if (map_fd_ < 0)
     throw BpfException("Map '" + map_name_ + "' is not initialized properly");
@@ -42,7 +61,7 @@ py::object BpfMap::lookup(const py::object &key) const {
                                        value_span.data(), value_size_, BPF_ANY);
   if (ret < 0) {
     if (ret == -ENOENT)
-      throw py::key_error("Key not found in map '" + map_name_ + "'");
+      return py::none();
     throw BpfException("Failed to lookup key in map '" + map_name_ +
                        "': " + std::strerror(-ret));
   }
@@ -215,7 +234,13 @@ void BpfMap::python_to_bytes_inplace(const py::object &obj,
   }
 }
 
-py::object BpfMap::bytes_to_python(std::span<const uint8_t> data) {
+py::object BpfMap::bytes_to_python(std::span<const uint8_t> data) const {
+  //  NOTE: Struct parsing for value type
+  if (struct_parser_ && !value_struct_name_.empty()) {
+    py::bytes py_data(reinterpret_cast<const char *>(data.data()), data.size());
+    return struct_parser_->parse(value_struct_name_, py_data);
+  }
+
   if (data.size() == 4) {
     uint32_t value;
     std::memcpy(&value, data.data(), 4);

src/core/bpf_map.h

@@ -9,6 +9,7 @@
 #include <vector>
 
 class BpfObject;
+class StructParser;
 
 namespace py = pybind11;
 
@@ -20,6 +21,11 @@ class BpfMap : public std::enable_shared_from_this<BpfMap> {
   std::string map_name_;
   __u32 key_size_, value_size_;
 
+  //  TODO: For now, we'll only support struct parsing for value types
+  // later we can extend this to keys
+  std::shared_ptr<StructParser> struct_parser_;
+  std::string value_struct_name_;
+
   template <size_t StackSize = 64> struct BufferManager {
     std::array<uint8_t, StackSize> stack_buf;
     std::vector<uint8_t> heap_buf;
@@ -52,6 +58,7 @@ class BpfMap : public std::enable_shared_from_this<BpfMap> {
   py::dict items() const;
   py::list keys() const;
   py::list values() const;
+  void set_value_struct(const std::string &struct_name);
 
   [[nodiscard]] std::string get_name() const { return map_name_; }
   [[nodiscard]] int get_fd() const { return map_fd_; }
@@ -62,11 +69,17 @@ class BpfMap : public std::enable_shared_from_this<BpfMap> {
   [[nodiscard]] std::shared_ptr<BpfObject> get_parent() const {
     return parent_obj_.lock();
   }
+  [[nodiscard]] std::string get_value_struct_name() const {
+    return value_struct_name_;
+  }
+  [[nodiscard]] bool has_struct_value() const {
+    return !value_struct_name_.empty();
+  }
 
 private:
   static void python_to_bytes_inplace(const py::object &obj,
                                       std::span<uint8_t> buffer);
-  static py::object bytes_to_python(std::span<const uint8_t> data);
+  py::object bytes_to_python(std::span<const uint8_t> data) const;
 };
 
 #endif // PYLIBBPF_BPF_MAP_H