fix next key not looping back to the beginning

libs/execution_context/ebpf_core.c

@@ -945,6 +945,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 pervious key was not found, return the first key (if it wasn't already looked up).
+    if (retval == EBPF_KEY_NOT_FOUND) {
+        ebpf_assert(previous_key_length != 0);
+        retval = ebpf_map_next_key(map, next_key_length, NULL, reply->next_key);
+    }
+
     reply->header.length = reply_length;
 
 Done:

libs/runtime/ebpf_hash_table.c

@@ -813,6 +813,14 @@ ebpf_hash_table_next_key_pointer_and_value(
                 found_entry = true;
             }
         }
+
+        // If there was a previous key that was not present in the bucket, then it is not present in the hash table
+        // either.
+        if (!found_entry && previous_key != NULL) {
+            result = EBPF_KEY_NOT_FOUND;
+            goto Done;
+        }
+
         if (next_entry) {
             break;
         }

tests/unit/libbpf_test.cpp

@@ -2672,7 +2672,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);
@@ -2743,6 +2743,31 @@ 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 BPF map if the previous key is not found.
+    // Add 3 entries into the now empty map.
+    for (key = 100; key < 400; key += 100) {
+        memset(&attr, 0, sizeof(attr));
+        attr.map_fd = map_fd;
+        attr.key = (uintptr_t)&key;
+        REQUIRE(bpf(BPF_MAP_UPDATE_ELEM, &attr, sizeof(attr)) == 0);
+    }
+    // Lookup the first key in the map, so to check it's returned later.
+    value = 0;
+    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;
+    // Lookup 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);
 }