writer: Correctly increase offset when writing ASes
[people/ms/libloc.git] / src / writer.c
1 /*
2         libloc - A library to determine the location of someone on the Internet
3
4         Copyright (C) 2017 IPFire Development Team <info@ipfire.org>
5
6         This library is free software; you can redistribute it and/or
7         modify it under the terms of the GNU Lesser General Public
8         License as published by the Free Software Foundation; either
9         version 2.1 of the License, or (at your option) any later version.
10
11         This library is distributed in the hope that it will be useful,
12         but WITHOUT ANY WARRANTY; without even the implied warranty of
13         MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14         Lesser General Public License for more details.
15 */
16
17 #include <endian.h>
18 #include <errno.h>
19 #include <stdio.h>
20 #include <stdlib.h>
21 #include <string.h>
22 #include <time.h>
23
24 #include <loc/libloc.h>
25 #include <loc/as.h>
26 #include <loc/format.h>
27 #include <loc/network.h>
28 #include <loc/private.h>
29 #include <loc/writer.h>
30
31 struct loc_writer {
32         struct loc_ctx* ctx;
33         int refcount;
34
35         struct loc_stringpool* pool;
36         off_t vendor;
37         off_t description;
38
39         struct loc_as** as;
40         size_t as_count;
41
42         struct loc_network_tree* networks;
43 };
44
45 LOC_EXPORT int loc_writer_new(struct loc_ctx* ctx, struct loc_writer** writer) {
46         struct loc_writer* w = calloc(1, sizeof(*w));
47         if (!w)
48                 return -ENOMEM;
49
50         w->ctx = loc_ref(ctx);
51         w->refcount = 1;
52
53         int r = loc_stringpool_new(ctx, &w->pool);
54         if (r) {
55                 loc_writer_unref(w);
56                 return r;
57         }
58
59         // Initialize the network tree
60         r = loc_network_tree_new(ctx, &w->networks);
61         if (r) {
62                 loc_writer_unref(w);
63                 return r;
64         }
65
66         *writer = w;
67         return 0;
68 }
69
70 LOC_EXPORT struct loc_writer* loc_writer_ref(struct loc_writer* writer) {
71         writer->refcount++;
72
73         return writer;
74 }
75
76 static void loc_writer_free(struct loc_writer* writer) {
77         DEBUG(writer->ctx, "Releasing writer at %p\n", writer);
78
79         // Unref all AS
80         for (unsigned int i = 0; i < writer->as_count; i++) {
81                 loc_as_unref(writer->as[i]);
82         }
83
84         // Release network tree
85         if (writer->networks)
86                 loc_network_tree_unref(writer->networks);
87
88         // Unref the string pool
89         loc_stringpool_unref(writer->pool);
90
91         loc_unref(writer->ctx);
92         free(writer);
93 }
94
95 LOC_EXPORT struct loc_writer* loc_writer_unref(struct loc_writer* writer) {
96         if (--writer->refcount > 0)
97                 return writer;
98
99         loc_writer_free(writer);
100
101         return NULL;
102 }
103
104 LOC_EXPORT const char* loc_writer_get_vendor(struct loc_writer* writer) {
105         return loc_stringpool_get(writer->pool, writer->vendor);
106 }
107
108 LOC_EXPORT int loc_writer_set_vendor(struct loc_writer* writer, const char* vendor) {
109         // Add the string to the string pool
110         off_t offset = loc_stringpool_add(writer->pool, vendor);
111         if (offset < 0)
112                 return offset;
113
114         writer->vendor = offset;
115         return 0;
116 }
117
118 LOC_EXPORT const char* loc_writer_get_description(struct loc_writer* writer) {
119         return loc_stringpool_get(writer->pool, writer->description);
120 }
121
122 LOC_EXPORT int loc_writer_set_description(struct loc_writer* writer, const char* description) {
123         // Add the string to the string pool
124         off_t offset = loc_stringpool_add(writer->pool, description);
125         if (offset < 0)
126                 return offset;
127
128         writer->description = offset;
129         return 0;
130 }
131
132 static int __loc_as_cmp(const void* as1, const void* as2) {
133         return loc_as_cmp(*(struct loc_as**)as1, *(struct loc_as**)as2);
134 }
135
136 LOC_EXPORT int loc_writer_add_as(struct loc_writer* writer, struct loc_as** as, uint32_t number) {
137         int r = loc_as_new(writer->ctx, writer->pool, as, number);
138         if (r)
139                 return r;
140
141         // We have a new AS to add
142         writer->as_count++;
143
144         // Make space
145         writer->as = realloc(writer->as, sizeof(*writer->as) * writer->as_count);
146         if (!writer->as)
147                 return -ENOMEM;
148
149         // Add as last element
150         writer->as[writer->as_count - 1] = loc_as_ref(*as);
151
152         // Sort everything
153         qsort(writer->as, writer->as_count, sizeof(*writer->as), __loc_as_cmp);
154
155         return 0;
156 }
157
158 LOC_EXPORT int loc_writer_add_network(struct loc_writer* writer, struct loc_network** network, const char* string) {
159         int r;
160
161         // Create a new network object
162         r = loc_network_new_from_string(writer->ctx, network, string);
163         if (r)
164                 return r;
165
166         // Add it to the local tree
167         return loc_network_tree_add_network(writer->networks, *network);
168 }
169
170 static void make_magic(struct loc_writer* writer, struct loc_database_magic* magic) {
171         // Copy magic bytes
172         for (unsigned int i = 0; i < strlen(LOC_DATABASE_MAGIC); i++)
173                 magic->magic[i] = LOC_DATABASE_MAGIC[i];
174
175         // Set version
176         magic->version = htobe16(LOC_DATABASE_VERSION);
177 }
178
179 static void align_page_boundary(off_t* offset, FILE* f) {
180         // Move to next page boundary
181         while (*offset % LOC_DATABASE_PAGE_SIZE > 0)
182                 *offset += fwrite("", 1, 1, f);
183 }
184
185 static int loc_database_write_pool(struct loc_writer* writer,
186                 struct loc_database_header_v0* header, off_t* offset, FILE* f) {
187         // Save the offset of the pool section
188         DEBUG(writer->ctx, "Pool starts at %jd bytes\n", *offset);
189         header->pool_offset = htobe32(*offset);
190
191         // Write the pool
192         size_t pool_length = loc_stringpool_write(writer->pool, f);
193         *offset += pool_length;
194
195         DEBUG(writer->ctx, "Pool has a length of %zu bytes\n", pool_length);
196         header->pool_length = htobe32(pool_length);
197
198         return 0;
199 }
200
201 static int loc_database_write_as_section(struct loc_writer* writer,
202                 struct loc_database_header_v0* header, off_t* offset, FILE* f) {
203         DEBUG(writer->ctx, "AS section starts at %jd bytes\n", *offset);
204         header->as_offset = htobe32(*offset);
205
206         size_t as_length = 0;
207
208         struct loc_database_as_v0 as;
209         for (unsigned int i = 0; i < writer->as_count; i++) {
210                 // Convert AS into database format
211                 loc_as_to_database_v0(writer->as[i], &as);
212
213                 // Write to disk
214                 *offset += fwrite(&as, 1, sizeof(as), f);
215                 as_length += sizeof(as);
216         }
217
218         DEBUG(writer->ctx, "AS section has a length of %zu bytes\n", as_length);
219         header->as_length = htobe32(as_length);
220
221         return 0;
222 }
223
224 static int loc_database_write_network_section(struct loc_network* network, void* data) {
225         FILE* f = (FILE*)data;
226
227         struct loc_database_network_v0 n;
228
229         int r = loc_network_to_database_v0(network, &n);
230         if (r)
231                 return r;
232
233         fwrite(&n, 1, sizeof(n), f);
234
235         return 0;
236 }
237
238 static int loc_database_write_networks_section(struct loc_writer* writer,
239                 struct loc_database_header_v0* header, off_t* offset, FILE* f) {
240         DEBUG(writer->ctx, "Networks section starts at %jd bytes\n", *offset);
241         header->networks_offset = htobe32(*offset);
242
243         size_t networks_length = sizeof(struct loc_database_network_v0)
244                 * loc_network_tree_count_networks(writer->networks);
245         offset += networks_length;
246
247         int r = loc_network_tree_walk(writer->networks, NULL, loc_database_write_network_section, f);
248         if (r)
249                 return r;
250
251         header->networks_length = htobe32(networks_length);
252
253         return 0;
254 }
255
256 LOC_EXPORT int loc_writer_write(struct loc_writer* writer, FILE* f) {
257         struct loc_database_magic magic;
258         make_magic(writer, &magic);
259
260         // Make the header
261         struct loc_database_header_v0 header;
262         header.vendor      = htobe32(writer->vendor);
263         header.description = htobe32(writer->description);
264
265         time_t now = time(NULL);
266         header.created_at = htobe64(now);
267
268         int r;
269         off_t offset = 0;
270
271         // Start writing at the beginning of the file
272         r = fseek(f, 0, SEEK_SET);
273         if (r)
274                 return r;
275
276         // Write the magic
277         offset += fwrite(&magic, 1, sizeof(magic), f);
278
279         // Skip the space we need to write the header later
280         r = fseek(f, sizeof(header), SEEK_CUR);
281         if (r) {
282                 DEBUG(writer->ctx, "Could not seek to position after header\n");
283                 return r;
284         }
285         offset += sizeof(header);
286
287         align_page_boundary(&offset, f);
288
289         // Write all ASes
290         r = loc_database_write_as_section(writer, &header, &offset, f);
291         if (r)
292                 return r;
293
294         align_page_boundary(&offset, f);
295
296         // Write all networks
297         r = loc_database_write_networks_section(writer, &header, &offset, f);
298         if (r)
299                 return r;
300
301         align_page_boundary(&offset, f);
302
303         // Write pool
304         r = loc_database_write_pool(writer, &header, &offset, f);
305         if (r)
306                 return r;
307
308         // Write the header
309         r = fseek(f, sizeof(magic), SEEK_SET);
310         if (r)
311                 return r;
312
313         offset += fwrite(&header, 1, sizeof(header), f);
314
315         return 0;
316 }