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