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