1 /*
2  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License").
5  * You may not use this file except in compliance with the License.
6  * A copy of the License is located at
7  *
8  *  http://aws.amazon.com/apache2.0
9  *
10  * or in the "license" file accompanying this file. This file is distributed
11  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12  * express or implied. See the License for the specific language governing
13  * permissions and limitations under the License.
14  */
15 
16 #include <string.h>
17 #include <stdio.h>
18 
19 #include "error/s2n_errno.h"
20 
21 #include "crypto/s2n_hash.h"
22 
23 #include "utils/s2n_safety.h"
24 #include "utils/s2n_result.h"
25 #include "utils/s2n_blob.h"
26 #include "utils/s2n_mem.h"
27 #include "utils/s2n_map.h"
28 #include "utils/s2n_map_internal.h"
29 
30 #include <s2n.h>
31 
32 #define S2N_INITIAL_TABLE_SIZE 1024
33 
s2n_map_slot(const struct s2n_map * map,struct s2n_blob * key,uint32_t * slot)34 static S2N_RESULT s2n_map_slot(const struct s2n_map *map, struct s2n_blob *key, uint32_t *slot)
35 {
36     union {
37         uint8_t u8[32];
38         uint32_t u32[8];
39     } digest;
40 
41     DEFER_CLEANUP(struct s2n_hash_state sha256 = {0}, s2n_hash_free);
42     RESULT_GUARD_POSIX(s2n_hash_new(&sha256));
43     RESULT_GUARD_POSIX(s2n_hash_init(&sha256, S2N_HASH_SHA256));
44     RESULT_GUARD_POSIX(s2n_hash_update(&sha256, key->data, key->size));
45     RESULT_GUARD_POSIX(s2n_hash_digest(&sha256, digest.u8, sizeof(digest)));
46 
47     *slot = digest.u32[0] % map->capacity;
48     return S2N_RESULT_OK;
49 }
50 
s2n_map_embiggen(struct s2n_map * map,uint32_t capacity)51 static S2N_RESULT s2n_map_embiggen(struct s2n_map *map, uint32_t capacity)
52 {
53     struct s2n_blob mem = {0};
54     struct s2n_map tmp = {0};
55 
56     RESULT_ENSURE(!map->immutable, S2N_ERR_MAP_IMMUTABLE);
57 
58     RESULT_GUARD_POSIX(s2n_alloc(&mem, (capacity * sizeof(struct s2n_map_entry))));
59     RESULT_GUARD_POSIX(s2n_blob_zero(&mem));
60 
61     tmp.capacity = capacity;
62     tmp.size = 0;
63     tmp.table = (void *) mem.data;
64     tmp.immutable = 0;
65 
66     for (uint32_t i = 0; i < map->capacity; i++) {
67         if (map->table[i].key.size) {
68             RESULT_GUARD(s2n_map_add(&tmp, &map->table[i].key, &map->table[i].value));
69             RESULT_GUARD_POSIX(s2n_free(&map->table[i].key));
70             RESULT_GUARD_POSIX(s2n_free(&map->table[i].value));
71         }
72     }
73     RESULT_GUARD_POSIX(s2n_free_object((uint8_t **)&map->table, map->capacity * sizeof(struct s2n_map_entry)));
74 
75     /* Clone the temporary map */
76     map->capacity = tmp.capacity;
77     map->size = tmp.size;
78     map->table = tmp.table;
79     map->immutable = 0;
80 
81     return S2N_RESULT_OK;
82 }
83 
s2n_map_new()84 struct s2n_map *s2n_map_new()
85 {
86     return s2n_map_new_with_initial_capacity(S2N_INITIAL_TABLE_SIZE);
87 }
88 
s2n_map_new_with_initial_capacity(uint32_t capacity)89 struct s2n_map *s2n_map_new_with_initial_capacity(uint32_t capacity)
90 {
91     PTR_ENSURE(capacity != 0, S2N_ERR_MAP_INVALID_MAP_SIZE);
92     struct s2n_blob mem = {0};
93     struct s2n_map *map;
94 
95     PTR_GUARD_POSIX(s2n_alloc(&mem, sizeof(struct s2n_map)));
96 
97     map = (void *) mem.data;
98     map->capacity = 0;
99     map->size = 0;
100     map->immutable = 0;
101     map->table = NULL;
102 
103     PTR_GUARD_RESULT(s2n_map_embiggen(map, capacity));
104 
105     return map;
106 }
107 
s2n_map_add(struct s2n_map * map,struct s2n_blob * key,struct s2n_blob * value)108 S2N_RESULT s2n_map_add(struct s2n_map *map, struct s2n_blob *key, struct s2n_blob *value)
109 {
110     RESULT_ENSURE(!map->immutable, S2N_ERR_MAP_IMMUTABLE);
111 
112     if (map->capacity < (map->size * 2)) {
113         /* Embiggen the map */
114         RESULT_GUARD(s2n_map_embiggen(map, map->capacity * 2));
115     }
116 
117     uint32_t slot = 0;
118     RESULT_GUARD(s2n_map_slot(map, key, &slot));
119 
120     /* Linear probing until we find an empty slot */
121     while(map->table[slot].key.size) {
122         if (key->size != map->table[slot].key.size ||
123             memcmp(key->data,  map->table[slot].key.data, key->size)) {
124             slot++;
125             slot %= map->capacity;
126             continue;
127         }
128 
129         /* We found a duplicate key */
130         RESULT_BAIL(S2N_ERR_MAP_DUPLICATE);
131     }
132 
133     RESULT_GUARD_POSIX(s2n_dup(key, &map->table[slot].key));
134     RESULT_GUARD_POSIX(s2n_dup(value, &map->table[slot].value));
135     map->size++;
136 
137     return S2N_RESULT_OK;
138 }
139 
s2n_map_put(struct s2n_map * map,struct s2n_blob * key,struct s2n_blob * value)140 S2N_RESULT s2n_map_put(struct s2n_map *map, struct s2n_blob *key, struct s2n_blob *value)
141 {
142     RESULT_ENSURE(!map->immutable, S2N_ERR_MAP_IMMUTABLE);
143 
144     if (map->capacity < (map->size * 2)) {
145         /* Embiggen the map */
146         RESULT_GUARD(s2n_map_embiggen(map, map->capacity * 2));
147     }
148 
149     uint32_t slot = 0;
150     RESULT_GUARD(s2n_map_slot(map, key, &slot));
151 
152     /* Linear probing until we find an empty slot */
153     while(map->table[slot].key.size) {
154         if (key->size != map->table[slot].key.size ||
155             memcmp(key->data,  map->table[slot].key.data, key->size)) {
156             slot++;
157             slot %= map->capacity;
158             continue;
159         }
160 
161         /* We found a duplicate key that will be overwritten */
162         RESULT_GUARD_POSIX(s2n_free(&map->table[slot].key));
163         RESULT_GUARD_POSIX(s2n_free(&map->table[slot].value));
164         map->size--;
165         break;
166     }
167 
168     RESULT_GUARD_POSIX(s2n_dup(key, &map->table[slot].key));
169     RESULT_GUARD_POSIX(s2n_dup(value, &map->table[slot].value));
170     map->size++;
171 
172     return S2N_RESULT_OK;
173 }
174 
s2n_map_complete(struct s2n_map * map)175 S2N_RESULT s2n_map_complete(struct s2n_map *map)
176 {
177     map->immutable = 1;
178 
179     return S2N_RESULT_OK;
180 }
181 
s2n_map_unlock(struct s2n_map * map)182 S2N_RESULT s2n_map_unlock(struct s2n_map *map)
183 {
184     map->immutable = 0;
185 
186     return S2N_RESULT_OK;
187 }
188 
s2n_map_lookup(const struct s2n_map * map,struct s2n_blob * key,struct s2n_blob * value,bool * key_found)189 S2N_RESULT s2n_map_lookup(const struct s2n_map *map, struct s2n_blob *key, struct s2n_blob *value, bool *key_found)
190 {
191     RESULT_ENSURE(map->immutable, S2N_ERR_MAP_MUTABLE);
192 
193     uint32_t slot = 0;
194     RESULT_GUARD(s2n_map_slot(map, key, &slot));
195     const uint32_t initial_slot = slot;
196 
197     while(map->table[slot].key.size) {
198         if (key->size != map->table[slot].key.size ||
199             memcmp(key->data,  map->table[slot].key.data, key->size)) {
200             slot++;
201             slot %= map->capacity;
202             /* We went over all the slots but found no match */
203             if (slot == initial_slot) {
204                 break;
205             }
206             continue;
207         }
208 
209         /* We found a match */
210         value->data = map->table[slot].value.data;
211         value->size = map->table[slot].value.size;
212 
213         *key_found = true;
214 
215         return S2N_RESULT_OK;
216     }
217 
218     *key_found = false;
219 
220     return S2N_RESULT_OK;
221 }
222 
s2n_map_free(struct s2n_map * map)223 S2N_RESULT s2n_map_free(struct s2n_map *map)
224 {
225     /* Free the keys and values */
226     for (uint32_t i = 0; i < map->capacity; i++) {
227         if (map->table[i].key.size) {
228             RESULT_GUARD_POSIX(s2n_free(&map->table[i].key));
229             RESULT_GUARD_POSIX(s2n_free(&map->table[i].value));
230         }
231     }
232 
233     /* Free the table */
234     RESULT_GUARD_POSIX(s2n_free_object((uint8_t **)&map->table, map->capacity * sizeof(struct s2n_map_entry)));
235 
236     /* And finally the map */
237     RESULT_GUARD_POSIX(s2n_free_object((uint8_t **)&map, sizeof(struct s2n_map)));
238 
239     return S2N_RESULT_OK;
240 }
241