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: %zu\n", length);
 254                 return -EINVAL;
 255         }
 256
 257         DEBUG(db->ctx, "Reading signature of %zu 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         // Remove mapped countries section
 450         if (db->countries_v1) {
 451                 r = munmap(db->countries_v1, db->countries_count * sizeof(*db->countries_v1));
 452                 if (r)
 453                         ERROR(db->ctx, "Could not unmap countries section: %s\n", strerror(errno));
 454         }
 455
 456         if (db->pool)
 457                 loc_stringpool_unref(db->pool);
 458
 459         // Free signature
 460         if (db->signature1)
 461                 free(db->signature1);
 462         if (db->signature2)
 463                 free(db->signature2);
 464
 465         // Close database file
 466         if (db->f)
 467                 fclose(db->f);
 468
 469         loc_unref(db->ctx);
 470         free(db);
 471 }
 472
 473 LOC_EXPORT struct loc_database* loc_database_unref(struct loc_database* db) {
 474         if (--db->refcount > 0)
 475                 return NULL;
 476
 477         loc_database_free(db);
 478         return NULL;
 479 }
 480
 481 LOC_EXPORT int loc_database_verify(struct loc_database* db, FILE* f) {
 482         // Cannot do this when no signature is available
 483         if (!db->signature1 && !db->signature2) {
 484                 DEBUG(db->ctx, "No signature available to verify\n");
 485                 return 1;
 486         }
 487
 488         // Start the stopwatch
 489         clock_t start = clock();
 490
 491         // Load public key
 492         EVP_PKEY* pkey = PEM_read_PUBKEY(f, NULL, NULL, NULL);
 493         if (!pkey) {
 494                 char* error = ERR_error_string(ERR_get_error(), NULL);
 495                 ERROR(db->ctx, "Could not parse public key: %s\n", error);
 496
 497                 return -1;
 498         }
 499
 500         int r = 0;
 501
 502         EVP_MD_CTX* mdctx = EVP_MD_CTX_new();
 503
 504         // Initialise hash function
 505         r = EVP_DigestVerifyInit(mdctx, NULL, NULL, NULL, pkey);
 506         if (r != 1) {
 507                 ERROR(db->ctx, "Error initializing signature validation: %s\n",
 508                         ERR_error_string(ERR_get_error(), NULL));
 509                 r = 1;
 510
 511                 goto CLEANUP;
 512         }
 513
 514         // Reset file to start
 515         rewind(db->f);
 516
 517         // Read magic
 518         struct loc_database_magic magic;
 519         fread(&magic, 1, sizeof(magic), db->f);
 520
 521         hexdump(db->ctx, &magic, sizeof(magic));
 522
 523         // Feed magic into the hash
 524         r = EVP_DigestVerifyUpdate(mdctx, &magic, sizeof(magic));
 525         if (r != 1) {
 526                 ERROR(db->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
 527                 r = 1;
 528
 529                 goto CLEANUP;
 530         }
 531
 532         // Read the header
 533         struct loc_database_header_v1 header_v1;
 534         size_t bytes_read;
 535
 536         switch (db->version) {
 537                 case LOC_DATABASE_VERSION_1:
 538                         bytes_read = fread(&header_v1, 1, sizeof(header_v1), db->f);
 539                         if (bytes_read < sizeof(header_v1)) {
 540                                 ERROR(db->ctx, "Could not read header\n");
 541                                 r = 1;
 542
 543                                 goto CLEANUP;
 544                         }
 545
 546                         // Clear signatures
 547                         memset(header_v1.signature1, '\0', sizeof(header_v1.signature1));
 548                         header_v1.signature1_length = 0;
 549                         memset(header_v1.signature2, '\0', sizeof(header_v1.signature2));
 550                         header_v1.signature2_length = 0;
 551
 552                         hexdump(db->ctx, &header_v1, sizeof(header_v1));
 553
 554                         // Feed header into the hash
 555                         r = EVP_DigestVerifyUpdate(mdctx, &header_v1, sizeof(header_v1));
 556                         if (r != 1) {
 557                                 ERROR(db->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
 558                                 r = 1;
 559
 560                                 goto CLEANUP;
 561                         }
 562                         break;
 563
 564                 default:
 565                         ERROR(db->ctx, "Cannot compute hash for database with format %d\n",
 566                                 db->version);
 567                         r = -EINVAL;
 568                         goto CLEANUP;
 569         }
 570
 571         // Walk through the file in chunks of 64kB
 572         char buffer[64 * 1024];
 573
 574         while (!feof(db->f)) {
 575                 bytes_read = fread(buffer, 1, sizeof(buffer), db->f);
 576
 577                 hexdump(db->ctx, buffer, bytes_read);
 578
 579                 r = EVP_DigestVerifyUpdate(mdctx, buffer, bytes_read);
 580                 if (r != 1) {
 581                         ERROR(db->ctx, "%s\n", ERR_error_string(ERR_get_error(), NULL));
 582                         r = 1;
 583
 584                         goto CLEANUP;
 585                 }
 586         }
 587
 588         // Check first signature
 589         if (db->signature1) {
 590                 hexdump(db->ctx, db->signature1, db->signature1_length);
 591
 592                 r = EVP_DigestVerifyFinal(mdctx,
 593                         (unsigned char*)db->signature1, db->signature1_length);
 594
 595                 if (r == 0) {
 596                         DEBUG(db->ctx, "The first signature is invalid\n");
 597                         r = 1;
 598                 } else if (r == 1) {
 599                         DEBUG(db->ctx, "The first signature is valid\n");
 600                         r = 0;
 601                 } else {
 602                         ERROR(db->ctx, "Error verifying the first signature: %s\n",
 603                                 ERR_error_string(ERR_get_error(), NULL));
 604                         r = -1;
 605                 }
 606         }
 607
 608         // Check second signature only when the first one was invalid
 609         if (r && db->signature2) {
 610                 hexdump(db->ctx, db->signature2, db->signature2_length);
 611
 612                 r = EVP_DigestVerifyFinal(mdctx,
 613                         (unsigned char*)db->signature2, db->signature2_length);
 614
 615                 if (r == 0) {
 616                         DEBUG(db->ctx, "The second signature is invalid\n");
 617                         r = 1;
 618                 } else if (r == 1) {
 619                         DEBUG(db->ctx, "The second signature is valid\n");
 620                         r = 0;
 621                 } else {
 622                         ERROR(db->ctx, "Error verifying the second signature: %s\n",
 623                                 ERR_error_string(ERR_get_error(), NULL));
 624                         r = -1;
 625                 }
 626         }
 627
 628         clock_t end = clock();
 629         INFO(db->ctx, "Signature checked in %.4fms\n",
 630                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 631
 632 CLEANUP:
 633         // Cleanup
 634         EVP_MD_CTX_free(mdctx);
 635         EVP_PKEY_free(pkey);
 636
 637         return r;
 638 }
 639
 640 LOC_EXPORT time_t loc_database_created_at(struct loc_database* db) {
 641         return db->created_at;
 642 }
 643
 644 LOC_EXPORT const char* loc_database_get_vendor(struct loc_database* db) {
 645         return loc_stringpool_get(db->pool, db->vendor);
 646 }
 647
 648 LOC_EXPORT const char* loc_database_get_description(struct loc_database* db) {
 649         return loc_stringpool_get(db->pool, db->description);
 650 }
 651
 652 LOC_EXPORT const char* loc_database_get_license(struct loc_database* db) {
 653         return loc_stringpool_get(db->pool, db->license);
 654 }
 655
 656 LOC_EXPORT size_t loc_database_count_as(struct loc_database* db) {
 657         return db->as_count;
 658 }
 659
 660 // Returns the AS at position pos
 661 static int loc_database_fetch_as(struct loc_database* db, struct loc_as** as, off_t pos) {
 662         if ((size_t)pos >= db->as_count)
 663                 return -EINVAL;
 664
 665         DEBUG(db->ctx, "Fetching AS at position %jd\n", (intmax_t)pos);
 666
 667         int r;
 668         switch (db->version) {
 669                 case LOC_DATABASE_VERSION_1:
 670                         r = loc_as_new_from_database_v1(db->ctx, db->pool, as, db->as_v1 + pos);
 671                         break;
 672
 673                 default:
 674                         return -1;
 675         }
 676
 677         if (r == 0) {
 678                 DEBUG(db->ctx, "Got AS%u\n", loc_as_get_number(*as));
 679         }
 680
 681         return r;
 682 }
 683
 684 // Performs a binary search to find the AS in the list
 685 LOC_EXPORT int loc_database_get_as(struct loc_database* db, struct loc_as** as, uint32_t number) {
 686         off_t lo = 0;
 687         off_t hi = db->as_count - 1;
 688
 689 #ifdef ENABLE_DEBUG
 690         // Save start time
 691         clock_t start = clock();
 692 #endif
 693
 694         while (lo <= hi) {
 695                 off_t i = (lo + hi) / 2;
 696
 697                 // Fetch AS in the middle between lo and hi
 698                 int r = loc_database_fetch_as(db, as, i);
 699                 if (r)
 700                         return r;
 701
 702                 // Check if this is a match
 703                 uint32_t as_number = loc_as_get_number(*as);
 704                 if (as_number == number) {
 705 #ifdef ENABLE_DEBUG
 706                         clock_t end = clock();
 707
 708                         // Log how fast this has been
 709                         DEBUG(db->ctx, "Found AS%u in %.4fms\n", as_number,
 710                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 711 #endif
 712
 713                         return 0;
 714                 }
 715
 716                 // If it wasn't, we release the AS and
 717                 // adjust our search pointers
 718                 loc_as_unref(*as);
 719
 720                 if (as_number < number) {
 721                         lo = i + 1;
 722                 } else
 723                         hi = i - 1;
 724         }
 725
 726         // Nothing found
 727         *as = NULL;
 728
 729         return 1;
 730 }
 731
 732 // Returns the network at position pos
 733 static int loc_database_fetch_network(struct loc_database* db, struct loc_network** network,
 734                 struct in6_addr* address, unsigned int prefix, off_t pos) {
 735         if ((size_t)pos >= db->networks_count) {
 736                 DEBUG(db->ctx, "Network ID out of range: %jd/%jd\n",
 737                         (intmax_t)pos, (intmax_t)db->networks_count);
 738                 return -EINVAL;
 739         }
 740
 741
 742         DEBUG(db->ctx, "Fetching network at position %jd\n", (intmax_t)pos);
 743
 744         int r;
 745         switch (db->version) {
 746                 case LOC_DATABASE_VERSION_1:
 747                         r = loc_network_new_from_database_v1(db->ctx, network,
 748                                 address, prefix, db->networks_v1 + pos);
 749                         break;
 750
 751                 default:
 752                         return -1;
 753         }
 754
 755 #ifdef ENABLE_DEBUG
 756         if (r == 0) {
 757                 char* string = loc_network_str(*network);
 758                 DEBUG(db->ctx, "Got network %s\n", string);
 759                 free(string);
 760         }
 761 #endif
 762
 763         return r;
 764 }
 765
 766 static int __loc_database_node_is_leaf(const struct loc_database_network_node_v1* node) {
 767         return (node->network != htobe32(0xffffffff));
 768 }
 769
 770 static int __loc_database_lookup_handle_leaf(struct loc_database* db, const struct in6_addr* address,
 771                 struct loc_network** network, struct in6_addr* network_address, unsigned int prefix,
 772                 const struct loc_database_network_node_v1* node) {
 773         off_t network_index = be32toh(node->network);
 774
 775         DEBUG(db->ctx, "Handling leaf node at %jd (%jd)\n", (intmax_t)(node - db->network_nodes_v1), (intmax_t)network_index);
 776
 777         // Fetch the network
 778         int r = loc_database_fetch_network(db, network,
 779                 network_address, prefix, network_index);
 780         if (r) {
 781                 ERROR(db->ctx, "Could not fetch network %jd from database\n", (intmax_t)network_index);
 782                 return r;
 783         }
 784
 785         // Check if the given IP address is inside the network
 786         if (!loc_network_match_address(*network, address)) {
 787                 DEBUG(db->ctx, "Searched address is not part of the network\n");
 788
 789                 loc_network_unref(*network);
 790                 *network = NULL;
 791                 return 1;
 792         }
 793
 794         // A network was found and the IP address matches
 795         return 0;
 796 }
 797
 798 // Searches for an exact match along the path
 799 static int __loc_database_lookup(struct loc_database* db, const struct in6_addr* address,
 800                 struct loc_network** network, struct in6_addr* network_address,
 801                 const struct loc_database_network_node_v1* node, unsigned int level) {
 802         int r;
 803         off_t node_index;
 804
 805         // Follow the path
 806         int bit = in6_addr_get_bit(address, level);
 807         in6_addr_set_bit(network_address, level, bit);
 808
 809         if (bit == 0)
 810                 node_index = be32toh(node->zero);
 811         else
 812                 node_index = be32toh(node->one);
 813
 814         // If the node index is zero, the tree ends here
 815         // and we cannot descend any further
 816         if (node_index > 0) {
 817                 // Check boundaries
 818                 if ((size_t)node_index >= db->network_nodes_count)
 819                         return -EINVAL;
 820
 821                 // Move on to the next node
 822                 r = __loc_database_lookup(db, address, network, network_address,
 823                         db->network_nodes_v1 + node_index, level + 1);
 824
 825                 // End here if a result was found
 826                 if (r == 0)
 827                         return r;
 828
 829                 // Raise any errors
 830                 else if (r < 0)
 831                         return r;
 832
 833                 DEBUG(db->ctx, "No match found below level %u\n", level);
 834         } else {
 835                 DEBUG(db->ctx, "Tree ended at level %u\n", level);
 836         }
 837
 838         // If this node has a leaf, we will check if it matches
 839         if (__loc_database_node_is_leaf(node)) {
 840                 r = __loc_database_lookup_handle_leaf(db, address, network, network_address, level, node);
 841                 if (r <= 0)
 842                         return r;
 843         }
 844
 845         return 1;
 846 }
 847
 848 LOC_EXPORT int loc_database_lookup(struct loc_database* db,
 849                 struct in6_addr* address, struct loc_network** network) {
 850         struct in6_addr network_address;
 851         memset(&network_address, 0, sizeof(network_address));
 852
 853         *network = NULL;
 854
 855 #ifdef ENABLE_DEBUG
 856         // Save start time
 857         clock_t start = clock();
 858 #endif
 859
 860         int r = __loc_database_lookup(db, address, network, &network_address,
 861                 db->network_nodes_v1, 0);
 862
 863 #ifdef ENABLE_DEBUG
 864         clock_t end = clock();
 865
 866         // Log how fast this has been
 867         DEBUG(db->ctx, "Executed network search in %.4fms\n",
 868                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 869 #endif
 870
 871         return r;
 872 }
 873
 874 LOC_EXPORT int loc_database_lookup_from_string(struct loc_database* db,
 875                 const char* string, struct loc_network** network) {
 876         struct in6_addr address;
 877
 878         int r = loc_parse_address(db->ctx, string, &address);
 879         if (r)
 880                 return r;
 881
 882         return loc_database_lookup(db, &address, network);
 883 }
 884
 885 // Returns the country at position pos
 886 static int loc_database_fetch_country(struct loc_database* db,
 887                 struct loc_country** country, off_t pos) {
 888         if ((size_t)pos >= db->countries_count)
 889                 return -EINVAL;
 890
 891         DEBUG(db->ctx, "Fetching country at position %jd\n", (intmax_t)pos);
 892
 893         int r;
 894         switch (db->version) {
 895                 case LOC_DATABASE_VERSION_1:
 896                         r = loc_country_new_from_database_v1(db->ctx, db->pool, country, db->countries_v1 + pos);
 897                         break;
 898
 899                 default:
 900                         return -1;
 901         }
 902
 903         if (r == 0) {
 904                 DEBUG(db->ctx, "Got country %s\n", loc_country_get_code(*country));
 905         }
 906
 907         return r;
 908 }
 909
 910 // Performs a binary search to find the country in the list
 911 LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 912                 struct loc_country** country, const char* code) {
 913         off_t lo = 0;
 914         off_t hi = db->countries_count - 1;
 915
 916 #ifdef ENABLE_DEBUG
 917         // Save start time
 918         clock_t start = clock();
 919 #endif
 920
 921         while (lo <= hi) {
 922                 off_t i = (lo + hi) / 2;
 923
 924                 // Fetch country in the middle between lo and hi
 925                 int r = loc_database_fetch_country(db, country, i);
 926                 if (r)
 927                         return r;
 928
 929                 // Check if this is a match
 930                 const char* cc = loc_country_get_code(*country);
 931                 int result = strcmp(code, cc);
 932
 933                 if (result == 0) {
 934 #ifdef ENABLE_DEBUG
 935                         clock_t end = clock();
 936
 937                         // Log how fast this has been
 938                         DEBUG(db->ctx, "Found country %s in %.4fms\n", cc,
 939                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 940 #endif
 941
 942                         return 0;
 943                 }
 944
 945                 // If it wasn't, we release the country and
 946                 // adjust our search pointers
 947                 loc_country_unref(*country);
 948
 949                 if (result > 0) {
 950                         lo = i + 1;
 951                 } else
 952                         hi = i - 1;
 953         }
 954
 955         // Nothing found
 956         *country = NULL;
 957
 958         return 1;
 959 }
 960
 961 // Enumerator
 962
 963 static void loc_database_enumerator_free(struct loc_database_enumerator* enumerator) {
 964         DEBUG(enumerator->ctx, "Releasing database enumerator %p\n", enumerator);
 965
 966         // Release all references
 967         loc_database_unref(enumerator->db);
 968         loc_unref(enumerator->ctx);
 969
 970         if (enumerator->string)
 971                 free(enumerator->string);
 972
 973         if (enumerator->countries)
 974                 loc_country_list_unref(enumerator->countries);
 975
 976         if (enumerator->asns)
 977                 loc_as_list_unref(enumerator->asns);
 978
 979         // Free network search
 980         free(enumerator->networks_visited);
 981
 982         // Free subnet stack
 983         if (enumerator->stack)
 984                 loc_network_list_unref(enumerator->stack);
 985
 986         free(enumerator);
 987 }
 988
 989 LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator,
 990                 struct loc_database* db, enum loc_database_enumerator_mode mode, int flags) {
 991         struct loc_database_enumerator* e = calloc(1, sizeof(*e));
 992         if (!e)
 993                 return -ENOMEM;
 994
 995         // Reference context
 996         e->ctx = loc_ref(db->ctx);
 997         e->db = loc_database_ref(db);
 998         e->mode = mode;
 999         e->refcount = 1;
1000
1001         // Flatten output?
1002         e->flatten = (flags & LOC_DB_ENUMERATOR_FLAGS_FLATTEN);
1003
1004         // Initialise graph search
1005         e->network_stack_depth = 1;
1006         e->networks_visited = calloc(db->network_nodes_count, sizeof(*e->networks_visited));
1007
1008         // Allocate stack
1009         int r = loc_network_list_new(e->ctx, &e->stack);
1010         if (r) {
1011                 loc_database_enumerator_free(e);
1012                 return r;
1013         }
1014
1015         DEBUG(e->ctx, "Database enumerator object allocated at %p\n", e);
1016
1017         *enumerator = e;
1018         return 0;
1019 }
1020
1021 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_ref(struct loc_database_enumerator* enumerator) {
1022         enumerator->refcount++;
1023
1024         return enumerator;
1025 }
1026
1027 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_unref(struct loc_database_enumerator* enumerator) {
1028         if (!enumerator)
1029                 return NULL;
1030
1031         if (--enumerator->refcount > 0)
1032                 return enumerator;
1033
1034         loc_database_enumerator_free(enumerator);
1035         return NULL;
1036 }
1037
1038 LOC_EXPORT int loc_database_enumerator_set_string(struct loc_database_enumerator* enumerator, const char* string) {
1039         enumerator->string = strdup(string);
1040
1041         // Make the string lowercase
1042         for (char *p = enumerator->string; *p; p++)
1043                 *p = tolower(*p);
1044
1045         return 0;
1046 }
1047
1048 LOC_EXPORT struct loc_country_list* loc_database_enumerator_get_countries(
1049                 struct loc_database_enumerator* enumerator) {
1050         if (!enumerator->countries)
1051                 return NULL;
1052
1053         return loc_country_list_ref(enumerator->countries);
1054 }
1055
1056 LOC_EXPORT int loc_database_enumerator_set_countries(
1057                 struct loc_database_enumerator* enumerator, struct loc_country_list* countries) {
1058         if (enumerator->countries)
1059                 loc_country_list_unref(enumerator->countries);
1060
1061         enumerator->countries = loc_country_list_ref(countries);
1062
1063         return 0;
1064 }
1065
1066 LOC_EXPORT struct loc_as_list* loc_database_enumerator_get_asns(
1067                 struct loc_database_enumerator* enumerator) {
1068         if (!enumerator->asns)
1069                 return NULL;
1070
1071         return loc_as_list_ref(enumerator->asns);
1072 }
1073
1074 LOC_EXPORT int loc_database_enumerator_set_asns(
1075                 struct loc_database_enumerator* enumerator, struct loc_as_list* asns) {
1076         if (enumerator->asns)
1077                 loc_as_list_unref(enumerator->asns);
1078
1079         enumerator->asns = loc_as_list_ref(asns);
1080
1081         return 0;
1082 }
1083
1084 LOC_EXPORT int loc_database_enumerator_set_flag(
1085                 struct loc_database_enumerator* enumerator, enum loc_network_flags flag) {
1086         enumerator->flags |= flag;
1087
1088         return 0;
1089 }
1090
1091 LOC_EXPORT int loc_database_enumerator_set_family(
1092                 struct loc_database_enumerator* enumerator, int family) {
1093         enumerator->family = family;
1094
1095         return 0;
1096 }
1097
1098 LOC_EXPORT int loc_database_enumerator_next_as(
1099                 struct loc_database_enumerator* enumerator, struct loc_as** as) {
1100         *as = NULL;
1101
1102         // Do not do anything if not in AS mode
1103         if (enumerator->mode != LOC_DB_ENUMERATE_ASES)
1104                 return 0;
1105
1106         struct loc_database* db = enumerator->db;
1107
1108         while (enumerator->as_index < db->as_count) {
1109                 // Fetch the next AS
1110                 int r = loc_database_fetch_as(db, as, enumerator->as_index++);
1111                 if (r)
1112                         return r;
1113
1114                 r = loc_as_match_string(*as, enumerator->string);
1115                 if (r == 1) {
1116                         DEBUG(enumerator->ctx, "AS%d (%s) matches %s\n",
1117                                 loc_as_get_number(*as), loc_as_get_name(*as), enumerator->string);
1118
1119                         return 0;
1120                 }
1121
1122                 // No match
1123                 loc_as_unref(*as);
1124                 *as = NULL;
1125         }
1126
1127         // Reset the index
1128         enumerator->as_index = 0;
1129
1130         // We have searched through all of them
1131         return 0;
1132 }
1133
1134 static int loc_database_enumerator_stack_push_node(
1135                 struct loc_database_enumerator* e, off_t offset, int i, int depth) {
1136         // Do not add empty nodes
1137         if (!offset)
1138                 return 0;
1139
1140         // Check if there is any space left on the stack
1141         if (e->network_stack_depth >= MAX_STACK_DEPTH) {
1142                 ERROR(e->ctx, "Maximum stack size reached: %d\n", e->network_stack_depth);
1143                 return -1;
1144         }
1145
1146         // Increase stack size
1147         int s = ++e->network_stack_depth;
1148
1149         DEBUG(e->ctx, "Added node %jd to stack (%d)\n", (intmax_t)offset, depth);
1150
1151         e->network_stack[s].offset = offset;
1152         e->network_stack[s].i = i;
1153         e->network_stack[s].depth = depth;
1154
1155         return 0;
1156 }
1157
1158 static int loc_database_enumerator_filter_network(
1159                 struct loc_database_enumerator* enumerator, struct loc_network* network) {
1160         // Skip if the family does not match
1161         if (enumerator->family && loc_network_address_family(network) != enumerator->family) {
1162                 DEBUG(enumerator->ctx, "Filtered network %p because of family not matching\n", network);
1163                 return 1;
1164         }
1165
1166         // Skip if the country code does not match
1167         if (enumerator->countries && !loc_country_list_empty(enumerator->countries)) {
1168                 const char* country_code = loc_network_get_country_code(network);
1169
1170                 if (!loc_country_list_contains_code(enumerator->countries, country_code)) {
1171                         DEBUG(enumerator->ctx, "Filtered network %p because of country code not matching\n", network);
1172                         return 1;
1173                 }
1174         }
1175
1176         // Skip if the ASN does not match
1177         if (enumerator->asns && !loc_as_list_empty(enumerator->asns)) {
1178                 uint32_t asn = loc_network_get_asn(network);
1179
1180                 if (!loc_as_list_contains_number(enumerator->asns, asn)) {
1181                         DEBUG(enumerator->ctx, "Filtered network %p because of ASN not matching\n", network);
1182                         return 1;
1183                 }
1184         }
1185
1186         // Skip if flags do not match
1187         if (enumerator->flags && !loc_network_match_flag(network, enumerator->flags)) {
1188                 DEBUG(enumerator->ctx, "Filtered network %p because of flags not matching\n", network);
1189                 return 1;
1190         }
1191
1192         // Do not filter
1193         return 0;
1194 }
1195
1196 static int __loc_database_enumerator_next_network(
1197                 struct loc_database_enumerator* enumerator, struct loc_network** network, int filter) {
1198         // Return top element from the stack
1199         while (1) {
1200                 *network = loc_network_list_pop_first(enumerator->stack);
1201
1202                 // Stack is empty
1203                 if (!*network)
1204                         break;
1205
1206                 // Throw away any networks by filter
1207                 if (filter && loc_database_enumerator_filter_network(enumerator, *network)) {
1208                         loc_network_unref(*network);
1209                         *network = NULL;
1210                         continue;
1211                 }
1212
1213                 // Return result
1214                 return 0;
1215         }
1216
1217         DEBUG(enumerator->ctx, "Called with a stack of %u nodes\n",
1218                 enumerator->network_stack_depth);
1219
1220         // Perform DFS
1221         while (enumerator->network_stack_depth > 0) {
1222                 DEBUG(enumerator->ctx, "Stack depth: %u\n", enumerator->network_stack_depth);
1223
1224                 // Get object from top of the stack
1225                 struct loc_node_stack* node = &enumerator->network_stack[enumerator->network_stack_depth];
1226
1227                 // Remove the node from the stack if we have already visited it
1228                 if (enumerator->networks_visited[node->offset]) {
1229                         enumerator->network_stack_depth--;
1230                         continue;
1231                 }
1232
1233                 // Mark the bits on the path correctly
1234                 in6_addr_set_bit(&enumerator->network_address,
1235                         (node->depth > 0) ? node->depth - 1 : 0, node->i);
1236
1237                 DEBUG(enumerator->ctx, "Looking at node %jd\n", (intmax_t)node->offset);
1238                 enumerator->networks_visited[node->offset]++;
1239
1240                 // Pop node from top of the stack
1241                 struct loc_database_network_node_v1* n =
1242                         enumerator->db->network_nodes_v1 + node->offset;
1243
1244                 // Add edges to stack
1245                 int r = loc_database_enumerator_stack_push_node(enumerator,
1246                         be32toh(n->one), 1, node->depth + 1);
1247
1248                 if (r)
1249                         return r;
1250
1251                 r = loc_database_enumerator_stack_push_node(enumerator,
1252                         be32toh(n->zero), 0, node->depth + 1);
1253
1254                 if (r)
1255                         return r;
1256
1257                 // Check if this node is a leaf and has a network object
1258                 if (__loc_database_node_is_leaf(n)) {
1259                         off_t network_index = be32toh(n->network);
1260
1261                         DEBUG(enumerator->ctx, "Node has a network at %jd\n", (intmax_t)network_index);
1262
1263                         // Fetch the network object
1264                         r = loc_database_fetch_network(enumerator->db, network,
1265                                 &enumerator->network_address, node->depth, network_index);
1266
1267                         // Break on any errors
1268                         if (r)
1269                                 return r;
1270
1271                         // Return all networks when the filter is disabled
1272                         if (!filter)
1273                                 return 0;
1274
1275                         // Check if we are interested in this network
1276                         if (loc_database_enumerator_filter_network(enumerator, *network)) {
1277                                 loc_network_unref(*network);
1278                                 *network = NULL;
1279
1280                                 continue;
1281                         }
1282
1283                         return 0;
1284                 }
1285         }
1286
1287         // Reached the end of the search
1288         return 0;
1289 }
1290
1291 static int __loc_database_enumerator_next_network_flattened(
1292                 struct loc_database_enumerator* enumerator, struct loc_network** network) {
1293         // Fetch the next network
1294         int r = __loc_database_enumerator_next_network(enumerator, network, 1);
1295         if (r)
1296                 return r;
1297
1298         // End if we could not read another network
1299         if (!*network)
1300                 return 0;
1301
1302         struct loc_network* subnet = NULL;
1303         struct loc_network_list* subnets;
1304
1305         // Create a list with all subnets
1306         r = loc_network_list_new(enumerator->ctx, &subnets);
1307         if (r)
1308                 return r;
1309
1310         // Search all subnets from the database
1311         while (1) {
1312                 // Fetch the next network in line
1313                 r = __loc_database_enumerator_next_network(enumerator, &subnet, 0);
1314                 if (r) {
1315                         loc_network_unref(subnet);
1316                         loc_network_list_unref(subnets);
1317
1318                         return r;
1319                 }
1320
1321                 // End if we did not receive another subnet
1322                 if (!subnet)
1323                         break;
1324
1325                 // Collect all subnets in a list
1326                 if (loc_network_is_subnet(*network, subnet)) {
1327                         r = loc_network_list_push(subnets, subnet);
1328                         if (r) {
1329                                 loc_network_unref(subnet);
1330                                 loc_network_list_unref(subnets);
1331
1332                                 return r;
1333                         }
1334
1335                         loc_network_unref(subnet);
1336                         continue;
1337                 }
1338
1339                 // If this is not a subnet, we push it back onto the stack and break
1340                 r = loc_network_list_push(enumerator->stack, subnet);
1341                 if (r) {
1342                         loc_network_unref(subnet);
1343                         loc_network_list_unref(subnets);
1344
1345                         return r;
1346                 }
1347
1348                 loc_network_unref(subnet);
1349                 break;
1350         }
1351
1352         DEBUG(enumerator->ctx, "Found %zu subnet(s)\n", loc_network_list_size(subnets));
1353
1354         // We can abort here if the network has no subnets
1355         if (loc_network_list_empty(subnets)) {
1356                 loc_network_list_unref(subnets);
1357
1358                 return 0;
1359         }
1360
1361         // If the network has any subnets, we will break it into smaller parts
1362         // without the subnets.
1363         struct loc_network_list* excluded = loc_network_exclude_list(*network, subnets);
1364         if (!excluded) {
1365                 loc_network_list_unref(subnets);
1366                 return -1;
1367         }
1368
1369         // Merge subnets onto the stack
1370         r = loc_network_list_merge(enumerator->stack, subnets);
1371         if (r) {
1372                 loc_network_list_unref(subnets);
1373                 loc_network_list_unref(excluded);
1374
1375                 return r;
1376         }
1377
1378         // Push excluded list onto the stack
1379         r = loc_network_list_merge(enumerator->stack, excluded);
1380         if (r) {
1381                 loc_network_list_unref(subnets);
1382                 loc_network_list_unref(excluded);
1383
1384                 return r;
1385         }
1386
1387         loc_network_list_unref(subnets);
1388         loc_network_list_unref(excluded);
1389
1390         // Drop the network and restart the whole process again to pick the next network
1391         loc_network_unref(*network);
1392
1393         return __loc_database_enumerator_next_network_flattened(enumerator, network);
1394 }
1395
1396 LOC_EXPORT int loc_database_enumerator_next_network(
1397                 struct loc_database_enumerator* enumerator, struct loc_network** network) {
1398         // Do not do anything if not in network mode
1399         if (enumerator->mode != LOC_DB_ENUMERATE_NETWORKS)
1400                 return 0;
1401
1402         // Flatten output?
1403         if (enumerator->flatten)
1404                 return __loc_database_enumerator_next_network_flattened(enumerator, network);
1405
1406         return __loc_database_enumerator_next_network(enumerator, network, 1);
1407 }
1408
1409 LOC_EXPORT int loc_database_enumerator_next_country(
1410                 struct loc_database_enumerator* enumerator, struct loc_country** country) {
1411         *country = NULL;
1412
1413         // Do not do anything if not in country mode
1414         if (enumerator->mode != LOC_DB_ENUMERATE_COUNTRIES)
1415                 return 0;
1416
1417         struct loc_database* db = enumerator->db;
1418
1419         while (enumerator->country_index < db->countries_count) {
1420                 // Fetch the next country
1421                 int r = loc_database_fetch_country(db, country, enumerator->country_index++);
1422                 if (r)
1423                         return r;
1424
1425                 // We do not filter here, so it always is a match
1426                 return 0;
1427         }
1428
1429         // Reset the index
1430         enumerator->country_index = 0;
1431
1432         // We have searched through all of them
1433         return 0;
1434 }