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