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