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
 489         switch (db->version) {
 490                 case LOC_DATABASE_VERSION_1:
 491                         fread(&header_v1, 1, sizeof(header_v1), db->f);
 492
 493                         // Clear signature
 494                         for (unsigned int i = 0; i < sizeof(header_v1.signature); i++) {
 495                                 header_v1.signature[i] = '\0';
 496                         }
 497
 498                         hexdump(db->ctx, &header_v1, sizeof(header_v1));
 499
 500                         // Feed header into the hash
 501                         r = EVP_DigestVerifyUpdate(mdctx, &header_v1, sizeof(header_v1));
 502                         if (r != 1) {
 503                                 ERROR(db->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
 504                                 r = 1;
 505
 506                                 goto CLEANUP;
 507                         }
 508                         break;
 509
 510                 default:
 511                         ERROR(db->ctx, "Cannot compute hash for database with format %d\n",
 512                                 db->version);
 513                         r = -EINVAL;
 514                         goto CLEANUP;
 515         }
 516
 517         // Walk through the file in chunks of 64kB
 518         char buffer[64 * 1024];
 519
 520         while (!feof(db->f)) {
 521                 size_t bytes_read = fread(buffer, 1, sizeof(buffer), db->f);
 522
 523                 hexdump(db->ctx, buffer, bytes_read);
 524
 525                 r = EVP_DigestVerifyUpdate(mdctx, buffer, bytes_read);
 526                 if (r != 1) {
 527                         ERROR(db->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
 528                         r = 1;
 529
 530                         goto CLEANUP;
 531                 }
 532         }
 533
 534         // Finish
 535         r = EVP_DigestVerifyFinal(mdctx,
 536                 (unsigned char*)db->signature, db->signature_length);
 537
 538         if (r == 0) {
 539                 DEBUG(db->ctx, "The signature is invalid\n");
 540                 r = 1;
 541         } else if (r == 1) {
 542                 DEBUG(db->ctx, "The signature is valid\n");
 543                 r = 0;
 544         } else {
 545                 ERROR(db->ctx, "Error verifying the signature: %s\n",
 546                         ERR_error_string(ERR_get_error(), NULL));
 547                 r = 1;
 548         }
 549
 550         // Dump signature
 551         hexdump(db->ctx, db->signature, db->signature_length);
 552
 553         clock_t end = clock();
 554         DEBUG(db->ctx, "Signature checked in %.4fms\n",
 555                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 556
 557 CLEANUP:
 558         // Cleanup
 559         EVP_MD_CTX_free(mdctx);
 560         EVP_PKEY_free(pkey);
 561
 562         return r;
 563 }
 564
 565 LOC_EXPORT time_t loc_database_created_at(struct loc_database* db) {
 566         return db->created_at;
 567 }
 568
 569 LOC_EXPORT const char* loc_database_get_vendor(struct loc_database* db) {
 570         return loc_stringpool_get(db->pool, db->vendor);
 571 }
 572
 573 LOC_EXPORT const char* loc_database_get_description(struct loc_database* db) {
 574         return loc_stringpool_get(db->pool, db->description);
 575 }
 576
 577 LOC_EXPORT const char* loc_database_get_license(struct loc_database* db) {
 578         return loc_stringpool_get(db->pool, db->license);
 579 }
 580
 581 LOC_EXPORT size_t loc_database_count_as(struct loc_database* db) {
 582         return db->as_count;
 583 }
 584
 585 // Returns the AS at position pos
 586 static int loc_database_fetch_as(struct loc_database* db, struct loc_as** as, off_t pos) {
 587         if ((size_t)pos >= db->as_count)
 588                 return -EINVAL;
 589
 590         DEBUG(db->ctx, "Fetching AS at position %jd\n", (intmax_t)pos);
 591
 592         int r;
 593         switch (db->version) {
 594                 case LOC_DATABASE_VERSION_1:
 595                         r = loc_as_new_from_database_v1(db->ctx, db->pool, as, db->as_v1 + pos);
 596                         break;
 597
 598                 default:
 599                         return -1;
 600         }
 601
 602         if (r == 0) {
 603                 DEBUG(db->ctx, "Got AS%u\n", loc_as_get_number(*as));
 604         }
 605
 606         return r;
 607 }
 608
 609 // Performs a binary search to find the AS in the list
 610 LOC_EXPORT int loc_database_get_as(struct loc_database* db, struct loc_as** as, uint32_t number) {
 611         off_t lo = 0;
 612         off_t hi = db->as_count - 1;
 613
 614         // Save start time
 615         clock_t start = clock();
 616
 617         while (lo <= hi) {
 618                 off_t i = (lo + hi) / 2;
 619
 620                 // Fetch AS in the middle between lo and hi
 621                 int r = loc_database_fetch_as(db, as, i);
 622                 if (r)
 623                         return r;
 624
 625                 // Check if this is a match
 626                 uint32_t as_number = loc_as_get_number(*as);
 627                 if (as_number == number) {
 628                         clock_t end = clock();
 629
 630                         // Log how fast this has been
 631                         DEBUG(db->ctx, "Found AS%u in %.4fms\n", as_number,
 632                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 633
 634                         return 0;
 635                 }
 636
 637                 // If it wasn't, we release the AS and
 638                 // adjust our search pointers
 639                 loc_as_unref(*as);
 640
 641                 if (as_number < number) {
 642                         lo = i + 1;
 643                 } else
 644                         hi = i - 1;
 645         }
 646
 647         // Nothing found
 648         *as = NULL;
 649
 650         return 1;
 651 }
 652
 653 // Returns the network at position pos
 654 static int loc_database_fetch_network(struct loc_database* db, struct loc_network** network,
 655                 struct in6_addr* address, unsigned int prefix, off_t pos) {
 656         if ((size_t)pos >= db->networks_count) {
 657                 DEBUG(db->ctx, "Network ID out of range: %jd/%jd\n",
 658                         (intmax_t)pos, (intmax_t)db->networks_count);
 659                 return -EINVAL;
 660         }
 661
 662
 663         DEBUG(db->ctx, "Fetching network at position %jd\n", (intmax_t)pos);
 664
 665         int r;
 666         switch (db->version) {
 667                 case LOC_DATABASE_VERSION_1:
 668                         r = loc_network_new_from_database_v1(db->ctx, network,
 669                                 address, prefix, db->networks_v1 + pos);
 670                         break;
 671
 672                 default:
 673                         return -1;
 674         }
 675
 676         if (r == 0) {
 677                 char* string = loc_network_str(*network);
 678                 DEBUG(db->ctx, "Got network %s\n", string);
 679                 free(string);
 680         }
 681
 682         return r;
 683 }
 684
 685 static int __loc_database_node_is_leaf(const struct loc_database_network_node_v1* node) {
 686         return (node->network != htobe32(0xffffffff));
 687 }
 688
 689 static int __loc_database_lookup_handle_leaf(struct loc_database* db, const struct in6_addr* address,
 690                 struct loc_network** network, struct in6_addr* network_address, unsigned int prefix,
 691                 const struct loc_database_network_node_v1* node) {
 692         off_t network_index = be32toh(node->network);
 693
 694         DEBUG(db->ctx, "Handling leaf node at %jd (%jd)\n", (intmax_t)(node - db->network_nodes_v1), (intmax_t)network_index);
 695
 696         // Fetch the network
 697         int r = loc_database_fetch_network(db, network,
 698                 network_address, prefix, network_index);
 699         if (r) {
 700                 ERROR(db->ctx, "Could not fetch network %jd from database\n", (intmax_t)network_index);
 701                 return r;
 702         }
 703
 704         // Check if the given IP address is inside the network
 705         r = loc_network_match_address(*network, address);
 706         if (r) {
 707                 DEBUG(db->ctx, "Searched address is not part of the network\n");
 708
 709                 loc_network_unref(*network);
 710                 *network = NULL;
 711                 return 1;
 712         }
 713
 714         // A network was found and the IP address matches
 715         return 0;
 716 }
 717
 718 // Searches for an exact match along the path
 719 static int __loc_database_lookup(struct loc_database* db, const struct in6_addr* address,
 720                 struct loc_network** network, struct in6_addr* network_address,
 721                 const struct loc_database_network_node_v1* node, unsigned int level) {
 722         int r;
 723         off_t node_index;
 724
 725         // Follow the path
 726         int bit = in6_addr_get_bit(address, level);
 727         in6_addr_set_bit(network_address, level, bit);
 728
 729         if (bit == 0)
 730                 node_index = be32toh(node->zero);
 731         else
 732                 node_index = be32toh(node->one);
 733
 734         // If the node index is zero, the tree ends here
 735         // and we cannot descend any further
 736         if (node_index > 0) {
 737                 // Check boundaries
 738                 if ((size_t)node_index >= db->network_nodes_count)
 739                         return -EINVAL;
 740
 741                 // Move on to the next node
 742                 r = __loc_database_lookup(db, address, network, network_address,
 743                         db->network_nodes_v1 + node_index, level + 1);
 744
 745                 // End here if a result was found
 746                 if (r == 0)
 747                         return r;
 748
 749                 // Raise any errors
 750                 else if (r < 0)
 751                         return r;
 752
 753                 DEBUG(db->ctx, "No match found below level %u\n", level);
 754         } else {
 755                 DEBUG(db->ctx, "Tree ended at level %u\n", level);
 756         }
 757
 758         // If this node has a leaf, we will check if it matches
 759         if (__loc_database_node_is_leaf(node)) {
 760                 r = __loc_database_lookup_handle_leaf(db, address, network, network_address, level, node);
 761                 if (r <= 0)
 762                         return r;
 763         }
 764
 765         return 1;
 766 }
 767
 768 LOC_EXPORT int loc_database_lookup(struct loc_database* db,
 769                 struct in6_addr* address, struct loc_network** network) {
 770         struct in6_addr network_address;
 771         memset(&network_address, 0, sizeof(network_address));
 772
 773         *network = NULL;
 774
 775         // Save start time
 776         clock_t start = clock();
 777
 778         int r = __loc_database_lookup(db, address, network, &network_address,
 779                 db->network_nodes_v1, 0);
 780
 781         clock_t end = clock();
 782
 783         // Log how fast this has been
 784         DEBUG(db->ctx, "Executed network search in %.4fms\n",
 785                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 786
 787         return r;
 788 }
 789
 790 LOC_EXPORT int loc_database_lookup_from_string(struct loc_database* db,
 791                 const char* string, struct loc_network** network) {
 792         struct in6_addr address;
 793
 794         int r = loc_parse_address(db->ctx, string, &address);
 795         if (r)
 796                 return r;
 797
 798         return loc_database_lookup(db, &address, network);
 799 }
 800
 801 // Returns the country at position pos
 802 static int loc_database_fetch_country(struct loc_database* db,
 803                 struct loc_country** country, off_t pos) {
 804         if ((size_t)pos >= db->countries_count)
 805                 return -EINVAL;
 806
 807         DEBUG(db->ctx, "Fetching country at position %jd\n", (intmax_t)pos);
 808
 809         int r;
 810         switch (db->version) {
 811                 case LOC_DATABASE_VERSION_1:
 812                         r = loc_country_new_from_database_v1(db->ctx, db->pool, country, db->countries_v1 + pos);
 813                         break;
 814
 815                 default:
 816                         return -1;
 817         }
 818
 819         if (r == 0) {
 820                 DEBUG(db->ctx, "Got country %s\n", loc_country_get_code(*country));
 821         }
 822
 823         return r;
 824 }
 825
 826 // Performs a binary search to find the country in the list
 827 LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 828                 struct loc_country** country, const char* code) {
 829         off_t lo = 0;
 830         off_t hi = db->countries_count - 1;
 831
 832         // Save start time
 833         clock_t start = clock();
 834
 835         while (lo <= hi) {
 836                 off_t i = (lo + hi) / 2;
 837
 838                 // Fetch country in the middle between lo and hi
 839                 int r = loc_database_fetch_country(db, country, i);
 840                 if (r)
 841                         return r;
 842
 843                 // Check if this is a match
 844                 const char* cc = loc_country_get_code(*country);
 845                 int result = strcmp(code, cc);
 846
 847                 if (result == 0) {
 848                         clock_t end = clock();
 849
 850                         // Log how fast this has been
 851                         DEBUG(db->ctx, "Found country %s in %.4fms\n", cc,
 852                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 853
 854                         return 0;
 855                 }
 856
 857                 // If it wasn't, we release the country and
 858                 // adjust our search pointers
 859                 loc_country_unref(*country);
 860
 861                 if (result > 0) {
 862                         lo = i + 1;
 863                 } else
 864                         hi = i - 1;
 865         }
 866
 867         // Nothing found
 868         *country = NULL;
 869
 870         return 1;
 871 }
 872
 873 // Enumerator
 874
 875 LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator,
 876                 struct loc_database* db, enum loc_database_enumerator_mode mode) {
 877         struct loc_database_enumerator* e = calloc(1, sizeof(*e));
 878         if (!e)
 879                 return -ENOMEM;
 880
 881         // Reference context
 882         e->ctx = loc_ref(db->ctx);
 883         e->db = loc_database_ref(db);
 884         e->mode = mode;
 885         e->refcount = 1;
 886
 887         // Initialise graph search
 888         //e->network_stack[++e->network_stack_depth] = 0;
 889         e->network_stack_depth = 1;
 890         e->networks_visited = calloc(db->network_nodes_count, sizeof(*e->networks_visited));
 891
 892         DEBUG(e->ctx, "Database enumerator object allocated at %p\n", e);
 893
 894         *enumerator = e;
 895         return 0;
 896 }
 897
 898 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_ref(struct loc_database_enumerator* enumerator) {
 899         enumerator->refcount++;
 900
 901         return enumerator;
 902 }
 903
 904 static void loc_database_enumerator_free(struct loc_database_enumerator* enumerator) {
 905         DEBUG(enumerator->ctx, "Releasing database enumerator %p\n", enumerator);
 906
 907         // Release all references
 908         loc_database_unref(enumerator->db);
 909         loc_unref(enumerator->ctx);
 910
 911         if (enumerator->string)
 912                 free(enumerator->string);
 913
 914         // Free network search
 915         free(enumerator->networks_visited);
 916
 917         free(enumerator);
 918 }
 919
 920 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_unref(struct loc_database_enumerator* enumerator) {
 921         if (!enumerator)
 922                 return NULL;
 923
 924         if (--enumerator->refcount > 0)
 925                 return enumerator;
 926
 927         loc_database_enumerator_free(enumerator);
 928         return NULL;
 929 }
 930
 931 LOC_EXPORT int loc_database_enumerator_set_string(struct loc_database_enumerator* enumerator, const char* string) {
 932         enumerator->string = strdup(string);
 933
 934         // Make the string lowercase
 935         for (char *p = enumerator->string; *p; p++)
 936                 *p = tolower(*p);
 937
 938         return 0;
 939 }
 940
 941 LOC_EXPORT int loc_database_enumerator_set_country_code(struct loc_database_enumerator* enumerator, const char* country_code) {
 942         // Set empty country code
 943         if (!country_code || !*country_code) {
 944                 *enumerator->country_code = '\0';
 945                 return 0;
 946         }
 947
 948         // Treat A1, A2, A3 as special country codes,
 949         // but perform search for flags instead
 950         if (strcmp(country_code, "A1") == 0) {
 951                 return loc_database_enumerator_set_flag(enumerator,
 952                         LOC_NETWORK_FLAG_ANONYMOUS_PROXY);
 953         } else if (strcmp(country_code, "A2") == 0) {
 954                 return loc_database_enumerator_set_flag(enumerator,
 955                         LOC_NETWORK_FLAG_SATELLITE_PROVIDER);
 956         } else if (strcmp(country_code, "A3") == 0) {
 957                 return loc_database_enumerator_set_flag(enumerator,
 958                         LOC_NETWORK_FLAG_ANYCAST);
 959         }
 960
 961         // Country codes must be two characters
 962         if (!loc_country_code_is_valid(country_code))
 963                 return -EINVAL;
 964
 965         for (unsigned int i = 0; i < 3; i++) {
 966                 enumerator->country_code[i] = country_code[i];
 967         }
 968
 969         return 0;
 970 }
 971
 972 LOC_EXPORT int loc_database_enumerator_set_asn(
 973                 struct loc_database_enumerator* enumerator, unsigned int asn) {
 974         enumerator->asn = asn;
 975
 976         return 0;
 977 }
 978
 979 LOC_EXPORT int loc_database_enumerator_set_flag(
 980                 struct loc_database_enumerator* enumerator, enum loc_network_flags flag) {
 981         enumerator->flags |= flag;
 982
 983         return 0;
 984 }
 985
 986 LOC_EXPORT int loc_database_enumerator_set_family(
 987                 struct loc_database_enumerator* enumerator, int family) {
 988         enumerator->family = family;
 989
 990         return 0;
 991 }
 992
 993 LOC_EXPORT int loc_database_enumerator_next_as(
 994                 struct loc_database_enumerator* enumerator, struct loc_as** as) {
 995         *as = NULL;
 996
 997         // Do not do anything if not in AS mode
 998         if (enumerator->mode != LOC_DB_ENUMERATE_ASES)
 999                 return 0;
1000
1001         struct loc_database* db = enumerator->db;
1002
1003         while (enumerator->as_index < db->as_count) {
1004                 // Fetch the next AS
1005                 int r = loc_database_fetch_as(db, as, enumerator->as_index++);
1006                 if (r)
1007                         return r;
1008
1009                 r = loc_as_match_string(*as, enumerator->string);
1010                 if (r == 1) {
1011                         DEBUG(enumerator->ctx, "AS%d (%s) matches %s\n",
1012                                 loc_as_get_number(*as), loc_as_get_name(*as), enumerator->string);
1013
1014                         return 0;
1015                 }
1016
1017                 // No match
1018                 loc_as_unref(*as);
1019                 *as = NULL;
1020         }
1021
1022         // Reset the index
1023         enumerator->as_index = 0;
1024
1025         // We have searched through all of them
1026         return 0;
1027 }
1028
1029 static int loc_database_enumerator_stack_push_node(
1030                 struct loc_database_enumerator* e, off_t offset, int i, int depth) {
1031         // Do not add empty nodes
1032         if (!offset)
1033                 return 0;
1034
1035         // Check if there is any space left on the stack
1036         if (e->network_stack_depth >= MAX_STACK_DEPTH) {
1037                 ERROR(e->ctx, "Maximum stack size reached: %d\n", e->network_stack_depth);
1038                 return -1;
1039         }
1040
1041         // Increase stack size
1042         int s = ++e->network_stack_depth;
1043
1044         DEBUG(e->ctx, "Added node %jd to stack (%d)\n", (intmax_t)offset, depth);
1045
1046         e->network_stack[s].offset = offset;
1047         e->network_stack[s].i = i;
1048         e->network_stack[s].depth = depth;
1049
1050         return 0;
1051 }
1052
1053 LOC_EXPORT int loc_database_enumerator_next_network(
1054                 struct loc_database_enumerator* enumerator, struct loc_network** network) {
1055         // Reset network
1056         *network = NULL;
1057
1058         // Do not do anything if not in network mode
1059         if (enumerator->mode != LOC_DB_ENUMERATE_NETWORKS)
1060                 return 0;
1061
1062         int r;
1063
1064         DEBUG(enumerator->ctx, "Called with a stack of %u nodes\n",
1065                 enumerator->network_stack_depth);
1066
1067         // Perform DFS
1068         while (enumerator->network_stack_depth > 0) {
1069                 DEBUG(enumerator->ctx, "Stack depth: %u\n", enumerator->network_stack_depth);
1070
1071                 // Get object from top of the stack
1072                 struct loc_node_stack* node = &enumerator->network_stack[enumerator->network_stack_depth];
1073
1074                 // Remove the node from the stack if we have already visited it
1075                 if (enumerator->networks_visited[node->offset]) {
1076                         enumerator->network_stack_depth--;
1077                         continue;
1078                 }
1079
1080                 // Mark the bits on the path correctly
1081                 in6_addr_set_bit(&enumerator->network_address,
1082                         (node->depth > 0) ? node->depth - 1 : 0, node->i);
1083
1084                 DEBUG(enumerator->ctx, "Looking at node %jd\n", (intmax_t)node->offset);
1085                 enumerator->networks_visited[node->offset]++;
1086
1087                 // Pop node from top of the stack
1088                 struct loc_database_network_node_v1* n =
1089                         enumerator->db->network_nodes_v1 + node->offset;
1090
1091                 // Add edges to stack
1092                 r = loc_database_enumerator_stack_push_node(enumerator,
1093                         be32toh(n->one), 1, node->depth + 1);
1094
1095                 if (r)
1096                         return r;
1097
1098                 r = loc_database_enumerator_stack_push_node(enumerator,
1099                         be32toh(n->zero), 0, node->depth + 1);
1100
1101                 if (r)
1102                         return r;
1103
1104                 // Check if this node is a leaf and has a network object
1105                 if (__loc_database_node_is_leaf(n)) {
1106                         off_t network_index = be32toh(n->network);
1107
1108                         DEBUG(enumerator->ctx, "Node has a network at %jd\n", (intmax_t)network_index);
1109
1110                         // Fetch the network object
1111                         r = loc_database_fetch_network(enumerator->db, network,
1112                                 &enumerator->network_address, node->depth, network_index);
1113
1114                         // Break on any errors
1115                         if (r)
1116                                 return r;
1117
1118                         // Check if we are interested in this network
1119
1120                         // Skip if the family does not match
1121                         if (enumerator->family && loc_network_address_family(*network) != enumerator->family) {
1122                                 loc_network_unref(*network);
1123                                 *network = NULL;
1124
1125                                 continue;
1126                         }
1127
1128                         // Skip if the country code does not match
1129                         if (*enumerator->country_code &&
1130                                         !loc_network_match_country_code(*network, enumerator->country_code)) {
1131                                 loc_network_unref(*network);
1132                                 *network = NULL;
1133
1134                                 continue;
1135                         }
1136
1137                         // Skip if the ASN does not match
1138                         if (enumerator->asn &&
1139                                         !loc_network_match_asn(*network, enumerator->asn)) {
1140                                 loc_network_unref(*network);
1141                                 *network = NULL;
1142
1143                                 continue;
1144                         }
1145
1146                         // Skip if flags do not match
1147                         if (enumerator->flags &&
1148                                         !loc_network_match_flag(*network, enumerator->flags)) {
1149                                 loc_network_unref(*network);
1150                                 *network = NULL;
1151                         }
1152
1153                         return 0;
1154                 }
1155         }
1156
1157         // Reached the end of the search
1158
1159         // Mark all nodes as non-visited
1160         for (unsigned int i = 0; i < enumerator->db->network_nodes_count; i++)
1161                 enumerator->networks_visited[i] = 0;
1162
1163         return 0;
1164 }