Update

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 <bpf_map/userspace/var_hash_map.hpp>
 #include <algorithm>
+#include <cerrno>
+#include <cstring>
 #include <functional>
 #include <unistd.h>
 
+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<bytes_vec>(),
 		   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,19 +52,56 @@ 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;
 }
 
 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) {

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

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

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 <boost/interprocess/creation_tags.hpp>
+#include <boost/interprocess/interprocess_fwd.hpp>
+#include <boost/interprocess/managed_shared_memory.hpp>
+#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);
+}