utils: add map iteration iterator

tests/unit/s2n_map_test.c

@@ -19,11 +19,15 @@
 
 #include "api/s2n.h"
 #include "s2n_test.h"
+#include "utils/s2n_map_internal.h"
+
+DEFINE_POINTER_CLEANUP_FUNC(struct s2n_map_iterator *, s2n_map_iterator_free);
 
 int main(int argc, char **argv)
 {
     char keystr[sizeof("ffff")];
     char valstr[sizeof("16384")];
+    uint32_t size;
     struct s2n_map *empty, *map;
     struct s2n_blob key = { 0 };
     struct s2n_blob val = { 0 };
@@ -33,6 +37,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));
@@ -78,6 +84,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, 8192);
 
     /* Try adding some duplicates */
     for (int i = 0; i < 10; i++) {
@@ -91,6 +99,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, 8192);
 
     /* Try replacing some entries */
     for (int i = 0; i < 10; i++) {
@@ -173,5 +183,110 @@ int main(int argc, char **argv)
 
     EXPECT_OK(s2n_map_free(map));
 
+    /* s2n_map_iterator test */
+    {
+        struct s2n_map *map = s2n_map_new();
+        EXPECT_NOT_NULL(map);
+        /* fail to initialize an iterator on a mutable map */
+        {
+            DEFER_CLEANUP(struct s2n_map_iterator *iter = NULL, s2n_map_iterator_free_pointer);
+            EXPECT_ERROR(s2n_map_iterator_new(map, &iter));
+        };
+
+        EXPECT_OK(s2n_map_complete(map));
+
+        /* has next is false on an empty map, and next returns an error */
+        {
+            DEFER_CLEANUP(struct s2n_map_iterator *iter = NULL, s2n_map_iterator_free_pointer);
+            EXPECT_OK(s2n_map_iterator_new(map, &iter));
+
+            bool has_next = false;
+            EXPECT_OK(s2n_map_iterator_has_next(iter, &has_next));
+            EXPECT_FALSE(has_next);
+
+            struct s2n_blob value = { 0 };
+            EXPECT_ERROR(s2n_map_iterator_next(iter, &value));
+        };
+
+        EXPECT_OK(s2n_map_unlock(map));
+        for (uint8_t i = 0; i < 10; 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[10] = { 0 };
+
+            DEFER_CLEANUP(struct s2n_map_iterator *iter = NULL, s2n_map_iterator_free_pointer);
+            EXPECT_OK(s2n_map_iterator_new(map, &iter));
+
+            bool has_next = false;
+            struct s2n_blob value = { 0 };
+            for (size_t i = 0; i < 10; i++) {
+                EXPECT_OK(s2n_map_iterator_has_next(iter, &has_next));
+                EXPECT_TRUE(has_next);
+
+                EXPECT_OK(s2n_map_iterator_next(iter, &value));
+                seen[*value.data] = true;
+            }
+
+            /* all elements have been iterated over */
+            EXPECT_OK(s2n_map_iterator_has_next(iter, &has_next));
+            EXPECT_FALSE(has_next);
+
+            EXPECT_ERROR(s2n_map_iterator_next(iter, &value));
+
+            /* all elements were seen */
+            for (size_t i = 0; i < 10; i++) {
+                EXPECT_TRUE(seen[i]);
+            }
+        }
+
+        EXPECT_OK(s2n_map_free(map));
+
+        /* test first and last slots in table */
+        {
+            struct s2n_blob blobs[4] = { 0 };
+            for (uint8_t i = 0; i < 4; i++) {
+                s2n_alloc(&blobs[i], 1);
+                *blobs[i].data = i;
+            }
+
+            struct s2n_map *map = s2n_map_new();
+            EXPECT_NOT_NULL(map);
+
+            /* set values in map to 0 and 1 */
+            map->table[0].value = blobs[0];
+            map->table[0].key = blobs[2];
+            map->table[map->capacity - 1].value = blobs[1];
+            map->table[map->capacity - 1].key = blobs[3];
+
+            map->size = 2;
+            EXPECT_OK(s2n_map_complete(map));
+
+            DEFER_CLEANUP(struct s2n_map_iterator *iter = NULL, s2n_map_iterator_free_pointer);
+            EXPECT_OK(s2n_map_iterator_new(map, &iter));
+            bool seen[2] = { 0 };
+
+            bool has_next = false;
+            struct s2n_blob value = { 0 };
+            for (size_t i = 0; i < 2; i++) {
+                EXPECT_OK(s2n_map_iterator_has_next(iter, &has_next));
+                EXPECT_TRUE(has_next);
+
+                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(map));
+        };
+    };
+
     END_TEST();
 }

