utils: add map iteration iterator (#4377)

tests/unit/s2n_map_iterator_test.c

@@ -0,0 +1,134 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+#include "api/s2n.h"
+#include "s2n_test.h"
+#include "utils/s2n_map.h"
+#include "utils/s2n_map_internal.h"
+
+#define TEST_VALUE_COUNT 10
+
+int main(int argc, char **argv)
+{
+    BEGIN_TEST();
+
+    /* s2n_map_iterator iteration test */
+    {
+        struct s2n_map *map = s2n_map_new();
+        EXPECT_NOT_NULL(map);
+        /* fail to initialize an iterator on a mutable map */
+        {
+            struct s2n_map_iterator iter = { 0 };
+            EXPECT_ERROR_WITH_ERRNO(s2n_map_iterator_init(&iter, map), S2N_ERR_MAP_MUTABLE);
+        };
+
+        EXPECT_OK(s2n_map_complete(map));
+
+        /* has next is false on an empty map, and next returns an error */
+        {
+            struct s2n_map_iterator iter = { 0 };
+            EXPECT_OK(s2n_map_iterator_init(&iter, map));
+
+            EXPECT_FALSE(s2n_map_iterator_has_next(&iter));
+
+            struct s2n_blob value = { 0 };
+            EXPECT_ERROR_WITH_ERRNO(s2n_map_iterator_next(&iter, &value), S2N_ERR_ARRAY_INDEX_OOB);
+        };
+
+        EXPECT_OK(s2n_map_unlock(map));
+        for (uint8_t i = 0; i < TEST_VALUE_COUNT; i++) {
+            struct s2n_blob key = { .size = 1, .data = &i };
+            struct s2n_blob val = { .size = 1, .data = &i };
+            EXPECT_OK(s2n_map_put(map, &key, &val));
+        }
+        EXPECT_OK(s2n_map_complete(map));
+
+        /* iterator goes over all elements */
+        {
+            bool seen[TEST_VALUE_COUNT] = { 0 };
+
+            struct s2n_map_iterator iter = { 0 };
+            EXPECT_OK(s2n_map_iterator_init(&iter, map));
+
+            struct s2n_blob value = { 0 };
+            for (size_t i = 0; i < TEST_VALUE_COUNT; i++) {
+                EXPECT_TRUE(s2n_map_iterator_has_next(&iter));
+
+                EXPECT_OK(s2n_map_iterator_next(&iter, &value));
+                seen[*value.data] = true;
+            }
+
+            /* all elements have been iterated over */
+            EXPECT_FALSE(s2n_map_iterator_has_next(&iter));
+            EXPECT_ERROR_WITH_ERRNO(s2n_map_iterator_next(&iter, &value), S2N_ERR_ARRAY_INDEX_OOB);
+
+            /* all elements were seen */
+            for (size_t i = 0; i < TEST_VALUE_COUNT; i++) {
+                EXPECT_TRUE(seen[i]);
+            }
+        };
+
+        /* next returns an error when the blob is null */
+        {
+            struct s2n_map_iterator iter = { 0 };
+            EXPECT_OK(s2n_map_iterator_init(&iter, map));
+
+            EXPECT_ERROR_WITH_ERRNO(s2n_map_iterator_next(&iter, NULL), S2N_ERR_NULL);
+        }
+
+        EXPECT_OK(s2n_map_free(map));
+    };
+
+    /* test first and last slots in table */
+    {
+        /* 2 (first and last slot) * 2 (key and value) */
+        struct s2n_blob blobs[2 * 2] = { 0 };
+        for (uint8_t i = 0; i < (2 * 2); i++) {
+            EXPECT_SUCCESS(s2n_alloc(&blobs[i], 1));
+            *blobs[i].data = i;
+        }
+
+        struct s2n_map *test_map = s2n_map_new();
+        EXPECT_NOT_NULL(test_map);
+
+        /* set values in map to 0 and 1 */
+        test_map->table[0].value = blobs[0];
+        test_map->table[0].key = blobs[2];
+        test_map->table[test_map->capacity - 1].value = blobs[1];
+        test_map->table[test_map->capacity - 1].key = blobs[3];
+
+        test_map->size = 2;
+        EXPECT_OK(s2n_map_complete(test_map));
+
+        struct s2n_map_iterator iter = { 0 };
+        EXPECT_OK(s2n_map_iterator_init(&iter, test_map));
+        bool seen[2] = { 0 };
+
+        struct s2n_blob value = { 0 };
+        for (size_t i = 0; i < 2; i++) {
+            EXPECT_TRUE(s2n_map_iterator_has_next(&iter));
+
+            EXPECT_OK(s2n_map_iterator_next(&iter, &value));
+            seen[*value.data] = true;
+        }
+
+        /* assert that 0 and 1 were both seen */
+        EXPECT_TRUE(seen[0] && seen[1]);
+
+        EXPECT_OK(s2n_map_free(test_map));
+    };
+
+    END_TEST();
+}

tests/unit/s2n_map_test.c

@@ -20,10 +20,13 @@
 #include "api/s2n.h"
 #include "s2n_test.h"
 
+#define TEST_VALUE_COUNT 8192
+
 int main(int argc, char **argv)
 {
     char keystr[sizeof("ffff")];
     char valstr[sizeof("16384")];
+    uint32_t size = 0;
     struct s2n_map *empty, *map;
     struct s2n_blob key = { 0 };
     struct s2n_blob val = { 0 };
@@ -33,6 +36,8 @@ int main(int argc, char **argv)
     EXPECT_SUCCESS(s2n_disable_tls13_in_test());
 
     EXPECT_NOT_NULL(empty = s2n_map_new());
+    EXPECT_OK(s2n_map_size(empty, &size));
+    EXPECT_EQUAL(size, 0);
 
     /* Try a lookup on an empty map. Expect an error because the map is still mutable. */
     EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", 1234));
@@ -67,7 +72,7 @@ int main(int argc, char **argv)
     EXPECT_NOT_NULL(map = s2n_map_new_with_initial_capacity(1));
 
     /* Insert 8k key value pairs of the form hex(i) -> dec(i) */
-    for (int i = 0; i < 8192; i++) {
+    for (int i = 0; i < TEST_VALUE_COUNT; i++) {
         EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", i));
         EXPECT_SUCCESS(snprintf(valstr, sizeof(valstr), "%05d", i));
 
@@ -78,6 +83,8 @@ int main(int argc, char **argv)
 
         EXPECT_OK(s2n_map_add(map, &key, &val));
     }
+    EXPECT_OK(s2n_map_size(map, &size));
+    EXPECT_EQUAL(size, TEST_VALUE_COUNT);
 
     /* Try adding some duplicates */
     for (int i = 0; i < 10; i++) {
@@ -91,6 +98,8 @@ int main(int argc, char **argv)
 
         EXPECT_ERROR(s2n_map_add(map, &key, &val));
     }
+    EXPECT_OK(s2n_map_size(map, &size));
+    EXPECT_EQUAL(size, TEST_VALUE_COUNT);
 
     /* Try replacing some entries */
     for (int i = 0; i < 10; i++) {
@@ -113,8 +122,8 @@ int main(int argc, char **argv)
     EXPECT_OK(s2n_map_complete(map));
 
     /* Make sure that add-after-complete fails */
-    EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", 8193));
-    EXPECT_SUCCESS(snprintf(valstr, sizeof(valstr), "%05d", 8193));
+    EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", TEST_VALUE_COUNT + 1));
+    EXPECT_SUCCESS(snprintf(valstr, sizeof(valstr), "%05d", TEST_VALUE_COUNT + 1));
 
     key.data = (void *) keystr;
     key.size = strlen(keystr) + 1;
@@ -124,7 +133,7 @@ int main(int argc, char **argv)
     EXPECT_ERROR(s2n_map_add(map, &key, &val));
 
     /* Check for equivalence */
-    for (int i = 0; i < 8192; i++) {
+    for (int i = 0; i < TEST_VALUE_COUNT; i++) {
         if (i >= 10) {
             EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", i));
             EXPECT_SUCCESS(snprintf(valstr, sizeof(valstr), "%05d", i));
@@ -144,7 +153,7 @@ int main(int argc, char **argv)
     }
 
     /* Check for a key that shouldn't be there */
-    EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", 8193));
+    EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", TEST_VALUE_COUNT + 1));
     key.data = (void *) keystr;
     key.size = strlen(keystr) + 1;
     EXPECT_OK(s2n_map_lookup(map, &key, &val, &key_found));
@@ -153,8 +162,8 @@ int main(int argc, char **argv)
     /* Make the map mutable */
     EXPECT_OK(s2n_map_unlock(map));
     /* Make sure that add-after-unlock succeeds */
-    EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", 8193));
-    EXPECT_SUCCESS(snprintf(valstr, sizeof(valstr), "%05d", 8193));
+    EXPECT_SUCCESS(snprintf(keystr, sizeof(keystr), "%04x", TEST_VALUE_COUNT + 1));
+    EXPECT_SUCCESS(snprintf(valstr, sizeof(valstr), "%05d", TEST_VALUE_COUNT + 1));
 
     key.data = (void *) keystr;
     key.size = strlen(keystr) + 1;

utils/s2n_map.c

@@ -209,8 +209,8 @@ S2N_RESULT s2n_map_lookup(const struct s2n_map *map, struct s2n_blob *key, struc
         }
 
         /* We found a match */
-        value->data = map->table[slot].value.data;
-        value->size = map->table[slot].value.size;
+        struct s2n_blob entry_value = map->table[slot].value;
+        RESULT_GUARD_POSIX(s2n_blob_init(value, entry_value.data, entry_value.size));
 
         *key_found = true;
 
@@ -246,3 +246,73 @@ S2N_RESULT s2n_map_free(struct s2n_map *map)
 
     return S2N_RESULT_OK;
 }
+
+S2N_RESULT s2n_map_size(struct s2n_map *map, uint32_t *size)
+{
+    RESULT_ENSURE_REF(map);
+    *size = map->size;
+    return S2N_RESULT_OK;
+}
+
+/* Update the internal state so that `current_index` will point to the next value
+ * in the table or set iter->consumed equal to true if there are no more elements
+ * in the map. 
+ */
+S2N_RESULT s2n_map_iterator_advance(struct s2n_map_iterator *iter)
+{
+    RESULT_ENSURE_REF(iter);
+    RESULT_ENSURE_REF(iter->map);
+    RESULT_ENSURE(s2n_map_iterator_has_next(iter), S2N_ERR_ARRAY_INDEX_OOB);
+
+    iter->current_index++;
+    while (iter->current_index < iter->map->capacity) {
+        /* a value was found in the map */
+        if (iter->map->table[iter->current_index].key.size) {
+            return S2N_RESULT_OK;
+        }
+        iter->current_index++;
+    }
+    /* no more values were found in the map */
+    iter->consumed = true;
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_map_iterator_init(struct s2n_map_iterator *iter, const struct s2n_map *map)
+{
+    RESULT_ENSURE_REF(iter);
+    RESULT_ENSURE_REF(map);
+    RESULT_ENSURE(map->immutable, S2N_ERR_MAP_MUTABLE);
+
+    iter->map = map;
+    iter->current_index = 0;
+
+    /* Advance the index to point to the first value in the map. This isn't 
+     * necessary if the slot at index 0 already contains a value.
+     */
+    if (iter->map->table[0].key.size == 0) {
+        RESULT_GUARD(s2n_map_iterator_advance(iter));
+    }
+
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_map_iterator_next(struct s2n_map_iterator *iter, struct s2n_blob *value)
+{
+    RESULT_ENSURE_REF(iter);
+    RESULT_ENSURE_REF(iter->map);
+    RESULT_ENSURE(iter->map->immutable, S2N_ERR_MAP_MUTABLE);
+    RESULT_ENSURE(s2n_map_iterator_has_next(iter), S2N_ERR_ARRAY_INDEX_OOB);
+
+    RESULT_ENSURE(iter->current_index < iter->map->capacity, S2N_ERR_ARRAY_INDEX_OOB);
+    struct s2n_blob entry_value = iter->map->table[iter->current_index].value;
+    RESULT_GUARD_POSIX(s2n_blob_init(value, entry_value.data, entry_value.size));
+
+    RESULT_GUARD(s2n_map_iterator_advance(iter));
+
+    return S2N_RESULT_OK;
+}
+
+bool s2n_map_iterator_has_next(const struct s2n_map_iterator *iter)
+{
+    return iter && !iter->consumed;
+}

utils/s2n_map.h

@@ -21,6 +21,12 @@
 #include "utils/s2n_result.h"
 
 struct s2n_map;
+struct s2n_map_iterator {
+    const struct s2n_map *map;
+    /* Index of the entry to be returned on the next `s2n_map_iterator_next()` call. */
+    uint32_t current_index;
+    bool consumed;
+};
 
 struct s2n_map *s2n_map_new();
 struct s2n_map *s2n_map_new_with_initial_capacity(uint32_t capacity);
@@ -30,3 +36,8 @@ S2N_RESULT s2n_map_complete(struct s2n_map *map);
 S2N_RESULT s2n_map_unlock(struct s2n_map *map);
 S2N_RESULT s2n_map_lookup(const struct s2n_map *map, struct s2n_blob *key, struct s2n_blob *value, bool *key_found);
 S2N_RESULT s2n_map_free(struct s2n_map *map);
+S2N_RESULT s2n_map_size(struct s2n_map *map, uint32_t *size);
+
+S2N_RESULT s2n_map_iterator_init(struct s2n_map_iterator *iter, const struct s2n_map *map);
+S2N_RESULT s2n_map_iterator_next(struct s2n_map_iterator *iter, struct s2n_blob *value);
+bool s2n_map_iterator_has_next(const struct s2n_map_iterator *iter);