Fix bpf_map_get_next_key "key not found" bug. (#3042)

* fix "prev key not found" bug

* update doc

* Apply comment suggestions from code review

Co-authored-by: Dave Thaler <dthaler1968@gmail.com>

* fix bug

---------

Co-authored-by: Dave Thaler <dthaler1968@gmail.com>

libs/execution_context/ebpf_core.c

@@ -940,6 +940,12 @@ _ebpf_core_protocol_map_get_next_key(
     retval = ebpf_map_next_key(
         map, next_key_length, previous_key_length == 0 ? NULL : request->previous_key, reply->next_key);
 
+    // If the previous key was not found, return the first key.
+    if (retval == EBPF_KEY_NOT_FOUND) {
+        ebpf_assert(previous_key_length != 0); // EBPF_KEY_NOT_FOUND is only returned if previous_key_length != 0.
+        retval = ebpf_map_next_key(map, next_key_length, NULL, reply->next_key);
+    }
+
     reply->header.length = reply_length;
 
 Done:

libs/execution_context/ebpf_maps.h

@@ -134,6 +134,7 @@ extern "C"
      * null indicates that the first key is to be returned.
      * @param[out] next_key Next key on success.
      * @retval EBPF_SUCCESS The operation was successful.
+     * @retval EBPF_KEY_NOT_FOUND The specified previous key was not found.
      * @retval EBPF_NO_MORE_KEYS There is no key following the specified
      * key in lexicographical order.
      */

libs/runtime/ebpf_hash_table.c

@@ -813,10 +813,20 @@ ebpf_hash_table_next_key_pointer_and_value(
                 found_entry = true;
             }
         }
+
         if (next_entry) {
             break;
         }
     }
+
+    // If we were given a previous key, and the searched key was not found in the hash table, we return
+    // EBPF_KEY_NOT_FOUND, so that the caller can detect that the key is missing, and return the first key (as per
+    // 'bpf_map_get_next_key' specs).
+    if (!found_entry && previous_key != NULL) {
+        result = EBPF_KEY_NOT_FOUND;
+        goto Done;
+    }
+
     if (!next_entry) {
         result = EBPF_NO_MORE_KEYS;
         goto Done;

tests/unit/libbpf_test.cpp

@@ -2694,7 +2694,7 @@ TEST_CASE("BPF_MAP_GET_NEXT_KEY etc.", "[libbpf]")
     attr.map_type = BPF_MAP_TYPE_HASH;
     attr.key_size = sizeof(uint32_t);
     attr.value_size = sizeof(uint32_t);
-    attr.max_entries = 2;
+    attr.max_entries = 3;
     attr.map_flags = 0;
     int map_fd = bpf(BPF_MAP_CREATE, &attr, sizeof(attr));
     REQUIRE(map_fd > 0);
@@ -2765,6 +2765,34 @@ TEST_CASE("BPF_MAP_GET_NEXT_KEY etc.", "[libbpf]")
     REQUIRE(bpf(BPF_MAP_GET_NEXT_KEY, &attr, sizeof(attr)) < 0);
     REQUIRE(errno == ENOENT);
 
+    // Test that the bpf_map_get_next_key returns the first key of the map if the previous key is not found.
+    // Add 3 entries into the now empty map.
+    for (key = 100; key < 400; key += 100) {
+        value = 0;
+        memset(&attr, 0, sizeof(attr));
+        attr.map_fd = map_fd;
+        attr.key = (uintptr_t)&key;
+        attr.value = (uintptr_t)&value;
+        REQUIRE(bpf(BPF_MAP_UPDATE_ELEM, &attr, sizeof(attr)) == 0);
+    }
+
+    // Look up the first key in the map, so we can check that it's returned later.
+    memset(&attr, 0, sizeof(attr));
+    attr.map_fd = map_fd;
+    attr.key = NULL;
+    attr.next_key = (uintptr_t)&next_key;
+    REQUIRE(bpf(BPF_MAP_GET_NEXT_KEY, &attr, sizeof(attr)) == 0);
+    uint64_t first_key = next_key;
+
+    // Look up a key that is not present in the map, and check that the first key is returned.
+    key = 123;
+    memset(&attr, 0, sizeof(attr));
+    attr.map_fd = map_fd;
+    attr.key = (uintptr_t)&key;
+    attr.next_key = (uintptr_t)&next_key;
+    REQUIRE(bpf(BPF_MAP_GET_NEXT_KEY, &attr, sizeof(attr)) == 0);
+    REQUIRE(next_key == first_key);
+
     Platform::_close(map_fd);
 }