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