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