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 <sys/queue.h>
  23 #include <time.h>
  24
  25 #include <loc/libloc.h>
  26 #include <loc/as.h>
  27 #include <loc/format.h>
  28 #include <loc/network.h>
  29 #include <loc/private.h>
  30 #include <loc/writer.h>
  31
  32 struct loc_writer {
  33         struct loc_ctx* ctx;
  34         int refcount;
  35
  36         struct loc_stringpool* pool;
  37         off_t vendor;
  38         off_t description;
  39         off_t license;
  40
  41         struct loc_as** as;
  42         size_t as_count;
  43
  44         struct loc_network_tree* networks;
  45 };
  46
  47 LOC_EXPORT int loc_writer_new(struct loc_ctx* ctx, struct loc_writer** writer) {
  48         struct loc_writer* w = calloc(1, sizeof(*w));
  49         if (!w)
  50                 return -ENOMEM;
  51
  52         w->ctx = loc_ref(ctx);
  53         w->refcount = 1;
  54
  55         int r = loc_stringpool_new(ctx, &w->pool);
  56         if (r) {
  57                 loc_writer_unref(w);
  58                 return r;
  59         }
  60
  61         // Initialize the network tree
  62         r = loc_network_tree_new(ctx, &w->networks);
  63         if (r) {
  64                 loc_writer_unref(w);
  65                 return r;
  66         }
  67
  68         *writer = w;
  69         return 0;
  70 }
  71
  72 LOC_EXPORT struct loc_writer* loc_writer_ref(struct loc_writer* writer) {
  73         writer->refcount++;
  74
  75         return writer;
  76 }
  77
  78 static void loc_writer_free(struct loc_writer* writer) {
  79         DEBUG(writer->ctx, "Releasing writer at %p\n", writer);
  80
  81         // Unref all AS
  82         for (unsigned int i = 0; i < writer->as_count; i++) {
  83                 loc_as_unref(writer->as[i]);
  84         }
  85
  86         // Release network tree
  87         if (writer->networks)
  88                 loc_network_tree_unref(writer->networks);
  89
  90         // Unref the string pool
  91         loc_stringpool_unref(writer->pool);
  92
  93         loc_unref(writer->ctx);
  94         free(writer);
  95 }
  96
  97 LOC_EXPORT struct loc_writer* loc_writer_unref(struct loc_writer* writer) {
  98         if (--writer->refcount > 0)
  99                 return writer;
 100
 101         loc_writer_free(writer);
 102
 103         return NULL;
 104 }
 105
 106 LOC_EXPORT const char* loc_writer_get_vendor(struct loc_writer* writer) {
 107         return loc_stringpool_get(writer->pool, writer->vendor);
 108 }
 109
 110 LOC_EXPORT int loc_writer_set_vendor(struct loc_writer* writer, const char* vendor) {
 111         // Add the string to the string pool
 112         off_t offset = loc_stringpool_add(writer->pool, vendor);
 113         if (offset < 0)
 114                 return offset;
 115
 116         writer->vendor = offset;
 117         return 0;
 118 }
 119
 120 LOC_EXPORT const char* loc_writer_get_description(struct loc_writer* writer) {
 121         return loc_stringpool_get(writer->pool, writer->description);
 122 }
 123
 124 LOC_EXPORT int loc_writer_set_description(struct loc_writer* writer, const char* description) {
 125         // Add the string to the string pool
 126         off_t offset = loc_stringpool_add(writer->pool, description);
 127         if (offset < 0)
 128                 return offset;
 129
 130         writer->description = offset;
 131         return 0;
 132 }
 133
 134 LOC_EXPORT const char* loc_writer_get_license(struct loc_writer* writer) {
 135         return loc_stringpool_get(writer->pool, writer->license);
 136 }
 137
 138 LOC_EXPORT int loc_writer_set_license(struct loc_writer* writer, const char* license) {
 139         // Add the string to the string pool
 140         off_t offset = loc_stringpool_add(writer->pool, license);
 141         if (offset < 0)
 142                 return offset;
 143
 144         writer->license = offset;
 145         return 0;
 146 }
 147
 148 static int __loc_as_cmp(const void* as1, const void* as2) {
 149         return loc_as_cmp(*(struct loc_as**)as1, *(struct loc_as**)as2);
 150 }
 151
 152 LOC_EXPORT int loc_writer_add_as(struct loc_writer* writer, struct loc_as** as, uint32_t number) {
 153         int r = loc_as_new(writer->ctx, as, number);
 154         if (r)
 155                 return r;
 156
 157         // We have a new AS to add
 158         writer->as_count++;
 159
 160         // Make space
 161         writer->as = realloc(writer->as, sizeof(*writer->as) * writer->as_count);
 162         if (!writer->as)
 163                 return -ENOMEM;
 164
 165         // Add as last element
 166         writer->as[writer->as_count - 1] = loc_as_ref(*as);
 167
 168         // Sort everything
 169         qsort(writer->as, writer->as_count, sizeof(*writer->as), __loc_as_cmp);
 170
 171         return 0;
 172 }
 173
 174 LOC_EXPORT int loc_writer_add_network(struct loc_writer* writer, struct loc_network** network, const char* string) {
 175         int r;
 176
 177         // Create a new network object
 178         r = loc_network_new_from_string(writer->ctx, network, string);
 179         if (r)
 180                 return r;
 181
 182         // Add it to the local tree
 183         return loc_network_tree_add_network(writer->networks, *network);
 184 }
 185
 186 static void make_magic(struct loc_writer* writer, struct loc_database_magic* magic) {
 187         // Copy magic bytes
 188         for (unsigned int i = 0; i < strlen(LOC_DATABASE_MAGIC); i++)
 189                 magic->magic[i] = LOC_DATABASE_MAGIC[i];
 190
 191         // Set version
 192         magic->version = htobe16(LOC_DATABASE_VERSION);
 193 }
 194
 195 static void align_page_boundary(off_t* offset, FILE* f) {
 196         // Move to next page boundary
 197         while (*offset % LOC_DATABASE_PAGE_SIZE > 0)
 198                 *offset += fwrite("", 1, 1, f);
 199 }
 200
 201 static int loc_database_write_pool(struct loc_writer* writer,
 202                 struct loc_database_header_v0* header, off_t* offset, FILE* f) {
 203         // Save the offset of the pool section
 204         DEBUG(writer->ctx, "Pool starts at %jd bytes\n", *offset);
 205         header->pool_offset = htobe32(*offset);
 206
 207         // Write the pool
 208         size_t pool_length = loc_stringpool_write(writer->pool, f);
 209         *offset += pool_length;
 210
 211         DEBUG(writer->ctx, "Pool has a length of %zu bytes\n", pool_length);
 212         header->pool_length = htobe32(pool_length);
 213
 214         return 0;
 215 }
 216
 217 static int loc_database_write_as_section(struct loc_writer* writer,
 218                 struct loc_database_header_v0* header, off_t* offset, FILE* f) {
 219         DEBUG(writer->ctx, "AS section starts at %jd bytes\n", *offset);
 220         header->as_offset = htobe32(*offset);
 221
 222         size_t as_length = 0;
 223
 224         struct loc_database_as_v0 as;
 225         for (unsigned int i = 0; i < writer->as_count; i++) {
 226                 // Convert AS into database format
 227                 loc_as_to_database_v0(writer->as[i], writer->pool, &as);
 228
 229                 // Write to disk
 230                 *offset += fwrite(&as, 1, sizeof(as), f);
 231                 as_length += sizeof(as);
 232         }
 233
 234         DEBUG(writer->ctx, "AS section has a length of %zu bytes\n", as_length);
 235         header->as_length = htobe32(as_length);
 236
 237         align_page_boundary(offset, f);
 238
 239         return 0;
 240 }
 241
 242 struct node {
 243         TAILQ_ENTRY(node) nodes;
 244
 245         struct loc_network_tree_node* node;
 246
 247         // Indices of the child nodes
 248         uint32_t index_zero;
 249         uint32_t index_one;
 250 };
 251
 252 static struct node* make_node(struct loc_network_tree_node* node) {
 253         struct node* n = malloc(sizeof(*n));
 254         if (!n)
 255                 return NULL;
 256
 257         n->node  = loc_network_tree_node_ref(node);
 258         n->index_zero = n->index_one = 0;
 259
 260         return n;
 261 }
 262
 263 static void free_node(struct node* node) {
 264         loc_network_tree_node_unref(node->node);
 265
 266         free(node);
 267 }
 268
 269 struct network {
 270         TAILQ_ENTRY(network) networks;
 271
 272         struct loc_network* network;
 273 };
 274
 275 static struct network* make_network(struct loc_network* network) {
 276         struct network* n = malloc(sizeof(*n));
 277         if (!n)
 278                 return NULL;
 279
 280         n->network = loc_network_ref(network);
 281
 282         return n;
 283 }
 284
 285 static void free_network(struct network* network) {
 286         loc_network_unref(network->network);
 287
 288         free(network);
 289 }
 290
 291 static int loc_database_write_networks(struct loc_writer* writer,
 292                 struct loc_database_header_v0* header, off_t* offset, FILE* f) {
 293         // Write the network tree
 294         DEBUG(writer->ctx, "Network tree starts at %jd bytes\n", *offset);
 295         header->network_tree_offset = htobe32(*offset);
 296
 297         size_t network_tree_length = 0;
 298         size_t network_data_length = 0;
 299
 300         struct node* node;
 301         struct node* child_node;
 302
 303         uint32_t index = 0;
 304         uint32_t network_index = 0;
 305
 306         struct loc_database_network_v0 db_network;
 307         struct loc_database_network_node_v0 db_node;
 308
 309         // Initialize queue for nodes
 310         TAILQ_HEAD(node_t, node) nodes;
 311         TAILQ_INIT(&nodes);
 312
 313         // Initialize queue for networks
 314         TAILQ_HEAD(network_t, network) networks;
 315         TAILQ_INIT(&networks);
 316
 317         // Add root
 318         struct loc_network_tree_node* root = loc_network_tree_get_root(writer->networks);
 319         node = make_node(root);
 320
 321         TAILQ_INSERT_TAIL(&nodes, node, nodes);
 322
 323         while (!TAILQ_EMPTY(&nodes)) {
 324                 // Pop first node in list
 325                 node = TAILQ_FIRST(&nodes);
 326                 TAILQ_REMOVE(&nodes, node, nodes);
 327
 328                 DEBUG(writer->ctx, "Processing node %p\n", node);
 329
 330                 // Get child nodes
 331                 struct loc_network_tree_node* node_zero = loc_network_tree_node_get(node->node, 0);
 332                 if (node_zero) {
 333                         node->index_zero = ++index;
 334
 335                         child_node = make_node(node_zero);
 336                         loc_network_tree_node_unref(node_zero);
 337
 338                         TAILQ_INSERT_TAIL(&nodes, child_node, nodes);
 339                 }
 340
 341                 struct loc_network_tree_node* node_one = loc_network_tree_node_get(node->node, 1);
 342                 if (node_one) {
 343                         node->index_one = ++index;
 344
 345                         child_node = make_node(node_one);
 346                         loc_network_tree_node_unref(node_one);
 347
 348                         TAILQ_INSERT_TAIL(&nodes, child_node, nodes);
 349                 }
 350
 351                 // Prepare what we are writing to disk
 352                 db_node.zero = htobe32(node->index_zero);
 353                 db_node.one  = htobe32(node->index_one);
 354
 355                 if (loc_network_tree_node_is_leaf(node->node)) {
 356                         struct loc_network* network = loc_network_tree_node_get_network(node->node);
 357
 358                         // Append network to be written out later
 359                         struct network* nw = make_network(network);
 360                         TAILQ_INSERT_TAIL(&networks, nw, networks);
 361
 362                         db_node.network = htobe32(network_index++);
 363                         loc_network_unref(network);
 364                 } else {
 365                         db_node.network = htobe32(0xffffffff);
 366                 }
 367
 368                 // Write the current node
 369                 DEBUG(writer->ctx, "Writing node %p (0 = %d, 1 = %d)\n",
 370                         node, node->index_zero, node->index_one);
 371
 372                 *offset += fwrite(&db_node, 1, sizeof(db_node), f);
 373                 network_tree_length += sizeof(db_node);
 374
 375                 free_node(node);
 376         }
 377
 378         loc_network_tree_node_unref(root);
 379
 380         header->network_tree_length = htobe32(network_tree_length);
 381
 382         align_page_boundary(offset, f);
 383
 384         DEBUG(writer->ctx, "Networks data section starts at %jd bytes\n", *offset);
 385         header->network_data_offset = htobe32(*offset);
 386
 387         // We have now written the entire tree and have all networks
 388         // in a queue in order as they are indexed
 389         while (!TAILQ_EMPTY(&networks)) {
 390                 struct network* nw = TAILQ_FIRST(&networks);
 391                 TAILQ_REMOVE(&networks, nw, networks);
 392
 393                 // Prepare what we are writing to disk
 394                 int r = loc_network_to_database_v0(nw->network, &db_network);
 395                 if (r)
 396                         return r;
 397
 398                 *offset += fwrite(&db_network, 1, sizeof(db_network), f);
 399                 network_data_length += sizeof(db_network);
 400
 401                 free_network(nw);
 402         }
 403
 404         header->network_data_length = htobe32(network_data_length);
 405
 406         align_page_boundary(offset, f);
 407
 408         return 0;
 409 }
 410
 411 LOC_EXPORT int loc_writer_write(struct loc_writer* writer, FILE* f) {
 412         struct loc_database_magic magic;
 413         make_magic(writer, &magic);
 414
 415         // Make the header
 416         struct loc_database_header_v0 header;
 417         header.vendor      = htobe32(writer->vendor);
 418         header.description = htobe32(writer->description);
 419         header.license     = htobe32(writer->license);
 420
 421         time_t now = time(NULL);
 422         header.created_at = htobe64(now);
 423
 424         int r;
 425         off_t offset = 0;
 426
 427         // Start writing at the beginning of the file
 428         r = fseek(f, 0, SEEK_SET);
 429         if (r)
 430                 return r;
 431
 432         // Write the magic
 433         offset += fwrite(&magic, 1, sizeof(magic), f);
 434
 435         // Skip the space we need to write the header later
 436         r = fseek(f, sizeof(header), SEEK_CUR);
 437         if (r) {
 438                 DEBUG(writer->ctx, "Could not seek to position after header\n");
 439                 return r;
 440         }
 441         offset += sizeof(header);
 442
 443         align_page_boundary(&offset, f);
 444
 445         // Write all ASes
 446         r = loc_database_write_as_section(writer, &header, &offset, f);
 447         if (r)
 448                 return r;
 449
 450         // Write all networks
 451         r = loc_database_write_networks(writer, &header, &offset, f);
 452         if (r)
 453                 return r;
 454
 455         // Write pool
 456         r = loc_database_write_pool(writer, &header, &offset, f);
 457         if (r)
 458                 return r;
 459
 460         // Write the header
 461         r = fseek(f, sizeof(magic), SEEK_SET);
 462         if (r)
 463                 return r;
 464
 465         offset += fwrite(&header, 1, sizeof(header), f);
 466
 467         return 0;
 468 }