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