diff --git a/runtime/src/bpf_map/userspace/var_hash_map.cpp b/runtime/src/bpf_map/userspace/var_hash_map.cpp index abfdcd61..d8acabdb 100644 --- a/runtime/src/bpf_map/userspace/var_hash_map.cpp +++ b/runtime/src/bpf_map/userspace/var_hash_map.cpp @@ -3,21 +3,31 @@ * Copyright (c) 2022, eunomia-bpf org * All rights reserved. */ +#include "linux/bpf.h" #include "spdlog/spdlog.h" #include #include +#include +#include #include #include +static bool is_good_update_flag(uint64_t flags) +{ + return flags == BPF_ANY || flags == BPF_EXIST || flags == BPF_NOEXIST; +} + namespace bpftime { var_size_hash_map_impl::var_size_hash_map_impl(managed_shared_memory &memory, uint32_t key_size, - uint32_t value_size) + uint32_t value_size, + uint32_t max_entries) : map_impl(10, bytes_vec_hasher(), std::equal_to(), bi_map_allocator(memory.get_segment_manager())), _key_size(key_size), _value_size(value_size), + _max_entries(max_entries), key_vec(key_size, memory.get_segment_manager()), value_vec(value_size, memory.get_segment_manager()) { @@ -42,8 +52,29 @@ void *var_size_hash_map_impl::elem_lookup(const void *key) long var_size_hash_map_impl::elem_update(const void *key, const void *value, uint64_t flags) { + // Check flags + if (!is_good_update_flag(flags)) { + errno = EINVAL; + return -1; + } key_vec.assign((uint8_t *)key, (uint8_t *)key + _key_size); value_vec.assign((uint8_t *)value, (uint8_t *)value + _value_size); + bool element_exists = map_impl.contains(key_vec); + // Check if the element exists... + if (flags == BPF_NOEXIST && element_exists) { + errno = EEXIST; + return -1; + } + if (flags == BPF_EXIST && !element_exists) { + errno = ENOENT; + return -1; + } + // Ensure max_entries + if (element_exists == false && map_impl.size() == _max_entries) { + errno = E2BIG; + return -1; + } + // map_impl. map_impl.insert_or_assign(key_vec, value_vec); return 0; } @@ -51,10 +82,26 @@ long var_size_hash_map_impl::elem_update(const void *key, const void *value, long var_size_hash_map_impl::elem_delete(const void *key) { key_vec.assign((uint8_t *)key, (uint8_t *)key + _key_size); - map_impl.erase(key_vec); + auto itr = map_impl.find(key_vec); + if (itr == map_impl.end()) { + errno = ENOENT; + return -1; + } + map_impl.erase(itr); + return 0; +} +int var_size_hash_map_impl::lookup_and_delete(const void *key, void *value_out) +{ + key_vec.assign((uint8_t *)key, (uint8_t *)key + _key_size); + auto itr = map_impl.find(key_vec); + if (itr == map_impl.end()) { + errno = ENOENT; + return -1; + } + memcpy(value_out, &itr->second[0], _value_size); + map_impl.erase(itr); return 0; } - int var_size_hash_map_impl::map_get_next_key(const void *key, void *next_key) { if (key == nullptr) { diff --git a/runtime/src/bpf_map/userspace/var_hash_map.hpp b/runtime/src/bpf_map/userspace/var_hash_map.hpp index 456cf266..1e927924 100644 --- a/runtime/src/bpf_map/userspace/var_hash_map.hpp +++ b/runtime/src/bpf_map/userspace/var_hash_map.hpp @@ -30,6 +30,7 @@ class var_size_hash_map_impl { shm_hash_map map_impl; uint32_t _key_size; uint32_t _value_size; + uint32_t _max_entries; // buffers used to access the key and value in hash map bytes_vec key_vec; @@ -38,7 +39,7 @@ class var_size_hash_map_impl { public: const static bool should_lock = true; var_size_hash_map_impl(managed_shared_memory &memory, uint32_t key_size, - uint32_t value_size); + uint32_t value_size, uint32_t max_entries); void *elem_lookup(const void *key); @@ -47,6 +48,8 @@ class var_size_hash_map_impl { long elem_delete(const void *key); int map_get_next_key(const void *key, void *next_key); + + int lookup_and_delete(const void *key, void *value_out); }; } // namespace bpftime diff --git a/runtime/unit-test/CMakeLists.txt b/runtime/unit-test/CMakeLists.txt index 95560844..d48ce7c6 100644 --- a/runtime/unit-test/CMakeLists.txt +++ b/runtime/unit-test/CMakeLists.txt @@ -19,7 +19,8 @@ set(TEST_SOURCES maps/test_shm_hash_maps.cpp maps/test_external_map_ops.cpp maps/test_bpftime_hash_map.cpp - + maps/kernel_unit_tests.cpp + test_bpftime_shm_json.cpp attach_with_ebpf/test_attach_filter_with_ebpf.cpp attach_with_ebpf/test_attach_uprobe_with_ebpf.cpp diff --git a/runtime/unit-test/maps/kernel_unit_tests.cpp b/runtime/unit-test/maps/kernel_unit_tests.cpp new file mode 100644 index 00000000..c181cb80 --- /dev/null +++ b/runtime/unit-test/maps/kernel_unit_tests.cpp @@ -0,0 +1,108 @@ +#include "catch2/catch_test_macros.hpp" +#include "linux/bpf.h" +#include "unit-test/common_def.hpp" +#include +#include +#include +#include "bpf_map/userspace/var_hash_map.hpp" +static const char *SHM_NAME = "_HASH_MAP_TEST"; + +using namespace boost::interprocess; +using namespace bpftime; + +TEST_CASE("Test hash map (kernel)") +{ + shm_remove remover(SHM_NAME); + managed_shared_memory mem(boost::interprocess::create_only, SHM_NAME, + 20 << 20); + + long long key, next_key, first_key, value; + var_size_hash_map_impl map(mem, sizeof(key), sizeof(value), 2); + key = 1; + value = 1234; + /* Insert key=1 element. */ + REQUIRE(map.elem_update(&key, &value, BPF_ANY) == 0); + + value = 0; + /* BPF_NOEXIST means add new element if it doesn't exist. */ + REQUIRE(map.elem_update(&key, &value, BPF_NOEXIST) < 0); + REQUIRE(/* key=1 already exists. */ errno == EEXIST); + + /* -1 is an invalid flag. */ + REQUIRE(map.elem_update(&key, &value, -1) < 0); + REQUIRE(errno == EINVAL); + + /* Check that key=1 can be found. */ + { + auto ptr = map.elem_lookup(&key); + REQUIRE(ptr != nullptr); + REQUIRE(*(long long *)ptr == 1234); + } + + key = 2; + value = 1234; + /* Insert key=2 element. */ + REQUIRE(map.elem_update(&key, &value, BPF_ANY) == 0); + + /* Check that key=2 matches the value and delete it */ + REQUIRE(map.lookup_and_delete(&key, &value) == 0); + REQUIRE(value == 1234); + + /* Check that key=2 is not found. */ + REQUIRE(map.elem_lookup(&key) == nullptr); + REQUIRE(errno == ENOENT); + + /* BPF_EXIST means update existing element. */ + REQUIRE(map.elem_update(&key, &value, BPF_EXIST) < 0); + + REQUIRE(/* key=2 is not there. */ + errno == ENOENT); + + /* Insert key=2 element. */ + REQUIRE(map.elem_update(&key, &value, BPF_NOEXIST) == 0); + + /* key=1 and key=2 were inserted, check that key=0 cannot be + * inserted due to max_entries limit. + */ + key = 0; + REQUIRE(map.elem_update(&key, &value, BPF_NOEXIST) < 0); + REQUIRE(errno == E2BIG); + + /* Update existing element, though the map is full. */ + key = 1; + REQUIRE(map.elem_update(&key, &value, BPF_EXIST) == 0); + key = 2; + REQUIRE(map.elem_update(&key, &value, BPF_ANY) == 0); + key = 3; + REQUIRE(map.elem_update(&key, &value, BPF_NOEXIST) < 0); + REQUIRE(errno == E2BIG); + + /* Check that key = 0 doesn't exist. */ + key = 0; + REQUIRE(map.elem_delete(&key) < 0); + REQUIRE(errno == ENOENT); + /* Iterate over two elements. */ + REQUIRE(map.map_get_next_key(NULL, &first_key) == 0); + REQUIRE((first_key == 1 || first_key == 2)); + REQUIRE(map.map_get_next_key(&key, &next_key) == 0); + REQUIRE((next_key == first_key)); + REQUIRE(map.map_get_next_key(&next_key, &next_key) == 0); + REQUIRE(((next_key == 1 || next_key == 2) && (next_key != first_key))); + REQUIRE(map.map_get_next_key(&next_key, &next_key) < 0); + REQUIRE(errno == ENOENT); + + /* Delete both elements. */ + key = 1; + REQUIRE(map.elem_delete(&key) == 0); + key = 2; + REQUIRE(map.elem_delete(&key) == 0); + REQUIRE(map.elem_delete(&key) < 0); + REQUIRE(errno == ENOENT); + + key = 0; + /* Check that map is empty. */ + REQUIRE(map.map_get_next_key(NULL, &next_key) < 0); + REQUIRE(errno == ENOENT); + REQUIRE(map.map_get_next_key(&key, &next_key) < 0); + REQUIRE(errno == ENOENT); +}