src/database.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 <arpa/inet.h>
  18 #include <ctype.h>
  19 #include <errno.h>
  20 #include <netinet/in.h>
  21 #include <stddef.h>
  22 #include <stdint.h>
  23 #include <stdio.h>
  24 #include <stdlib.h>
  25 #include <string.h>
  26 #include <sys/mman.h>
  27 #include <sys/types.h>
  28 #include <time.h>
  29 #include <unistd.h>
  30
  31 #ifdef HAVE_ENDIAN_H
  32 #  include <endian.h>
  33 #endif
  34
  35 #include <openssl/evp.h>
  36
  37 #include <loc/libloc.h>
  38 #include <loc/as.h>
  39 #include <loc/compat.h>
  40 #include <loc/country.h>
  41 #include <loc/database.h>
  42 #include <loc/format.h>
  43 #include <loc/network.h>
  44 #include <loc/private.h>
  45 #include <loc/stringpool.h>
  46
  47 struct loc_database {
  48         struct loc_ctx* ctx;
  49         int refcount;
  50
  51         FILE* f;
  52
  53         unsigned int version;
  54         time_t created_at;
  55         off_t vendor;
  56         off_t description;
  57         off_t license;
  58
  59         // ASes in the database
  60         struct loc_database_as_v0* as_v0;
  61         size_t as_count;
  62
  63         // Network tree
  64         struct loc_database_network_node_v0* network_nodes_v0;
  65         size_t network_nodes_count;
  66
  67         // Networks
  68         struct loc_database_network_v0* networks_v0;
  69         size_t networks_count;
  70
  71         // Countries
  72         struct loc_database_country_v0* countries_v0;
  73         size_t countries_count;
  74
  75         struct loc_stringpool* pool;
  76 };
  77
  78 #define MAX_STACK_DEPTH 256
  79
  80 struct loc_node_stack {
  81         off_t offset;
  82         int i; // Is this node 0 or 1?
  83         int depth;
  84 };
  85
  86 struct loc_database_enumerator {
  87         struct loc_ctx* ctx;
  88         struct loc_database* db;
  89         enum loc_database_enumerator_mode mode;
  90         int refcount;
  91
  92         // Search string
  93         char* string;
  94         char country_code[3];
  95         uint32_t asn;
  96         enum loc_network_flags flags;
  97
  98         // Index of the AS we are looking at
  99         unsigned int as_index;
 100
 101         // Network state
 102         struct in6_addr network_address;
 103         struct loc_node_stack network_stack[MAX_STACK_DEPTH];
 104         int network_stack_depth;
 105         unsigned int* networks_visited;
 106 };
 107
 108 static int loc_database_read_magic(struct loc_database* db) {
 109         struct loc_database_magic magic;
 110
 111         // Read from file
 112         size_t bytes_read = fread(&magic, 1, sizeof(magic), db->f);
 113
 114         // Check if we have been able to read enough data
 115         if (bytes_read < sizeof(magic)) {
 116                 ERROR(db->ctx, "Could not read enough data to validate magic bytes\n");
 117                 DEBUG(db->ctx, "Read %zu bytes, but needed %zu\n", bytes_read, sizeof(magic));
 118                 return -ENOMSG;
 119         }
 120
 121         // Compare magic bytes
 122         if (memcmp(LOC_DATABASE_MAGIC, magic.magic, strlen(LOC_DATABASE_MAGIC)) == 0) {
 123                 DEBUG(db->ctx, "Magic value matches\n");
 124
 125                 // Parse version
 126                 db->version = be16toh(magic.version);
 127                 DEBUG(db->ctx, "Database version is %u\n", db->version);
 128
 129                 return 0;
 130         }
 131
 132         ERROR(db->ctx, "Database format is not compatible\n");
 133
 134         // Return an error
 135         return 1;
 136 }
 137
 138 static int loc_database_read_as_section_v0(struct loc_database* db,
 139                 const struct loc_database_header_v0* header) {
 140         off_t as_offset  = be32toh(header->as_offset);
 141         size_t as_length = be32toh(header->as_length);
 142
 143         DEBUG(db->ctx, "Reading AS section from %jd (%zu bytes)\n", (intmax_t)as_offset, as_length);
 144
 145         if (as_length > 0) {
 146                 db->as_v0 = mmap(NULL, as_length, PROT_READ,
 147                         MAP_SHARED, fileno(db->f), as_offset);
 148
 149                 if (db->as_v0 == MAP_FAILED)
 150                         return -errno;
 151         }
 152
 153         db->as_count = as_length / sizeof(*db->as_v0);
 154
 155         INFO(db->ctx, "Read %zu ASes from the database\n", db->as_count);
 156
 157         return 0;
 158 }
 159
 160 static int loc_database_read_network_nodes_section_v0(struct loc_database* db,
 161                 const struct loc_database_header_v0* header) {
 162         off_t network_nodes_offset  = be32toh(header->network_tree_offset);
 163         size_t network_nodes_length = be32toh(header->network_tree_length);
 164
 165         DEBUG(db->ctx, "Reading network nodes section from %jd (%zu bytes)\n",
 166                 (intmax_t)network_nodes_offset, network_nodes_length);
 167
 168         if (network_nodes_length > 0) {
 169                 db->network_nodes_v0 = mmap(NULL, network_nodes_length, PROT_READ,
 170                         MAP_SHARED, fileno(db->f), network_nodes_offset);
 171
 172                 if (db->network_nodes_v0 == MAP_FAILED)
 173                         return -errno;
 174         }
 175
 176         db->network_nodes_count = network_nodes_length / sizeof(*db->network_nodes_v0);
 177
 178         INFO(db->ctx, "Read %zu network nodes from the database\n", db->network_nodes_count);
 179
 180         return 0;
 181 }
 182
 183 static int loc_database_read_networks_section_v0(struct loc_database* db,
 184                 const struct loc_database_header_v0* header) {
 185         off_t networks_offset  = be32toh(header->network_data_offset);
 186         size_t networks_length = be32toh(header->network_data_length);
 187
 188         DEBUG(db->ctx, "Reading networks section from %jd (%zu bytes)\n",
 189                 (intmax_t)networks_offset, networks_length);
 190
 191         if (networks_length > 0) {
 192                 db->networks_v0 = mmap(NULL, networks_length, PROT_READ,
 193                         MAP_SHARED, fileno(db->f), networks_offset);
 194
 195                 if (db->networks_v0 == MAP_FAILED)
 196                         return -errno;
 197         }
 198
 199         db->networks_count = networks_length / sizeof(*db->networks_v0);
 200
 201         INFO(db->ctx, "Read %zu networks from the database\n", db->networks_count);
 202
 203         return 0;
 204 }
 205
 206 static int loc_database_read_countries_section_v0(struct loc_database* db,
 207                 const struct loc_database_header_v0* header) {
 208         off_t countries_offset  = be32toh(header->countries_offset);
 209         size_t countries_length = be32toh(header->countries_length);
 210
 211         DEBUG(db->ctx, "Reading countries section from %jd (%zu bytes)\n",
 212                 (intmax_t)countries_offset, countries_length);
 213
 214         if (countries_length > 0) {
 215                 db->countries_v0 = mmap(NULL, countries_length, PROT_READ,
 216                         MAP_SHARED, fileno(db->f), countries_offset);
 217
 218                 if (db->countries_v0 == MAP_FAILED)
 219                         return -errno;
 220         }
 221
 222         db->countries_count = countries_length / sizeof(*db->countries_v0);
 223
 224         INFO(db->ctx, "Read %zu countries from the database\n",
 225                 db->countries_count);
 226
 227         return 0;
 228 }
 229
 230 static int loc_database_read_header_v0(struct loc_database* db) {
 231         struct loc_database_header_v0 header;
 232
 233         // Read from file
 234         size_t size = fread(&header, 1, sizeof(header), db->f);
 235
 236         if (size < sizeof(header)) {
 237                 ERROR(db->ctx, "Could not read enough data for header\n");
 238                 return -ENOMSG;
 239         }
 240
 241         // Copy over data
 242         db->created_at  = be64toh(header.created_at);
 243         db->vendor      = be32toh(header.vendor);
 244         db->description = be32toh(header.description);
 245         db->license     = be32toh(header.license);
 246
 247         // Open pool
 248         off_t pool_offset  = be32toh(header.pool_offset);
 249         size_t pool_length = be32toh(header.pool_length);
 250
 251         int r = loc_stringpool_open(db->ctx, &db->pool,
 252                 db->f, pool_length, pool_offset);
 253         if (r)
 254                 return r;
 255
 256         // AS section
 257         r = loc_database_read_as_section_v0(db, &header);
 258         if (r)
 259                 return r;
 260
 261         // Network Nodes
 262         r = loc_database_read_network_nodes_section_v0(db, &header);
 263         if (r)
 264                 return r;
 265
 266         // Networks
 267         r = loc_database_read_networks_section_v0(db, &header);
 268         if (r)
 269                 return r;
 270
 271         // countries
 272         r = loc_database_read_countries_section_v0(db, &header);
 273         if (r)
 274                 return r;
 275
 276         return 0;
 277 }
 278
 279 static int loc_database_read_header(struct loc_database* db) {
 280         switch (db->version) {
 281                 case 0:
 282                         return loc_database_read_header_v0(db);
 283
 284                 default:
 285                         ERROR(db->ctx, "Incompatible database version: %u\n", db->version);
 286                         return 1;
 287         }
 288 }
 289
 290 static int loc_database_read(struct loc_database* db, FILE* f) {
 291         clock_t start = clock();
 292
 293         int fd = fileno(f);
 294
 295         // Clone file descriptor
 296         fd = dup(fd);
 297         if (!fd) {
 298                 ERROR(db->ctx, "Could not duplicate file descriptor\n");
 299                 return -1;
 300         }
 301
 302         // Reopen the file so that we can keep our own file handle
 303         db->f = fdopen(fd, "r");
 304         if (!db->f) {
 305                 ERROR(db->ctx, "Could not re-open database file\n");
 306                 return -1;
 307         }
 308
 309         // Rewind to the start of the file
 310         rewind(db->f);
 311
 312         // Read magic bytes
 313         int r = loc_database_read_magic(db);
 314         if (r)
 315                 return r;
 316
 317         // Read the header
 318         r = loc_database_read_header(db);
 319         if (r)
 320                 return r;
 321
 322         clock_t end = clock();
 323
 324         INFO(db->ctx, "Opened database in %.4fms\n",
 325                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 326
 327         return 0;
 328 }
 329
 330 LOC_EXPORT int loc_database_new(struct loc_ctx* ctx, struct loc_database** database, FILE* f) {
 331         // Fail on invalid file handle
 332         if (!f)
 333                 return -EINVAL;
 334
 335         struct loc_database* db = calloc(1, sizeof(*db));
 336         if (!db)
 337                 return -ENOMEM;
 338
 339         // Reference context
 340         db->ctx = loc_ref(ctx);
 341         db->refcount = 1;
 342
 343         DEBUG(db->ctx, "Database object allocated at %p\n", db);
 344
 345         int r = loc_database_read(db, f);
 346         if (r) {
 347                 loc_database_unref(db);
 348                 return r;
 349         }
 350
 351         *database = db;
 352
 353         return 0;
 354 }
 355
 356 LOC_EXPORT struct loc_database* loc_database_ref(struct loc_database* db) {
 357         db->refcount++;
 358
 359         return db;
 360 }
 361
 362 static void loc_database_free(struct loc_database* db) {
 363         int r;
 364
 365         DEBUG(db->ctx, "Releasing database %p\n", db);
 366
 367         // Removing all ASes
 368         if (db->as_v0) {
 369                 r = munmap(db->as_v0, db->as_count * sizeof(*db->as_v0));
 370                 if (r)
 371                         ERROR(db->ctx, "Could not unmap AS section: %s\n", strerror(errno));
 372         }
 373
 374         // Remove mapped network sections
 375         if (db->networks_v0) {
 376                 r = munmap(db->networks_v0, db->networks_count * sizeof(*db->networks_v0));
 377                 if (r)
 378                         ERROR(db->ctx, "Could not unmap networks section: %s\n", strerror(errno));
 379         }
 380
 381         // Remove mapped network nodes section
 382         if (db->network_nodes_v0) {
 383                 r = munmap(db->network_nodes_v0, db->network_nodes_count * sizeof(*db->network_nodes_v0));
 384                 if (r)
 385                         ERROR(db->ctx, "Could not unmap network nodes section: %s\n", strerror(errno));
 386         }
 387
 388         loc_stringpool_unref(db->pool);
 389
 390         // Close database file
 391         if (db->f)
 392                 fclose(db->f);
 393
 394         loc_unref(db->ctx);
 395         free(db);
 396 }
 397
 398 LOC_EXPORT struct loc_database* loc_database_unref(struct loc_database* db) {
 399         if (--db->refcount > 0)
 400                 return NULL;
 401
 402         loc_database_free(db);
 403         return NULL;
 404 }
 405
 406 static int loc_database_hash(struct loc_database* db,
 407                 unsigned char* hash, unsigned int* length) {
 408         int r = 0;
 409
 410         EVP_MD_CTX* mdctx = EVP_MD_CTX_new();
 411
 412         // Select SHA512
 413         const EVP_MD* md = EVP_sha512();
 414
 415         // Initialise hash function
 416         EVP_DigestInit_ex(mdctx, md, NULL);
 417
 418         // Reset file to start
 419         rewind(db->f);
 420
 421         // Read magic
 422         struct loc_database_magic magic;
 423         fread(&magic, 1, sizeof(magic), db->f);
 424
 425         // Feed magic into the hash
 426         EVP_DigestUpdate(mdctx, &magic, sizeof(magic));
 427
 428         // Read the header
 429         struct loc_database_header_v0 header_v0;
 430
 431         switch (db->version) {
 432                 case 0:
 433                         fread(&header_v0, 1, sizeof(header_v0), db->f);
 434
 435                         // Clear signature
 436                         for (unsigned int i = 0; i < sizeof(header_v0.signature); i++) {
 437                                 header_v0.signature[i] = '\0';
 438                         }
 439
 440                         // Feed header into the hash
 441                         EVP_DigestUpdate(mdctx, &header_v0, sizeof(header_v0));
 442                         break;
 443
 444                 default:
 445                         ERROR(db->ctx, "Cannot compute hash for database with format %d\n",
 446                                 db->version);
 447                         r = -EINVAL;
 448                         goto CLEANUP;
 449         }
 450
 451         // Walk through the file in chunks of 100kB
 452         char buffer[1024 * 100];
 453
 454         while (!feof(db->f)) {
 455                 size_t bytes_read = fread(buffer, 1, sizeof(buffer), db->f);
 456
 457                 EVP_DigestUpdate(mdctx, buffer, bytes_read);
 458         }
 459
 460         // Finish hash
 461         EVP_DigestFinal_ex(mdctx, hash, length);
 462
 463 #ifdef ENABLE_DEBUG
 464         char hash_string[(EVP_MAX_MD_SIZE * 2) + 1];
 465
 466         for (unsigned int i = 0; i < *length; i++) {
 467                 sprintf(&hash_string[i*2], "%02x", hash[i]);
 468         }
 469
 470         DEBUG(db->ctx, "Database hash: %s\n", hash_string);
 471 #endif
 472
 473 CLEANUP:
 474         // Cleanup
 475         EVP_MD_CTX_free(mdctx);
 476
 477         return r;
 478 }
 479
 480 LOC_EXPORT int loc_database_verify(struct loc_database* db) {
 481         // Hash
 482         unsigned char hash[EVP_MAX_MD_SIZE];
 483         unsigned int hash_length;
 484
 485         // Compute hash of the database
 486         int r = loc_database_hash(db, hash, &hash_length);
 487         if (r) {
 488                 ERROR(db->ctx, "Could not compute hash of the database\n");
 489                 return r;
 490         }
 491
 492         # warning TODO Check signature against hash
 493
 494         return 0;
 495 }
 496
 497 LOC_EXPORT time_t loc_database_created_at(struct loc_database* db) {
 498         return db->created_at;
 499 }
 500
 501 LOC_EXPORT const char* loc_database_get_vendor(struct loc_database* db) {
 502         return loc_stringpool_get(db->pool, db->vendor);
 503 }
 504
 505 LOC_EXPORT const char* loc_database_get_description(struct loc_database* db) {
 506         return loc_stringpool_get(db->pool, db->description);
 507 }
 508
 509 LOC_EXPORT const char* loc_database_get_license(struct loc_database* db) {
 510         return loc_stringpool_get(db->pool, db->license);
 511 }
 512
 513 LOC_EXPORT size_t loc_database_count_as(struct loc_database* db) {
 514         return db->as_count;
 515 }
 516
 517 // Returns the AS at position pos
 518 static int loc_database_fetch_as(struct loc_database* db, struct loc_as** as, off_t pos) {
 519         if ((size_t)pos >= db->as_count)
 520                 return -EINVAL;
 521
 522         DEBUG(db->ctx, "Fetching AS at position %jd\n", (intmax_t)pos);
 523
 524         int r;
 525         switch (db->version) {
 526                 case 0:
 527                         r = loc_as_new_from_database_v0(db->ctx, db->pool, as, db->as_v0 + pos);
 528                         break;
 529
 530                 default:
 531                         return -1;
 532         }
 533
 534         if (r == 0) {
 535                 DEBUG(db->ctx, "Got AS%u\n", loc_as_get_number(*as));
 536         }
 537
 538         return r;
 539 }
 540
 541 // Performs a binary search to find the AS in the list
 542 LOC_EXPORT int loc_database_get_as(struct loc_database* db, struct loc_as** as, uint32_t number) {
 543         off_t lo = 0;
 544         off_t hi = db->as_count - 1;
 545
 546         // Save start time
 547         clock_t start = clock();
 548
 549         while (lo <= hi) {
 550                 off_t i = (lo + hi) / 2;
 551
 552                 // Fetch AS in the middle between lo and hi
 553                 int r = loc_database_fetch_as(db, as, i);
 554                 if (r)
 555                         return r;
 556
 557                 // Check if this is a match
 558                 uint32_t as_number = loc_as_get_number(*as);
 559                 if (as_number == number) {
 560                         clock_t end = clock();
 561
 562                         // Log how fast this has been
 563                         DEBUG(db->ctx, "Found AS%u in %.4fms\n", as_number,
 564                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 565
 566                         return 0;
 567                 }
 568
 569                 // If it wasn't, we release the AS and
 570                 // adjust our search pointers
 571                 loc_as_unref(*as);
 572
 573                 if (as_number < number) {
 574                         lo = i + 1;
 575                 } else
 576                         hi = i - 1;
 577         }
 578
 579         // Nothing found
 580         *as = NULL;
 581
 582         return 1;
 583 }
 584
 585 // Returns the network at position pos
 586 static int loc_database_fetch_network(struct loc_database* db, struct loc_network** network,
 587                 struct in6_addr* address, unsigned int prefix, off_t pos) {
 588         if ((size_t)pos >= db->networks_count) {
 589                 DEBUG(db->ctx, "Network ID out of range: %jd/%jd\n",
 590                         (intmax_t)pos, (intmax_t)db->networks_count);
 591                 return -EINVAL;
 592         }
 593
 594
 595         DEBUG(db->ctx, "Fetching network at position %jd\n", (intmax_t)pos);
 596
 597         int r;
 598         switch (db->version) {
 599                 case 0:
 600                         r = loc_network_new_from_database_v0(db->ctx, network,
 601                                 address, prefix, db->networks_v0 + pos);
 602                         break;
 603
 604                 default:
 605                         return -1;
 606         }
 607
 608         if (r == 0) {
 609                 char* string = loc_network_str(*network);
 610                 DEBUG(db->ctx, "Got network %s\n", string);
 611                 free(string);
 612         }
 613
 614         return r;
 615 }
 616
 617 static int __loc_database_node_is_leaf(const struct loc_database_network_node_v0* node) {
 618         return (node->network != htobe32(0xffffffff));
 619 }
 620
 621 static int __loc_database_lookup_handle_leaf(struct loc_database* db, const struct in6_addr* address,
 622                 struct loc_network** network, struct in6_addr* network_address, unsigned int prefix,
 623                 const struct loc_database_network_node_v0* node) {
 624         off_t network_index = be32toh(node->network);
 625
 626         DEBUG(db->ctx, "Handling leaf node at %jd (%jd)\n", (intmax_t)(node - db->network_nodes_v0), (intmax_t)network_index);
 627
 628         // Fetch the network
 629         int r = loc_database_fetch_network(db, network,
 630                 network_address, prefix, network_index);
 631         if (r) {
 632                 ERROR(db->ctx, "Could not fetch network %jd from database\n", (intmax_t)network_index);
 633                 return r;
 634         }
 635
 636         // Check if the given IP address is inside the network
 637         r = loc_network_match_address(*network, address);
 638         if (r) {
 639                 DEBUG(db->ctx, "Searched address is not part of the network\n");
 640
 641                 loc_network_unref(*network);
 642                 *network = NULL;
 643                 return 1;
 644         }
 645
 646         // A network was found and the IP address matches
 647         return 0;
 648 }
 649
 650 // Searches for an exact match along the path
 651 static int __loc_database_lookup(struct loc_database* db, const struct in6_addr* address,
 652                 struct loc_network** network, struct in6_addr* network_address,
 653                 const struct loc_database_network_node_v0* node, unsigned int level) {
 654         int r;
 655         off_t node_index;
 656
 657         // Follow the path
 658         int bit = in6_addr_get_bit(address, level);
 659         in6_addr_set_bit(network_address, level, bit);
 660
 661         if (bit == 0)
 662                 node_index = be32toh(node->zero);
 663         else
 664                 node_index = be32toh(node->one);
 665
 666         // If the node index is zero, the tree ends here
 667         // and we cannot descend any further
 668         if (node_index > 0) {
 669                 // Check boundaries
 670                 if ((size_t)node_index >= db->network_nodes_count)
 671                         return -EINVAL;
 672
 673                 // Move on to the next node
 674                 r = __loc_database_lookup(db, address, network, network_address,
 675                         db->network_nodes_v0 + node_index, level + 1);
 676
 677                 // End here if a result was found
 678                 if (r == 0)
 679                         return r;
 680
 681                 // Raise any errors
 682                 else if (r < 0)
 683                         return r;
 684
 685                 DEBUG(db->ctx, "No match found below level %u\n", level);
 686         } else {
 687                 DEBUG(db->ctx, "Tree ended at level %u\n", level);
 688         }
 689
 690         // If this node has a leaf, we will check if it matches
 691         if (__loc_database_node_is_leaf(node)) {
 692                 r = __loc_database_lookup_handle_leaf(db, address, network, network_address, level, node);
 693                 if (r <= 0)
 694                         return r;
 695         }
 696
 697         return 1;
 698 }
 699
 700 LOC_EXPORT int loc_database_lookup(struct loc_database* db,
 701                 struct in6_addr* address, struct loc_network** network) {
 702         struct in6_addr network_address;
 703         memset(&network_address, 0, sizeof(network_address));
 704
 705         *network = NULL;
 706
 707         // Save start time
 708         clock_t start = clock();
 709
 710         int r = __loc_database_lookup(db, address, network, &network_address,
 711                 db->network_nodes_v0, 0);
 712
 713         clock_t end = clock();
 714
 715         // Log how fast this has been
 716         DEBUG(db->ctx, "Executed network search in %.4fms\n",
 717                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 718
 719         return r;
 720 }
 721
 722 LOC_EXPORT int loc_database_lookup_from_string(struct loc_database* db,
 723                 const char* string, struct loc_network** network) {
 724         struct in6_addr address;
 725
 726         int r = loc_parse_address(db->ctx, string, &address);
 727         if (r)
 728                 return r;
 729
 730         return loc_database_lookup(db, &address, network);
 731 }
 732
 733 // Returns the country at position pos
 734 static int loc_database_fetch_country(struct loc_database* db,
 735                 struct loc_country** country, off_t pos) {
 736         if ((size_t)pos >= db->countries_count)
 737                 return -EINVAL;
 738
 739         DEBUG(db->ctx, "Fetching country at position %jd\n", (intmax_t)pos);
 740
 741         int r;
 742         switch (db->version) {
 743                 case 0:
 744                         r = loc_country_new_from_database_v0(db->ctx, db->pool, country, db->countries_v0 + pos);
 745                         break;
 746
 747                 default:
 748                         return -1;
 749         }
 750
 751         if (r == 0) {
 752                 DEBUG(db->ctx, "Got country %s\n", loc_country_get_code(*country));
 753         }
 754
 755         return r;
 756 }
 757
 758 // Performs a binary search to find the country in the list
 759 LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 760                 struct loc_country** country, const char* code) {
 761         off_t lo = 0;
 762         off_t hi = db->countries_count - 1;
 763
 764         // Save start time
 765         clock_t start = clock();
 766
 767         while (lo <= hi) {
 768                 off_t i = (lo + hi) / 2;
 769
 770                 // Fetch country in the middle between lo and hi
 771                 int r = loc_database_fetch_country(db, country, i);
 772                 if (r)
 773                         return r;
 774
 775                 // Check if this is a match
 776                 const char* cc = loc_country_get_code(*country);
 777                 int result = strcmp(code, cc);
 778
 779                 if (result == 0) {
 780                         clock_t end = clock();
 781
 782                         // Log how fast this has been
 783                         DEBUG(db->ctx, "Found country %s in %.4fms\n", cc,
 784                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 785
 786                         return 0;
 787                 }
 788
 789                 // If it wasn't, we release the country and
 790                 // adjust our search pointers
 791                 loc_country_unref(*country);
 792
 793                 if (result > 0) {
 794                         lo = i + 1;
 795                 } else
 796                         hi = i - 1;
 797         }
 798
 799         // Nothing found
 800         *country = NULL;
 801
 802         return 1;
 803 }
 804
 805 // Enumerator
 806
 807 LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator,
 808                 struct loc_database* db, enum loc_database_enumerator_mode mode) {
 809         struct loc_database_enumerator* e = calloc(1, sizeof(*e));
 810         if (!e)
 811                 return -ENOMEM;
 812
 813         // Reference context
 814         e->ctx = loc_ref(db->ctx);
 815         e->db = loc_database_ref(db);
 816         e->mode = mode;
 817         e->refcount = 1;
 818
 819         // Initialise graph search
 820         //e->network_stack[++e->network_stack_depth] = 0;
 821         e->network_stack_depth = 1;
 822         e->networks_visited = calloc(db->network_nodes_count, sizeof(*e->networks_visited));
 823
 824         DEBUG(e->ctx, "Database enumerator object allocated at %p\n", e);
 825
 826         *enumerator = e;
 827         return 0;
 828 }
 829
 830 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_ref(struct loc_database_enumerator* enumerator) {
 831         enumerator->refcount++;
 832
 833         return enumerator;
 834 }
 835
 836 static void loc_database_enumerator_free(struct loc_database_enumerator* enumerator) {
 837         DEBUG(enumerator->ctx, "Releasing database enumerator %p\n", enumerator);
 838
 839         // Release all references
 840         loc_database_unref(enumerator->db);
 841         loc_unref(enumerator->ctx);
 842
 843         if (enumerator->string)
 844                 free(enumerator->string);
 845
 846         // Free network search
 847         free(enumerator->networks_visited);
 848
 849         free(enumerator);
 850 }
 851
 852 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_unref(struct loc_database_enumerator* enumerator) {
 853         if (!enumerator)
 854                 return NULL;
 855
 856         if (--enumerator->refcount > 0)
 857                 return enumerator;
 858
 859         loc_database_enumerator_free(enumerator);
 860         return NULL;
 861 }
 862
 863 LOC_EXPORT int loc_database_enumerator_set_string(struct loc_database_enumerator* enumerator, const char* string) {
 864         enumerator->string = strdup(string);
 865
 866         // Make the string lowercase
 867         for (char *p = enumerator->string; *p; p++)
 868                 *p = tolower(*p);
 869
 870         return 0;
 871 }
 872
 873 LOC_EXPORT int loc_database_enumerator_set_country_code(struct loc_database_enumerator* enumerator, const char* country_code) {
 874         // Set empty country code
 875         if (!country_code || !*country_code) {
 876                 *enumerator->country_code = '\0';
 877                 return 0;
 878         }
 879
 880         // Treat A1, A2, A3 as special country codes,
 881         // but perform search for flags instead
 882         if (strcmp(country_code, "A1") == 0) {
 883                 return loc_database_enumerator_set_flag(enumerator,
 884                         LOC_NETWORK_FLAG_ANONYMOUS_PROXY);
 885         } else if (strcmp(country_code, "A2") == 0) {
 886                 return loc_database_enumerator_set_flag(enumerator,
 887                         LOC_NETWORK_FLAG_SATELLITE_PROVIDER);
 888         } else if (strcmp(country_code, "A3") == 0) {
 889                 return loc_database_enumerator_set_flag(enumerator,
 890                         LOC_NETWORK_FLAG_ANYCAST);
 891         }
 892
 893         // Country codes must be two characters
 894         if (!loc_country_code_is_valid(country_code))
 895                 return -EINVAL;
 896
 897         for (unsigned int i = 0; i < 3; i++) {
 898                 enumerator->country_code[i] = country_code[i];
 899         }
 900
 901         return 0;
 902 }
 903
 904 LOC_EXPORT int loc_database_enumerator_set_asn(
 905                 struct loc_database_enumerator* enumerator, unsigned int asn) {
 906         enumerator->asn = asn;
 907
 908         return 0;
 909 }
 910
 911 LOC_EXPORT int loc_database_enumerator_set_flag(
 912                 struct loc_database_enumerator* enumerator, enum loc_network_flags flag) {
 913         enumerator->flags |= flag;
 914
 915         return 0;
 916 }
 917
 918 LOC_EXPORT int loc_database_enumerator_next_as(
 919                 struct loc_database_enumerator* enumerator, struct loc_as** as) {
 920         *as = NULL;
 921
 922         // Do not do anything if not in AS mode
 923         if (enumerator->mode != LOC_DB_ENUMERATE_ASES)
 924                 return 0;
 925
 926         struct loc_database* db = enumerator->db;
 927
 928         while (enumerator->as_index < db->as_count) {
 929                 // Fetch the next AS
 930                 int r = loc_database_fetch_as(db, as, enumerator->as_index++);
 931                 if (r)
 932                         return r;
 933
 934                 r = loc_as_match_string(*as, enumerator->string);
 935                 if (r == 1) {
 936                         DEBUG(enumerator->ctx, "AS%d (%s) matches %s\n",
 937                                 loc_as_get_number(*as), loc_as_get_name(*as), enumerator->string);
 938
 939                         return 0;
 940                 }
 941
 942                 // No match
 943                 loc_as_unref(*as);
 944                 *as = NULL;
 945         }
 946
 947         // Reset the index
 948         enumerator->as_index = 0;
 949
 950         // We have searched through all of them
 951         return 0;
 952 }
 953
 954 static int loc_database_enumerator_stack_push_node(
 955                 struct loc_database_enumerator* e, off_t offset, int i, int depth) {
 956         // Do not add empty nodes
 957         if (!offset)
 958                 return 0;
 959
 960         // Check if there is any space left on the stack
 961         if (e->network_stack_depth >= MAX_STACK_DEPTH) {
 962                 ERROR(e->ctx, "Maximum stack size reached: %d\n", e->network_stack_depth);
 963                 return -1;
 964         }
 965
 966         // Increase stack size
 967         int s = ++e->network_stack_depth;
 968
 969         DEBUG(e->ctx, "Added node %jd to stack (%d)\n", (intmax_t)offset, depth);
 970
 971         e->network_stack[s].offset = offset;
 972         e->network_stack[s].i = i;
 973         e->network_stack[s].depth = depth;
 974
 975         return 0;
 976 }
 977
 978 LOC_EXPORT int loc_database_enumerator_next_network(
 979                 struct loc_database_enumerator* enumerator, struct loc_network** network) {
 980         // Reset network
 981         *network = NULL;
 982
 983         // Do not do anything if not in network mode
 984         if (enumerator->mode != LOC_DB_ENUMERATE_NETWORKS)
 985                 return 0;
 986
 987         int r;
 988
 989         DEBUG(enumerator->ctx, "Called with a stack of %u nodes\n",
 990                 enumerator->network_stack_depth);
 991
 992         // Perform DFS
 993         while (enumerator->network_stack_depth > 0) {
 994                 DEBUG(enumerator->ctx, "Stack depth: %u\n", enumerator->network_stack_depth);
 995
 996                 // Get object from top of the stack
 997                 struct loc_node_stack* node = &enumerator->network_stack[enumerator->network_stack_depth];
 998
 999                 // Remove the node from the stack if we have already visited it
1000                 if (enumerator->networks_visited[node->offset]) {
1001                         enumerator->network_stack_depth--;
1002                         continue;
1003                 }
1004
1005                 // Mark the bits on the path correctly
1006                 in6_addr_set_bit(&enumerator->network_address,
1007                         (node->depth > 0) ? node->depth - 1 : 0, node->i);
1008
1009                 DEBUG(enumerator->ctx, "Looking at node %jd\n", (intmax_t)node->offset);
1010                 enumerator->networks_visited[node->offset]++;
1011
1012                 // Pop node from top of the stack
1013                 struct loc_database_network_node_v0* n =
1014                         enumerator->db->network_nodes_v0 + node->offset;
1015
1016                 // Add edges to stack
1017                 r = loc_database_enumerator_stack_push_node(enumerator,
1018                         be32toh(n->one), 1, node->depth + 1);
1019
1020                 if (r)
1021                         return r;
1022
1023                 r = loc_database_enumerator_stack_push_node(enumerator,
1024                         be32toh(n->zero), 0, node->depth + 1);
1025
1026                 if (r)
1027                         return r;
1028
1029                 // Check if this node is a leaf and has a network object
1030                 if (__loc_database_node_is_leaf(n)) {
1031                         off_t network_index = be32toh(n->network);
1032
1033                         DEBUG(enumerator->ctx, "Node has a network at %jd\n", (intmax_t)network_index);
1034
1035                         // Fetch the network object
1036                         r = loc_database_fetch_network(enumerator->db, network,
1037                                 &enumerator->network_address, node->depth, network_index);
1038
1039                         // Break on any errors
1040                         if (r)
1041                                 return r;
1042
1043                         // Check if we are interested in this network
1044
1045                         // Skip if the country code does not match
1046                         if (*enumerator->country_code &&
1047                                         !loc_network_match_country_code(*network, enumerator->country_code)) {
1048                                 loc_network_unref(*network);
1049                                 *network = NULL;
1050
1051                                 continue;
1052                         }
1053
1054                         // Skip if the ASN does not match
1055                         if (enumerator->asn &&
1056                                         !loc_network_match_asn(*network, enumerator->asn)) {
1057                                 loc_network_unref(*network);
1058                                 *network = NULL;
1059
1060                                 continue;
1061                         }
1062
1063                         // Skip if flags do not match
1064                         if (enumerator->flags &&
1065                                         !loc_network_match_flag(*network, enumerator->flags)) {
1066                                 loc_network_unref(*network);
1067                                 *network = NULL;
1068                         }
1069
1070                         return 0;
1071                 }
1072         }
1073
1074         // Reached the end of the search
1075
1076         // Mark all nodes as non-visited
1077         for (unsigned int i = 0; i < enumerator->db->network_nodes_count; i++)
1078                 enumerator->networks_visited[i] = 0;
1079
1080         return 0;
1081 }