1 /*
2  * Copyright (c) 2022 Stefan Sperling <stsp@openbsd.org>
3  *
4  * Permission to use, copy, modify, and distribute this software for any
5  * purpose with or without fee is hereby granted, provided that the above
6  * copyright notice and this permission notice appear in all copies.
7  *
8  * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
9  * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
10  * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
11  * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
12  * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
13  * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
14  * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
15  */
16 
17 #include <stdlib.h>
18 #include <stdint.h>
19 #include <string.h>
20 #include <stdio.h>
21 #include <zlib.h>
22 #include <limits.h>
23 #include <time.h>
24 #include <errno.h>
25 
26 #include "got_compat.h"
27 
28 #include "got_object.h"
29 #include "got_error.h"
30 
31 #include "got_lib_delta.h"
32 #include "got_lib_inflate.h"
33 #include "got_lib_object.h"
34 #include "got_lib_object_idset.h"
35 #include "got_lib_object_parse.h"
36 
37 #define GOT_OBJECT_IDSET_MIN_BUCKETS	64
38 
39 struct got_object_idset {
40 	struct got_object_id_queue *ids;
41 	size_t nbuckets;
42 	unsigned int totelem;
43 	unsigned int flags;
44 #define GOT_OBJECT_IDSET_F_TRAVERSAL	0x01
45 #define GOT_OBJECT_IDSET_F_NOMEM	0x02
46 	SIPHASH_KEY key;
47 };
48 
49 struct got_object_idset *
50 got_object_idset_alloc(void)
51 {
52 	struct got_object_idset *set;
53 	int i;
54 
55 	set = malloc(sizeof(*set));
56 	if (set == NULL)
57 		return NULL;
58 
59 	set->ids = calloc(sizeof(set->ids[0]), GOT_OBJECT_IDSET_MIN_BUCKETS);
60 	if (set->ids == NULL) {
61 		free(set);
62 		return NULL;
63 	}
64 	for (i = 0; i < GOT_OBJECT_IDSET_MIN_BUCKETS; i++)
65 		STAILQ_INIT(&set->ids[i]);
66 
67 	set->totelem = 0;
68 	set->nbuckets = GOT_OBJECT_IDSET_MIN_BUCKETS;
69 	set->flags = 0;
70 	arc4random_buf(&set->key, sizeof(set->key));
71 	return set;
72 }
73 
74 void
75 got_object_idset_free(struct got_object_idset *set)
76 {
77 	size_t i;
78 	struct got_object_qid *qid;
79 
80 	for (i = 0; i < set->nbuckets; i++) {
81 		while (!STAILQ_EMPTY(&set->ids[i])) {
82 			qid = STAILQ_FIRST(&set->ids[i]);
83 			STAILQ_REMOVE(&set->ids[i], qid, got_object_qid, entry);
84 			got_object_qid_free(qid);
85 		}
86 	}
87 	/* User data should be freed by caller. */
88 	free(set->ids);
89 	free(set);
90 }
91 
92 static uint64_t
93 idset_hash(struct got_object_idset *set, struct got_object_id *id)
94 {
95 	return SipHash24(&set->key, id->sha1, sizeof(id->sha1));
96 }
97 
98 static const struct got_error *
99 idset_resize(struct got_object_idset *set, size_t nbuckets)
100 {
101 	struct got_object_id_queue *ids;
102 	size_t i;
103 
104 	ids = calloc(nbuckets, sizeof(ids[0]));
105 	if (ids == NULL) {
106 		if (errno != ENOMEM)
107 			return got_error_from_errno("calloc");
108 		/* Proceed with our current amount of hash buckets. */
109 		set->flags |= GOT_OBJECT_IDSET_F_NOMEM;
110 		return NULL;
111 	}
112 
113 	for (i = 0; i < nbuckets; i++)
114 		STAILQ_INIT(&ids[i]);
115 
116 	arc4random_buf(&set->key, sizeof(set->key));
117 
118 	for (i = 0; i < set->nbuckets; i++) {
119 		while (!STAILQ_EMPTY(&set->ids[i])) {
120 			struct got_object_qid *qid;
121 			uint64_t idx;
122 			qid = STAILQ_FIRST(&set->ids[i]);
123 			STAILQ_REMOVE(&set->ids[i], qid, got_object_qid, entry);
124 			idx = idset_hash(set, &qid->id) % nbuckets;
125 			STAILQ_INSERT_HEAD(&ids[idx], qid, entry);
126 		}
127 	}
128 
129 	free(set->ids);
130 	set->ids = ids;
131 	set->nbuckets = nbuckets;
132 	return NULL;
133 }
134 
135 static const struct got_error *
136 idset_grow(struct got_object_idset *set)
137 {
138 	size_t nbuckets;
139 
140 	if (set->flags & GOT_OBJECT_IDSET_F_NOMEM)
141 		return NULL;
142 
143 	if (set->nbuckets >= UINT_MAX / 2)
144 		nbuckets = UINT_MAX;
145 	else
146 		nbuckets = set->nbuckets * 2;
147 
148 	return idset_resize(set, nbuckets);
149 }
150 
151 const struct got_error *
152 got_object_idset_add(struct got_object_idset *set, struct got_object_id *id,
153     void *data)
154 {
155 	const struct got_error *err;
156 	struct got_object_qid *qid;
157 	uint64_t idx;
158 	struct got_object_id_queue *head;
159 
160 	/* This function may resize the set. */
161 	if (set->flags & GOT_OBJECT_IDSET_F_TRAVERSAL)
162 		return got_error_msg(GOT_ERR_NOT_IMPL,
163 		    "cannot add elements to idset during traversal");
164 
165 	if (set->totelem == UINT_MAX)
166 		return got_error(GOT_ERR_NO_SPACE);
167 		
168 	err = got_object_qid_alloc_partial(&qid);
169 	if (err)
170 		return err;
171 	memcpy(&qid->id, id, sizeof(qid->id));
172 	qid->data = data;
173 
174 	idx = idset_hash(set, id) % set->nbuckets;
175 	head = &set->ids[idx];
176 	STAILQ_INSERT_HEAD(head, qid, entry);
177 	set->totelem++;
178 
179 	if (set->nbuckets < set->totelem)
180 		err = idset_grow(set);
181 
182 	return err;
183 }
184 
185 static struct got_object_qid *
186 find_element(struct got_object_idset *set, struct got_object_id *id)
187 {
188 	uint64_t idx = idset_hash(set, id) % set->nbuckets;
189 	struct got_object_id_queue *head = &set->ids[idx];
190 	struct got_object_qid *qid;
191 
192 	STAILQ_FOREACH(qid, head, entry) {
193 		if (got_object_id_cmp(&qid->id, id) == 0)
194 			return qid;
195 	}
196 
197 	return NULL;
198 }
199 
200 void *
201 got_object_idset_get(struct got_object_idset *set, struct got_object_id *id)
202 {
203 	struct got_object_qid *qid = find_element(set, id);
204 	return qid ? qid->data : NULL;
205 }
206 
207 const struct got_error *
208 got_object_idset_remove(void **data, struct got_object_idset *set,
209     struct got_object_id *id)
210 {
211 	uint64_t idx;
212 	struct got_object_id_queue *head;
213 	struct got_object_qid *qid;
214 
215 	if (data)
216 		*data = NULL;
217 
218 	if (set->totelem == 0)
219 		return got_error(GOT_ERR_NO_OBJ);
220 
221 	if (id == NULL) {
222 		/* Remove a "random" element. */
223 		for (idx = 0; idx < set->nbuckets; idx++) {
224 			head = &set->ids[idx];
225 			qid = STAILQ_FIRST(head);
226 			if (qid)
227 				break;
228 		}
229 	} else {
230 		idx = idset_hash(set, id) % set->nbuckets;
231 		head = &set->ids[idx];
232 		STAILQ_FOREACH(qid, head, entry) {
233 			if (got_object_id_cmp(&qid->id, id) == 0)
234 				break;
235 		}
236 		if (qid == NULL)
237 			return got_error_no_obj(id);
238 	}
239 
240 	if (data)
241 		*data = qid->data;
242 	STAILQ_REMOVE(head, qid, got_object_qid, entry);
243 	got_object_qid_free(qid);
244 	set->totelem--;
245 
246 	return NULL;
247 }
248 
249 int
250 got_object_idset_contains(struct got_object_idset *set,
251     struct got_object_id *id)
252 {
253 	struct got_object_qid *qid = find_element(set, id);
254 	return qid ? 1 : 0;
255 }
256 
257 const struct got_error *
258 got_object_idset_for_each(struct got_object_idset *set,
259     const struct got_error *(*cb)(struct got_object_id *, void *, void *),
260     void *arg)
261 {
262 	const struct got_error *err = NULL;
263 	struct got_object_id_queue *head;
264 	struct got_object_qid *qid, *tmp;
265 	size_t i;
266 
267 	set->flags |= GOT_OBJECT_IDSET_F_TRAVERSAL;
268 	for (i = 0; i < set->nbuckets; i++) {
269 		head = &set->ids[i];
270 		STAILQ_FOREACH_SAFE(qid, head, entry, tmp) {
271 			err = (*cb)(&qid->id, qid->data, arg);
272 			if (err)
273 				goto done;
274 		}
275 	}
276 done:
277 	set->flags &= ~GOT_OBJECT_IDSET_F_TRAVERSAL;
278 	return err;
279 }
280 
281 int
282 got_object_idset_num_elements(struct got_object_idset *set)
283 {
284 	return set->totelem;
285 }