From 865479d336fda5be804f08bd2f736189137e248f Mon Sep 17 00:00:00 2001 From: James Mayclin Date: Thu, 25 Jan 2024 04:49:13 +0000 Subject: [PATCH] utils: add map iteration iterator --- tests/unit/s2n_map_test.c | 115 ++++++++++++++++++++++++++++++++++++++ utils/s2n_map.c | 88 +++++++++++++++++++++++++++++ utils/s2n_map.h | 7 +++ utils/s2n_map_internal.h | 6 ++ 4 files changed, 216 insertions(+) diff --git a/tests/unit/s2n_map_test.c b/tests/unit/s2n_map_test.c index 8f1d8bb8437..9f00d4fa8fb 100644 --- a/tests/unit/s2n_map_test.c +++ b/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(); } diff --git a/utils/s2n_map.c b/utils/s2n_map.c index e95679553e0..b5171d94681 100644 --- a/utils/s2n_map.c +++ b/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; +} diff --git a/utils/s2n_map.h b/utils/s2n_map.h index 259082936cf..c49b95a3676 100644 --- a/utils/s2n_map.h +++ b/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); diff --git a/utils/s2n_map_internal.h b/utils/s2n_map_internal.h index 01d629fb16f..dfc42062b32 100644 --- a/utils/s2n_map_internal.h +++ b/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; +};