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 }