src/database.c

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