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