From 722cf83941879c52ebea5e5a1692b2976de6ad62 Mon Sep 17 00:00:00 2001 From: masibw Date: Tue, 22 Jun 2021 15:18:23 +0900 Subject: [PATCH] Allow the use of custom keys in BPF_HASH_OF_MAPS (#3500) - Allow the use of custom keys in BPF_HASH_OF_MAPS - Add both python and C++ tests --- src/cc/api/BPF.cc | 7 -- src/cc/api/BPF.h | 9 ++- src/cc/api/BPFTable.cc | 21 ------ src/cc/api/BPFTable.h | 26 +++++-- src/cc/export/helpers.h | 13 +++- tests/cc/test_map_in_map.cc | 119 +++++++++++++++++++++++++++++++- tests/python/CMakeLists.txt | 2 + tests/python/test_map_in_map.py | 90 ++++++++++++++++++++++-- 8 files changed, 241 insertions(+), 46 deletions(-) diff --git a/src/cc/api/BPF.cc b/src/cc/api/BPF.cc index 04ef3bf7..87d4a331 100644 --- a/src/cc/api/BPF.cc +++ b/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)) diff --git a/src/cc/api/BPF.h b/src/cc/api/BPF.h index d6f3b2a9..c266828e 100644 --- a/src/cc/api/BPF.h +++ b/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 + BPFMapInMapTable 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({}); + } bool add_module(std::string module); diff --git a/src/cc/api/BPFTable.cc b/src/cc/api/BPFTable.cc index d1050512..f27e7125 100644 --- a/src/cc/api/BPFTable.cc +++ b/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(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(&index), const_cast(&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(&index))) - return StatusTuple(-1, "Error removing value: %s", std::strerror(errno)); - return StatusTuple::OK(); -} - BPFSockmapTable::BPFSockmapTable(const TableDesc& desc) : BPFTableBase(desc) { if(desc.type != BPF_MAP_TYPE_SOCKMAP) diff --git a/src/cc/api/BPFTable.h b/src/cc/api/BPFTable.h index d63f3c5d..8786a3f9 100644 --- a/src/cc/api/BPFTable.h +++ b/src/cc/api/BPFTable.h @@ -479,12 +479,26 @@ public: StatusTuple remove_value(const int& index); }; -class BPFMapInMapTable : public BPFTableBase { -public: - BPFMapInMapTable(const TableDesc& desc); - - StatusTuple update_value(const int& index, const int& inner_map_fd); - StatusTuple remove_value(const int& index); +template +class BPFMapInMapTable : public BPFTableBase { + public: + BPFMapInMapTable(const TableDesc& desc) : BPFTableBase(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(&key), + const_cast(&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(&key))) + return StatusTuple(-1, "Error removing value: %s", std::strerror(errno)); + return StatusTuple::OK(); + } }; class BPFSockmapTable : public BPFTableBase { diff --git a/src/cc/export/helpers.h b/src/cc/export/helpers.h index 12072b06..a3283d2e 100644 --- a/src/cc/export/helpers.h +++ b/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 { \ diff --git a/tests/cc/test_map_in_map.cc b/tests/cc/test_map_in_map.cc index f8c1a0b6..7a383de9 100644 --- a/tests/cc/test_map_in_map.cc +++ b/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("maps_hash"); auto ex1_table = bpf.get_array_table("ex1"); auto ex2_table = bpf.get_array_table("ex2"); auto ex3_table = bpf.get_array_table("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("maps_hash"); + auto ex1_table = bpf.get_hash_table("ex1"); + auto ex2_table = bpf.get_hash_table("ex2"); + auto cntl_table = bpf.get_array_table("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("maps_array"); auto ex1_table = bpf.get_hash_table("ex1"); auto ex2_table = bpf.get_hash_table("ex2"); auto ex3_table = diff --git a/tests/python/CMakeLists.txt b/tests/python/CMakeLists.txt index 7e60413b..3cd96031 100644 --- a/tests/python/CMakeLists.txt +++ b/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) diff --git a/tests/python/test_map_in_map.py b/tests/python/test_map_in_map.py index 0f596043..bd909d84 100755 --- a/tests/python/test_map_in_map.py +++ b/tests/python/test_map_in_map.py @@ -12,6 +12,13 @@ from unittest import main, skipUnless, TestCase 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 @@ class TestUDST(TestCase): 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 @@ class TestUDST(TestCase): 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 @@ class TestUDST(TestCase): 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 @@ class TestUDST(TestCase): 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) -- 2.34.1