utils/s2n_map.c

@@ -246,3 +246,91 @@ 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 will equal `map->capacity` 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);
+
+    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;
+        }
+    }
+    /* no more values were found in the map */
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_map_iterator_new(struct s2n_map *map, struct s2n_map_iterator **iter_out)
+{
+    RESULT_ENSURE_REF(map);
+    RESULT_ENSURE(map->immutable, S2N_ERR_MAP_MUTABLE);
+
+    struct s2n_blob mem = { 0 };
+    struct s2n_map_iterator *iter;
+
+    RESULT_GUARD_POSIX(s2n_alloc(&mem, sizeof(struct s2n_map_iterator)));
+
+    iter = (void *) mem.data;
+    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));
+    }
+
+    *iter_out = 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(iter->map->immutable, S2N_ERR_MAP_MUTABLE);
+
+    bool has_next = false;
+    RESULT_GUARD(s2n_map_iterator_has_next(iter, &has_next));
+    RESULT_ENSURE(has_next, S2N_ERR_SAFETY);
+
+    value->data = iter->map->table[iter->current_index].value.data;
+    value->size = iter->map->table[iter->current_index].value.size;
+
+    RESULT_GUARD(s2n_map_iterator_advance(iter));
+
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_map_iterator_has_next(const struct s2n_map_iterator *iter, bool *has_next)
+{
+    RESULT_ENSURE_REF(iter);
+    RESULT_ENSURE_REF(iter->map);
+    *has_next = iter->current_index < iter->map->capacity;
+
+    return S2N_RESULT_OK;
+}
+
+int s2n_map_iterator_free(struct s2n_map_iterator *iter)
+{
+    if (iter == NULL) {
+        return S2N_SUCCESS;
+    }
+    POSIX_GUARD(s2n_free_object((uint8_t **) &iter, sizeof(struct s2n_map_iterator)));
+
+    return S2N_SUCCESS;
+}

utils/s2n_map.h

@@ -21,6 +21,7 @@
 #include "utils/s2n_result.h"
 
 struct s2n_map;
+struct s2n_map_iterator;
 
 struct s2n_map *s2n_map_new();
 struct s2n_map *s2n_map_new_with_initial_capacity(uint32_t capacity);
@@ -30,3 +31,9 @@ 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_new(struct s2n_map *map, struct s2n_map_iterator **iter);
+S2N_RESULT s2n_map_iterator_next(struct s2n_map_iterator *iter, struct s2n_blob *value);
+S2N_RESULT s2n_map_iterator_has_next(const struct s2n_map_iterator *iter, bool *has_next);
+int s2n_map_iterator_free(struct s2n_map_iterator *iter);

utils/s2n_map_internal.h

@@ -34,3 +34,9 @@ struct s2n_map {
     /* Pointer to the hash-table, should be capacity * sizeof(struct s2n_map_entry) */
     struct s2n_map_entry *table;
 };
+
+struct s2n_map_iterator {
+    struct s2n_map* map;
+    /* Index of the entry to be returned on the next `s2n_map_iterator_next()` call. */
+    uint32_t current_index;
+};