Allow the use of custom keys in BPF_HASH_OF_MAPS

src/cc/api/BPF.cc

@@ -833,13 +833,6 @@ BPFStackBuildIdTable BPF::get_stackbuildid_table(const std::string &name, bool u
   return BPFStackBuildIdTable({}, use_debug_file, check_debug_file_crc, get_bsymcache());
 }
 
-BPFMapInMapTable BPF::get_map_in_map_table(const std::string& name) {
-  TableStorage::iterator it;
-  if (bpf_module_->table_storage().Find(Path({bpf_module_->id(), name}), it))
-    return BPFMapInMapTable(it->second);
-  return BPFMapInMapTable({});
-}
-
 BPFSockmapTable BPF::get_sockmap_table(const std::string& name) {
   TableStorage::iterator it;
   if (bpf_module_->table_storage().Find(Path({bpf_module_->id(), name}), it))

src/cc/api/BPF.h

@@ -211,8 +211,13 @@ class BPF {
   BPFStackBuildIdTable get_stackbuildid_table(const std::string &name,
                                               bool use_debug_file = true,
                                               bool check_debug_file_crc = true);
-
-  BPFMapInMapTable get_map_in_map_table(const std::string& name);
+  template <class KeyType>
+  BPFMapInMapTable<KeyType> get_map_in_map_table(const std::string& name){
+      TableStorage::iterator it;
+      if (bpf_module_->table_storage().Find(Path({bpf_module_->id(), name}), it))
+        return BPFMapInMapTable<KeyType>(it->second);
+      return BPFMapInMapTable<KeyType>({});
+  }
 
   bool add_module(std::string module);
 

src/cc/api/BPFTable.cc

@@ -689,27 +689,6 @@ StatusTuple BPFXskmapTable::remove_value(const int& index) {
     return StatusTuple::OK();
 }
 
-BPFMapInMapTable::BPFMapInMapTable(const TableDesc& desc)
-    : BPFTableBase<int, int>(desc) {
-    if(desc.type != BPF_MAP_TYPE_ARRAY_OF_MAPS &&
-       desc.type != BPF_MAP_TYPE_HASH_OF_MAPS)
-      throw std::invalid_argument("Table '" + desc.name +
-                                  "' is not a map-in-map table");
-}
-
-StatusTuple BPFMapInMapTable::update_value(const int& index,
-                                           const int& inner_map_fd) {
-    if (!this->update(const_cast<int*>(&index), const_cast<int*>(&inner_map_fd)))
-      return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
-    return StatusTuple::OK();
-}
-
-StatusTuple BPFMapInMapTable::remove_value(const int& index) {
-    if (!this->remove(const_cast<int*>(&index)))
-      return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
-    return StatusTuple::OK();
-}
-
 BPFSockmapTable::BPFSockmapTable(const TableDesc& desc)
     : BPFTableBase<int, int>(desc) {
     if(desc.type != BPF_MAP_TYPE_SOCKMAP)

src/cc/api/BPFTable.h

@@ -479,12 +479,26 @@ class BPFXskmapTable : public BPFTableBase<int, int> {
   StatusTuple remove_value(const int& index);
 };
 
-class BPFMapInMapTable : public BPFTableBase<int, int> {
-public:
-  BPFMapInMapTable(const TableDesc& desc);
-
-  StatusTuple update_value(const int& index, const int& inner_map_fd);
-  StatusTuple remove_value(const int& index);
+template <class KeyType>
+class BPFMapInMapTable : public BPFTableBase<KeyType, int> {
+ public:
+  BPFMapInMapTable(const TableDesc& desc) : BPFTableBase<KeyType, int>(desc) {
+    if (desc.type != BPF_MAP_TYPE_ARRAY_OF_MAPS &&
+        desc.type != BPF_MAP_TYPE_HASH_OF_MAPS)
+      throw std::invalid_argument("Table '" + desc.name +
+                                  "' is not a map-in-map table");
+  }
+  virtual StatusTuple update_value(const KeyType& key, const int& inner_map_fd) {
+    if (!this->update(const_cast<KeyType*>(&key),
+                      const_cast<int*>(&inner_map_fd)))
+      return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
+    return StatusTuple::OK();
+  }
+  virtual StatusTuple remove_value(const KeyType& key) {
+    if (!this->remove(const_cast<KeyType*>(&key)))
+      return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
+    return StatusTuple::OK();
+  }
 };
 
 class BPFSockmapTable : public BPFTableBase<int, int> {

src/cc/export/helpers.h

@@ -384,8 +384,17 @@ struct _name##_table_t _name = { .max_entries = (_max_entries) }
 #define BPF_ARRAY_OF_MAPS(_name, _inner_map_name, _max_entries) \
   BPF_TABLE("array_of_maps$" _inner_map_name, int, int, _name, _max_entries)
 
-#define BPF_HASH_OF_MAPS(_name, _inner_map_name, _max_entries) \
-  BPF_TABLE("hash_of_maps$" _inner_map_name, int, int, _name, _max_entries)
+#define BPF_HASH_OF_MAPS2(_name, _inner_map_name) \
+  BPF_TABLE("hash_of_maps$" _inner_map_name, int, int, _name, 10240)
+#define BPF_HASH_OF_MAPS3(_name, _key_type, _inner_map_name) \
+  BPF_TABLE("hash_of_maps$" _inner_map_name, _key_type, int, _name, 10240)
+#define BPF_HASH_OF_MAPS4(_name, _key_type, _inner_map_name, _max_entries) \
+  BPF_TABLE("hash_of_maps$" _inner_map_name, _key_type, int, _name, _max_entries)
+
+#define BPF_HASH_OF_MAPSX(_name, _2, _3, _4, NAME, ...) NAME
+
+#define BPF_HASH_OF_MAPS(...) \
+  BPF_HASH_OF_MAPSX(__VA_ARGS__, BPF_HASH_OF_MAPS4, BPF_HASH_OF_MAPS3, BPF_HASH_OF_MAPS2)(__VA_ARGS__)
 
 #define BPF_SK_STORAGE(_name, _leaf_type) \
 struct _name##_table_t { \

tests/cc/test_map_in_map.cc

@@ -30,7 +30,7 @@ TEST_CASE("test hash of maps", "[hash_of_maps]") {
       BPF_ARRAY(ex1, int, 1024);
       BPF_ARRAY(ex2, int, 1024);
       BPF_ARRAY(ex3, u64, 1024);
-      BPF_HASH_OF_MAPS(maps_hash, "ex1", 10);
+      BPF_HASH_OF_MAPS(maps_hash, int, "ex1", 10);
 
       int syscall__getuid(void *ctx) {
          int key = 0, data, *val, cntl_val;
@@ -63,7 +63,7 @@ TEST_CASE("test hash of maps", "[hash_of_maps]") {
     res = bpf.init(BPF_PROGRAM);
     REQUIRE(res.code() == 0);
 
-    auto t = bpf.get_map_in_map_table("maps_hash");
+    auto t = bpf.get_map_in_map_table<int>("maps_hash");
     auto ex1_table = bpf.get_array_table<int>("ex1");
     auto ex2_table = bpf.get_array_table<int>("ex2");
     auto ex3_table = bpf.get_array_table<unsigned long long>("ex3");
@@ -115,6 +115,119 @@ TEST_CASE("test hash of maps", "[hash_of_maps]") {
   }
 }
 
+TEST_CASE("test hash of maps using custom key", "[hash_of_maps_custom_key]") {
+  {
+    const std::string BPF_PROGRAM = R"(
+        struct custom_key {
+          int value_1;
+          int value_2;
+        };
+
+        BPF_ARRAY(cntl, int, 1);
+        BPF_TABLE("hash", int, int, ex1, 1024);
+        BPF_TABLE("hash", int, int, ex2, 1024);
+        BPF_HASH_OF_MAPS(maps_hash, struct custom_key, "ex1", 10);
+
+        int syscall__getuid(void *ctx) {
+          struct custom_key hash_key = {1, 0};
+          int key = 0, data, *val, cntl_val;
+          void *inner_map;
+
+          val = cntl.lookup(&key);
+          if (!val || *val == 0)
+            return 0;
+
+          hash_key.value_2 = *val;
+          inner_map = maps_hash.lookup(&hash_key);
+          if (!inner_map)
+            return 0;
+
+          val = bpf_map_lookup_elem(inner_map, &key);
+          if (!val) {
+            data = 1;
+            bpf_map_update_elem(inner_map, &key, &data, 0);
+          } else {
+            data = 1 + *val;
+            bpf_map_update_elem(inner_map, &key, &data, 0);
+          }
+
+          return 0;
+        }
+    )";
+
+    struct custom_key {
+      int value_1;
+      int value_2;
+    };
+
+    ebpf::BPF bpf;
+    ebpf::StatusTuple res(0);
+    res = bpf.init(BPF_PROGRAM);
+    REQUIRE(res.code() == 0);
+
+    auto t = bpf.get_map_in_map_table<struct custom_key>("maps_hash");
+    auto ex1_table = bpf.get_hash_table<int, int>("ex1");
+    auto ex2_table = bpf.get_hash_table<int, int>("ex2");
+    auto cntl_table = bpf.get_array_table<int>("cntl");
+    int ex1_fd = ex1_table.get_fd();
+    int ex2_fd = ex2_table.get_fd();
+
+    // test effectiveness of map-in-map
+    std::string getuid_fnname = bpf.get_syscall_fnname("getuid");
+    res = bpf.attach_kprobe(getuid_fnname, "syscall__getuid");
+    REQUIRE(res.code() == 0);
+
+    struct custom_key hash_key = {1, 1};
+
+    res = t.update_value(hash_key, ex1_fd);
+    REQUIRE(res.code() == 0);
+
+    struct custom_key hash_key2 = {1, 2};
+    res = t.update_value(hash_key2, ex2_fd);
+    REQUIRE(res.code() == 0);
+
+    int key = 0, value = 0, value2 = 0;
+
+    // Can't get value when value didn't set.
+    res = ex1_table.get_value(key, value);
+    REQUIRE(res.code() != 0);
+    REQUIRE(value == 0);
+
+    // Call syscall__getuid, then set value to ex1_table
+    res = cntl_table.update_value(key, 1);
+    REQUIRE(res.code() == 0);
+    REQUIRE(getuid() >= 0);
+
+    // Now we can get value from ex1_table
+    res = ex1_table.get_value(key, value);
+    REQUIRE(res.code() == 0);
+    REQUIRE(value >= 1);
+
+    // Can't get value when value didn't set.
+    res = ex2_table.get_value(key, value2);
+    REQUIRE(res.code() != 0);
+    REQUIRE(value2 == 0);
+
+    // Call syscall__getuid, then set value to ex2_table
+    res = cntl_table.update_value(key, 2);
+    REQUIRE(res.code() == 0);
+    REQUIRE(getuid() >= 0);
+
+    // Now we can get value from ex2_table
+    res = ex2_table.get_value(key, value2);
+    REQUIRE(res.code() == 0);
+    REQUIRE(value > 0);
+
+    res = bpf.detach_kprobe(getuid_fnname);
+    REQUIRE(res.code() == 0);
+
+    res = t.remove_value(hash_key);
+    REQUIRE(res.code() == 0);
+    res = t.remove_value(hash_key2);
+    REQUIRE(res.code() == 0);
+  }
+}
+
 TEST_CASE("test array of maps", "[array_of_maps]") {
   {
     const std::string BPF_PROGRAM = R"(
@@ -158,7 +271,7 @@ TEST_CASE("test array of maps", "[array_of_maps]") {
     res = bpf.init(BPF_PROGRAM);
     REQUIRE(res.code() == 0);
 
-    auto t = bpf.get_map_in_map_table("maps_array");
+    auto t = bpf.get_map_in_map_table<int>("maps_array");
     auto ex1_table = bpf.get_hash_table<int, int>("ex1");
     auto ex2_table = bpf.get_hash_table<int, int>("ex2");
     auto ex3_table =

tests/python/CMakeLists.txt

@@ -95,3 +95,5 @@ add_test(NAME py_queuestack WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_queuestack sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_queuestack.py)
 add_test(NAME py_test_map_batch_ops WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test_map_batch_ops sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_map_batch_ops.py)
+add_test(NAME py_test_map_in_map WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test_map_in_map sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_map_in_map.py)

tests/python/test_map_in_map.py

@@ -12,6 +12,13 @@
 import ctypes as ct
 import os
 
+
+class CustomKey(ct.Structure):
+  _fields_ = [
+    ("value_1", ct.c_int),
+    ("value_2", ct.c_int)
+  ]
+
 def kernel_version_ge(major, minor):
     # True if running kernel is >= X.Y
     version = distutils.version.LooseVersion(os.uname()[2]).version
@@ -30,7 +37,7 @@ def test_hash_table(self):
       BPF_ARRAY(cntl, int, 1);
       BPF_TABLE("hash", int, int, ex1, 1024);
       BPF_TABLE("hash", int, int, ex2, 1024);
-      BPF_HASH_OF_MAPS(maps_hash, "ex1", 10);
+      BPF_HASH_OF_MAPS(maps_hash, int, "ex1", 10);
 
       int syscall__getuid(void *ctx) {
          int key = 0, data, *val, cntl_val;
@@ -77,7 +84,7 @@ def test_hash_table(self):
 
         cntl_map[0] = ct.c_int(1)
         os.getuid()
-        assert(ex1_map[ct.c_int(0)] >= 1)
+        assert(ex1_map[ct.c_int(0)].value >= 1)
 
         try:
           ex2_map[ct.c_int(0)]
@@ -87,12 +94,85 @@ def test_hash_table(self):
 
         cntl_map[0] = ct.c_int(2)
         os.getuid()
-        assert(ex2_map[ct.c_int(0)] >= 1)
+        assert(ex2_map[ct.c_int(0)].value >= 1)
 
         b.detach_kprobe(event=syscall_fnname)
         del hash_maps[ct.c_int(1)]
         del hash_maps[ct.c_int(2)]
 
+    def test_hash_table_custom_key(self):
+        bpf_text = """
+        struct custom_key {
+          int value_1;
+          int value_2;
+        };
+
+        BPF_ARRAY(cntl, int, 1);
+        BPF_TABLE("hash", int, int, ex1, 1024);
+        BPF_TABLE("hash", int, int, ex2, 1024);
+        BPF_HASH_OF_MAPS(maps_hash, struct custom_key, "ex1", 10);
+
+        int syscall__getuid(void *ctx) {
+          struct custom_key hash_key = {1, 0};
+          int key = 0, data, *val, cntl_val;
+          void *inner_map;
+
+          val = cntl.lookup(&key);
+          if (!val || *val == 0)
+            return 0;
+
+          hash_key.value_2 = *val;
+          inner_map = maps_hash.lookup(&hash_key);
+          if (!inner_map)
+            return 0;
+
+          val = bpf_map_lookup_elem(inner_map, &key);
+          if (!val) {
+            data = 1;
+            bpf_map_update_elem(inner_map, &key, &data, 0);
+          } else {
+            data = 1 + *val;
+            bpf_map_update_elem(inner_map, &key, &data, 0);
+          }
+
+          return 0;
+        }
+"""
+        b = BPF(text=bpf_text)
+        cntl_map = b.get_table("cntl")
+        ex1_map = b.get_table("ex1")
+        ex2_map = b.get_table("ex2")
+        hash_maps = b.get_table("maps_hash")
+
+        hash_maps[CustomKey(1, 1)] = ct.c_int(ex1_map.get_fd())
+        hash_maps[CustomKey(1, 2)] = ct.c_int(ex2_map.get_fd())
+        syscall_fnname = b.get_syscall_fnname("getuid")
+        b.attach_kprobe(event=syscall_fnname, fn_name="syscall__getuid")
+
+        try:
+          ex1_map[ct.c_int(0)]
+          raise Exception("Unexpected success for ex1_map[0]")
+        except KeyError:
+          pass
+
+        cntl_map[0] = ct.c_int(1)
+        os.getuid()
+        assert(ex1_map[ct.c_int(0)].value >= 1)
+
+        try:
+          ex2_map[ct.c_int(0)]
+          raise Exception("Unexpected success for ex2_map[0]")
+        except KeyError:
+          pass
+
+        cntl_map[0] = ct.c_int(2)
+        os.getuid()
+        assert(ex2_map[ct.c_int(0)].value >= 1)
+
+        b.detach_kprobe(event=syscall_fnname)
+        del hash_maps[CustomKey(1, 1)]
+        del hash_maps[CustomKey(1, 2)]
+
     def test_array_table(self):
         bpf_text = """
       BPF_ARRAY(cntl, int, 1);
@@ -136,11 +216,11 @@ def test_array_table(self):
 
         cntl_map[0] = ct.c_int(1)
         os.getuid()
-        assert(ex1_map[ct.c_int(0)] >= 1)
+        assert(ex1_map[ct.c_int(0)].value >= 1)
 
         cntl_map[0] = ct.c_int(2)
         os.getuid()
-        assert(ex2_map[ct.c_int(0)] >= 1)
+        assert(ex2_map[ct.c_int(0)].value >= 1)
 
         b.detach_kprobe(event=syscall_fnname)