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 <loc/libloc.h>
  36 #include <loc/as.h>
  37 #include <loc/compat.h>
  38 #include <loc/country.h>
  39 #include <loc/database.h>
  40 #include <loc/format.h>
  41 #include <loc/network.h>
  42 #include <loc/private.h>
  43 #include <loc/stringpool.h>
  44
  45 struct loc_database {
  46         struct loc_ctx* ctx;
  47         int refcount;
  48
  49         unsigned int version;
  50         time_t created_at;
  51         off_t vendor;
  52         off_t description;
  53         off_t license;
  54
  55         // ASes in the database
  56         struct loc_database_as_v0* as_v0;
  57         size_t as_count;
  58
  59         // Network tree
  60         struct loc_database_network_node_v0* network_nodes_v0;
  61         size_t network_nodes_count;
  62
  63         // Networks
  64         struct loc_database_network_v0* networks_v0;
  65         size_t networks_count;
  66
  67         // Countries
  68         struct loc_database_country_v0* countries_v0;
  69         size_t countries_count;
  70
  71         struct loc_stringpool* pool;
  72 };
  73
  74 #define MAX_STACK_DEPTH 256
  75
  76 struct loc_node_stack {
  77         off_t offset;
  78         int i; // Is this node 0 or 1?
  79         int depth;
  80 };
  81
  82 struct loc_database_enumerator {
  83         struct loc_ctx* ctx;
  84         struct loc_database* db;
  85         enum loc_database_enumerator_mode mode;
  86         int refcount;
  87
  88         // Search string
  89         char* string;
  90         char country_code[3];
  91         uint32_t asn;
  92         enum loc_network_flags flags;
  93
  94         // Index of the AS we are looking at
  95         unsigned int as_index;
  96
  97         // Network state
  98         struct in6_addr network_address;
  99         struct loc_node_stack network_stack[MAX_STACK_DEPTH];
 100         int network_stack_depth;
 101         unsigned int* networks_visited;
 102 };
 103
 104 static int loc_database_read_magic(struct loc_database* db, FILE* f) {
 105         struct loc_database_magic magic;
 106
 107         // Read from file
 108         size_t bytes_read = fread(&magic, 1, sizeof(magic), f);
 109
 110         // Check if we have been able to read enough data
 111         if (bytes_read < sizeof(magic)) {
 112                 ERROR(db->ctx, "Could not read enough data to validate magic bytes\n");
 113                 DEBUG(db->ctx, "Read %zu bytes, but needed %zu\n", bytes_read, sizeof(magic));
 114                 return -ENOMSG;
 115         }
 116
 117         // Compare magic bytes
 118         if (memcmp(LOC_DATABASE_MAGIC, magic.magic, strlen(LOC_DATABASE_MAGIC)) == 0) {
 119                 DEBUG(db->ctx, "Magic value matches\n");
 120
 121                 // Parse version
 122                 db->version = be16toh(magic.version);
 123                 DEBUG(db->ctx, "Database version is %u\n", db->version);
 124
 125                 return 0;
 126         }
 127
 128         ERROR(db->ctx, "Database format is not compatible\n");
 129
 130         // Return an error
 131         return 1;
 132 }
 133
 134 static int loc_database_read_as_section_v0(struct loc_database* db,
 135                 FILE* f, const struct loc_database_header_v0* header) {
 136         off_t as_offset  = be32toh(header->as_offset);
 137         size_t as_length = be32toh(header->as_length);
 138
 139         DEBUG(db->ctx, "Reading AS section from %jd (%zu bytes)\n", (intmax_t)as_offset, as_length);
 140
 141         if (as_length > 0) {
 142                 db->as_v0 = mmap(NULL, as_length, PROT_READ,
 143                         MAP_SHARED, fileno(f), as_offset);
 144
 145                 if (db->as_v0 == MAP_FAILED)
 146                         return -errno;
 147         }
 148
 149         db->as_count = as_length / sizeof(*db->as_v0);
 150
 151         INFO(db->ctx, "Read %zu ASes from the database\n", db->as_count);
 152
 153         return 0;
 154 }
 155
 156 static int loc_database_read_network_nodes_section_v0(struct loc_database* db,
 157                 FILE* f, const struct loc_database_header_v0* header) {
 158         off_t network_nodes_offset  = be32toh(header->network_tree_offset);
 159         size_t network_nodes_length = be32toh(header->network_tree_length);
 160
 161         DEBUG(db->ctx, "Reading network nodes section from %jd (%zu bytes)\n",
 162                 (intmax_t)network_nodes_offset, network_nodes_length);
 163
 164         if (network_nodes_length > 0) {
 165                 db->network_nodes_v0 = mmap(NULL, network_nodes_length, PROT_READ,
 166                         MAP_SHARED, fileno(f), network_nodes_offset);
 167
 168                 if (db->network_nodes_v0 == MAP_FAILED)
 169                         return -errno;
 170         }
 171
 172         db->network_nodes_count = network_nodes_length / sizeof(*db->network_nodes_v0);
 173
 174         INFO(db->ctx, "Read %zu network nodes from the database\n", db->network_nodes_count);
 175
 176         return 0;
 177 }
 178
 179 static int loc_database_read_networks_section_v0(struct loc_database* db,
 180                 FILE* f, const struct loc_database_header_v0* header) {
 181         off_t networks_offset  = be32toh(header->network_data_offset);
 182         size_t networks_length = be32toh(header->network_data_length);
 183
 184         DEBUG(db->ctx, "Reading networks section from %jd (%zu bytes)\n",
 185                 (intmax_t)networks_offset, networks_length);
 186
 187         if (networks_length > 0) {
 188                 db->networks_v0 = mmap(NULL, networks_length, PROT_READ,
 189                         MAP_SHARED, fileno(f), networks_offset);
 190
 191                 if (db->networks_v0 == MAP_FAILED)
 192                         return -errno;
 193         }
 194
 195         db->networks_count = networks_length / sizeof(*db->networks_v0);
 196
 197         INFO(db->ctx, "Read %zu networks from the database\n", db->networks_count);
 198
 199         return 0;
 200 }
 201
 202 static int loc_database_read_countries_section_v0(struct loc_database* db,
 203                 FILE* f, const struct loc_database_header_v0* header) {
 204         off_t countries_offset  = be32toh(header->countries_offset);
 205         size_t countries_length = be32toh(header->countries_length);
 206
 207         DEBUG(db->ctx, "Reading countries section from %jd (%zu bytes)\n",
 208                 (intmax_t)countries_offset, countries_length);
 209
 210         if (countries_length > 0) {
 211                 db->countries_v0 = mmap(NULL, countries_length, PROT_READ,
 212                         MAP_SHARED, fileno(f), countries_offset);
 213
 214                 if (db->countries_v0 == MAP_FAILED)
 215                         return -errno;
 216         }
 217
 218         db->countries_count = countries_length / sizeof(*db->countries_v0);
 219
 220         INFO(db->ctx, "Read %zu countries from the database\n",
 221                 db->countries_count);
 222
 223         return 0;
 224 }
 225
 226 static int loc_database_read_header_v0(struct loc_database* db, FILE* f) {
 227         struct loc_database_header_v0 header;
 228
 229         // Read from file
 230         size_t size = fread(&header, 1, sizeof(header), f);
 231
 232         if (size < sizeof(header)) {
 233                 ERROR(db->ctx, "Could not read enough data for header\n");
 234                 return -ENOMSG;
 235         }
 236
 237         // Copy over data
 238         db->created_at  = be64toh(header.created_at);
 239         db->vendor      = be32toh(header.vendor);
 240         db->description = be32toh(header.description);
 241         db->license     = be32toh(header.license);
 242
 243         // Open pool
 244         off_t pool_offset  = be32toh(header.pool_offset);
 245         size_t pool_length = be32toh(header.pool_length);
 246
 247         int r = loc_stringpool_open(db->ctx, &db->pool,
 248                 f, pool_length, pool_offset);
 249         if (r)
 250                 return r;
 251
 252         // AS section
 253         r = loc_database_read_as_section_v0(db, f, &header);
 254         if (r)
 255                 return r;
 256
 257         // Network Nodes
 258         r = loc_database_read_network_nodes_section_v0(db, f, &header);
 259         if (r)
 260                 return r;
 261
 262         // Networks
 263         r = loc_database_read_networks_section_v0(db, f, &header);
 264         if (r)
 265                 return r;
 266
 267         // countries
 268         r = loc_database_read_countries_section_v0(db, f, &header);
 269         if (r)
 270                 return r;
 271
 272         return 0;
 273 }
 274
 275 static int loc_database_read_header(struct loc_database* db, FILE* f) {
 276         switch (db->version) {
 277                 case 0:
 278                         return loc_database_read_header_v0(db, f);
 279
 280                 default:
 281                         ERROR(db->ctx, "Incompatible database version: %u\n", db->version);
 282                         return 1;
 283         }
 284 }
 285
 286 static int loc_database_read(struct loc_database* db, FILE* f) {
 287         clock_t start = clock();
 288
 289         // Read magic bytes
 290         int r = loc_database_read_magic(db, f);
 291         if (r)
 292                 return r;
 293
 294         // Read the header
 295         r = loc_database_read_header(db, f);
 296         if (r)
 297                 return r;
 298
 299         clock_t end = clock();
 300
 301         INFO(db->ctx, "Opened database in %.4fms\n",
 302                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 303
 304         return 0;
 305 }
 306
 307 LOC_EXPORT int loc_database_new(struct loc_ctx* ctx, struct loc_database** database, FILE* f) {
 308         // Fail on invalid file handle
 309         if (!f)
 310                 return -EINVAL;
 311
 312         struct loc_database* db = calloc(1, sizeof(*db));
 313         if (!db)
 314                 return -ENOMEM;
 315
 316         // Reference context
 317         db->ctx = loc_ref(ctx);
 318         db->refcount = 1;
 319
 320         DEBUG(db->ctx, "Database object allocated at %p\n", db);
 321
 322         int r = loc_database_read(db, f);
 323         if (r) {
 324                 loc_database_unref(db);
 325                 return r;
 326         }
 327
 328         *database = db;
 329
 330         return 0;
 331 }
 332
 333 LOC_EXPORT struct loc_database* loc_database_ref(struct loc_database* db) {
 334         db->refcount++;
 335
 336         return db;
 337 }
 338
 339 static void loc_database_free(struct loc_database* db) {
 340         int r;
 341
 342         DEBUG(db->ctx, "Releasing database %p\n", db);
 343
 344         // Removing all ASes
 345         if (db->as_v0) {
 346                 r = munmap(db->as_v0, db->as_count * sizeof(*db->as_v0));
 347                 if (r)
 348                         ERROR(db->ctx, "Could not unmap AS section: %s\n", strerror(errno));
 349         }
 350
 351         // Remove mapped network sections
 352         if (db->networks_v0) {
 353                 r = munmap(db->networks_v0, db->networks_count * sizeof(*db->networks_v0));
 354                 if (r)
 355                         ERROR(db->ctx, "Could not unmap networks section: %s\n", strerror(errno));
 356         }
 357
 358         // Remove mapped network nodes section
 359         if (db->network_nodes_v0) {
 360                 r = munmap(db->network_nodes_v0, db->network_nodes_count * sizeof(*db->network_nodes_v0));
 361                 if (r)
 362                         ERROR(db->ctx, "Could not unmap network nodes section: %s\n", strerror(errno));
 363         }
 364
 365         loc_stringpool_unref(db->pool);
 366
 367         loc_unref(db->ctx);
 368         free(db);
 369 }
 370
 371 LOC_EXPORT struct loc_database* loc_database_unref(struct loc_database* db) {
 372         if (--db->refcount > 0)
 373                 return NULL;
 374
 375         loc_database_free(db);
 376         return NULL;
 377 }
 378
 379 LOC_EXPORT time_t loc_database_created_at(struct loc_database* db) {
 380         return db->created_at;
 381 }
 382
 383 LOC_EXPORT const char* loc_database_get_vendor(struct loc_database* db) {
 384         return loc_stringpool_get(db->pool, db->vendor);
 385 }
 386
 387 LOC_EXPORT const char* loc_database_get_description(struct loc_database* db) {
 388         return loc_stringpool_get(db->pool, db->description);
 389 }
 390
 391 LOC_EXPORT const char* loc_database_get_license(struct loc_database* db) {
 392         return loc_stringpool_get(db->pool, db->license);
 393 }
 394
 395 LOC_EXPORT size_t loc_database_count_as(struct loc_database* db) {
 396         return db->as_count;
 397 }
 398
 399 // Returns the AS at position pos
 400 static int loc_database_fetch_as(struct loc_database* db, struct loc_as** as, off_t pos) {
 401         if ((size_t)pos >= db->as_count)
 402                 return -EINVAL;
 403
 404         DEBUG(db->ctx, "Fetching AS at position %jd\n", (intmax_t)pos);
 405
 406         int r;
 407         switch (db->version) {
 408                 case 0:
 409                         r = loc_as_new_from_database_v0(db->ctx, db->pool, as, db->as_v0 + pos);
 410                         break;
 411
 412                 default:
 413                         return -1;
 414         }
 415
 416         if (r == 0) {
 417                 DEBUG(db->ctx, "Got AS%u\n", loc_as_get_number(*as));
 418         }
 419
 420         return r;
 421 }
 422
 423 // Performs a binary search to find the AS in the list
 424 LOC_EXPORT int loc_database_get_as(struct loc_database* db, struct loc_as** as, uint32_t number) {
 425         off_t lo = 0;
 426         off_t hi = db->as_count - 1;
 427
 428         // Save start time
 429         clock_t start = clock();
 430
 431         while (lo <= hi) {
 432                 off_t i = (lo + hi) / 2;
 433
 434                 // Fetch AS in the middle between lo and hi
 435                 int r = loc_database_fetch_as(db, as, i);
 436                 if (r)
 437                         return r;
 438
 439                 // Check if this is a match
 440                 uint32_t as_number = loc_as_get_number(*as);
 441                 if (as_number == number) {
 442                         clock_t end = clock();
 443
 444                         // Log how fast this has been
 445                         DEBUG(db->ctx, "Found AS%u in %.4fms\n", as_number,
 446                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 447
 448                         return 0;
 449                 }
 450
 451                 // If it wasn't, we release the AS and
 452                 // adjust our search pointers
 453                 loc_as_unref(*as);
 454
 455                 if (as_number < number) {
 456                         lo = i + 1;
 457                 } else
 458                         hi = i - 1;
 459         }
 460
 461         // Nothing found
 462         *as = NULL;
 463
 464         return 1;
 465 }
 466
 467 // Returns the network at position pos
 468 static int loc_database_fetch_network(struct loc_database* db, struct loc_network** network,
 469                 struct in6_addr* address, unsigned int prefix, off_t pos) {
 470         if ((size_t)pos >= db->networks_count) {
 471                 DEBUG(db->ctx, "Network ID out of range: %jd/%jd\n",
 472                         (intmax_t)pos, (intmax_t)db->networks_count);
 473                 return -EINVAL;
 474         }
 475
 476
 477         DEBUG(db->ctx, "Fetching network at position %jd\n", (intmax_t)pos);
 478
 479         int r;
 480         switch (db->version) {
 481                 case 0:
 482                         r = loc_network_new_from_database_v0(db->ctx, network,
 483                                 address, prefix, db->networks_v0 + pos);
 484                         break;
 485
 486                 default:
 487                         return -1;
 488         }
 489
 490         if (r == 0) {
 491                 char* string = loc_network_str(*network);
 492                 DEBUG(db->ctx, "Got network %s\n", string);
 493                 free(string);
 494         }
 495
 496         return r;
 497 }
 498
 499 static int __loc_database_node_is_leaf(const struct loc_database_network_node_v0* node) {
 500         return (node->network != htobe32(0xffffffff));
 501 }
 502
 503 static int __loc_database_lookup_handle_leaf(struct loc_database* db, const struct in6_addr* address,
 504                 struct loc_network** network, struct in6_addr* network_address, unsigned int prefix,
 505                 const struct loc_database_network_node_v0* node) {
 506         off_t network_index = be32toh(node->network);
 507
 508         DEBUG(db->ctx, "Handling leaf node at %jd (%jd)\n", (intmax_t)(node - db->network_nodes_v0), (intmax_t)network_index);
 509
 510         // Fetch the network
 511         int r = loc_database_fetch_network(db, network,
 512                 network_address, prefix, network_index);
 513         if (r) {
 514                 ERROR(db->ctx, "Could not fetch network %jd from database\n", (intmax_t)network_index);
 515                 return r;
 516         }
 517
 518         // Check if the given IP address is inside the network
 519         r = loc_network_match_address(*network, address);
 520         if (r) {
 521                 DEBUG(db->ctx, "Searched address is not part of the network\n");
 522
 523                 loc_network_unref(*network);
 524                 *network = NULL;
 525                 return 1;
 526         }
 527
 528         // A network was found and the IP address matches
 529         return 0;
 530 }
 531
 532 // Searches for an exact match along the path
 533 static int __loc_database_lookup(struct loc_database* db, const struct in6_addr* address,
 534                 struct loc_network** network, struct in6_addr* network_address,
 535                 const struct loc_database_network_node_v0* node, unsigned int level) {
 536         int r;
 537         off_t node_index;
 538
 539         // Follow the path
 540         int bit = in6_addr_get_bit(address, level);
 541         in6_addr_set_bit(network_address, level, bit);
 542
 543         if (bit == 0)
 544                 node_index = be32toh(node->zero);
 545         else
 546                 node_index = be32toh(node->one);
 547
 548         // If the node index is zero, the tree ends here
 549         // and we cannot descend any further
 550         if (node_index > 0) {
 551                 // Check boundaries
 552                 if ((size_t)node_index >= db->network_nodes_count)
 553                         return -EINVAL;
 554
 555                 // Move on to the next node
 556                 r = __loc_database_lookup(db, address, network, network_address,
 557                         db->network_nodes_v0 + node_index, level + 1);
 558
 559                 // End here if a result was found
 560                 if (r == 0)
 561                         return r;
 562
 563                 // Raise any errors
 564                 else if (r < 0)
 565                         return r;
 566
 567                 DEBUG(db->ctx, "No match found below level %u\n", level);
 568         } else {
 569                 DEBUG(db->ctx, "Tree ended at level %u\n", level);
 570         }
 571
 572         // If this node has a leaf, we will check if it matches
 573         if (__loc_database_node_is_leaf(node)) {
 574                 r = __loc_database_lookup_handle_leaf(db, address, network, network_address, level, node);
 575                 if (r <= 0)
 576                         return r;
 577         }
 578
 579         return 1;
 580 }
 581
 582 LOC_EXPORT int loc_database_lookup(struct loc_database* db,
 583                 struct in6_addr* address, struct loc_network** network) {
 584         struct in6_addr network_address;
 585         memset(&network_address, 0, sizeof(network_address));
 586
 587         *network = NULL;
 588
 589         // Save start time
 590         clock_t start = clock();
 591
 592         int r = __loc_database_lookup(db, address, network, &network_address,
 593                 db->network_nodes_v0, 0);
 594
 595         clock_t end = clock();
 596
 597         // Log how fast this has been
 598         DEBUG(db->ctx, "Executed network search in %.4fms\n",
 599                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 600
 601         return r;
 602 }
 603
 604 LOC_EXPORT int loc_database_lookup_from_string(struct loc_database* db,
 605                 const char* string, struct loc_network** network) {
 606         struct in6_addr address;
 607
 608         int r = loc_parse_address(db->ctx, string, &address);
 609         if (r)
 610                 return r;
 611
 612         return loc_database_lookup(db, &address, network);
 613 }
 614
 615 // Returns the country at position pos
 616 static int loc_database_fetch_country(struct loc_database* db,
 617                 struct loc_country** country, off_t pos) {
 618         if ((size_t)pos >= db->countries_count)
 619                 return -EINVAL;
 620
 621         DEBUG(db->ctx, "Fetching country at position %jd\n", (intmax_t)pos);
 622
 623         int r;
 624         switch (db->version) {
 625                 case 0:
 626                         r = loc_country_new_from_database_v0(db->ctx, db->pool, country, db->countries_v0 + pos);
 627                         break;
 628
 629                 default:
 630                         return -1;
 631         }
 632
 633         if (r == 0) {
 634                 DEBUG(db->ctx, "Got country %s\n", loc_country_get_code(*country));
 635         }
 636
 637         return r;
 638 }
 639
 640 // Performs a binary search to find the country in the list
 641 LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 642                 struct loc_country** country, const char* code) {
 643         off_t lo = 0;
 644         off_t hi = db->countries_count - 1;
 645
 646         // Save start time
 647         clock_t start = clock();
 648
 649         while (lo <= hi) {
 650                 off_t i = (lo + hi) / 2;
 651
 652                 // Fetch country in the middle between lo and hi
 653                 int r = loc_database_fetch_country(db, country, i);
 654                 if (r)
 655                         return r;
 656
 657                 // Check if this is a match
 658                 const char* cc = loc_country_get_code(*country);
 659                 int result = strcmp(code, cc);
 660
 661                 if (result == 0) {
 662                         clock_t end = clock();
 663
 664                         // Log how fast this has been
 665                         DEBUG(db->ctx, "Found country %s in %.4fms\n", cc,
 666                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 667
 668                         return 0;
 669                 }
 670
 671                 // If it wasn't, we release the country and
 672                 // adjust our search pointers
 673                 loc_country_unref(*country);
 674
 675                 if (result > 0) {
 676                         lo = i + 1;
 677                 } else
 678                         hi = i - 1;
 679         }
 680
 681         // Nothing found
 682         *country = NULL;
 683
 684         return 1;
 685 }
 686
 687 // Enumerator
 688
 689 LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator,
 690                 struct loc_database* db, enum loc_database_enumerator_mode mode) {
 691         struct loc_database_enumerator* e = calloc(1, sizeof(*e));
 692         if (!e)
 693                 return -ENOMEM;
 694
 695         // Reference context
 696         e->ctx = loc_ref(db->ctx);
 697         e->db = loc_database_ref(db);
 698         e->mode = mode;
 699         e->refcount = 1;
 700
 701         // Initialise graph search
 702         //e->network_stack[++e->network_stack_depth] = 0;
 703         e->network_stack_depth = 1;
 704         e->networks_visited = calloc(db->network_nodes_count, sizeof(*e->networks_visited));
 705
 706         DEBUG(e->ctx, "Database enumerator object allocated at %p\n", e);
 707
 708         *enumerator = e;
 709         return 0;
 710 }
 711
 712 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_ref(struct loc_database_enumerator* enumerator) {
 713         enumerator->refcount++;
 714
 715         return enumerator;
 716 }
 717
 718 static void loc_database_enumerator_free(struct loc_database_enumerator* enumerator) {
 719         DEBUG(enumerator->ctx, "Releasing database enumerator %p\n", enumerator);
 720
 721         // Release all references
 722         loc_database_unref(enumerator->db);
 723         loc_unref(enumerator->ctx);
 724
 725         if (enumerator->string)
 726                 free(enumerator->string);
 727
 728         // Free network search
 729         free(enumerator->networks_visited);
 730
 731         free(enumerator);
 732 }
 733
 734 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_unref(struct loc_database_enumerator* enumerator) {
 735         if (!enumerator)
 736                 return NULL;
 737
 738         if (--enumerator->refcount > 0)
 739                 return enumerator;
 740
 741         loc_database_enumerator_free(enumerator);
 742         return NULL;
 743 }
 744
 745 LOC_EXPORT int loc_database_enumerator_set_string(struct loc_database_enumerator* enumerator, const char* string) {
 746         enumerator->string = strdup(string);
 747
 748         // Make the string lowercase
 749         for (char *p = enumerator->string; *p; p++)
 750                 *p = tolower(*p);
 751
 752         return 0;
 753 }
 754
 755 LOC_EXPORT int loc_database_enumerator_set_country_code(struct loc_database_enumerator* enumerator, const char* country_code) {
 756         // Set empty country code
 757         if (!country_code || !*country_code) {
 758                 *enumerator->country_code = '\0';
 759                 return 0;
 760         }
 761
 762         // Treat A1, A2, A3 as special country codes,
 763         // but perform search for flags instead
 764         if (strcmp(country_code, "A1") == 0) {
 765                 return loc_database_enumerator_set_flag(enumerator,
 766                         LOC_NETWORK_FLAG_ANONYMOUS_PROXY);
 767         } else if (strcmp(country_code, "A2") == 0) {
 768                 return loc_database_enumerator_set_flag(enumerator,
 769                         LOC_NETWORK_FLAG_SATELLITE_PROVIDER);
 770         } else if (strcmp(country_code, "A3") == 0) {
 771                 return loc_database_enumerator_set_flag(enumerator,
 772                         LOC_NETWORK_FLAG_ANYCAST);
 773         }
 774
 775         // Country codes must be two characters
 776         if (!loc_country_code_is_valid(country_code))
 777                 return -EINVAL;
 778
 779         for (unsigned int i = 0; i < 3; i++) {
 780                 enumerator->country_code[i] = country_code[i];
 781         }
 782
 783         return 0;
 784 }
 785
 786 LOC_EXPORT int loc_database_enumerator_set_asn(
 787                 struct loc_database_enumerator* enumerator, unsigned int asn) {
 788         enumerator->asn = asn;
 789
 790         return 0;
 791 }
 792
 793 LOC_EXPORT int loc_database_enumerator_set_flag(
 794                 struct loc_database_enumerator* enumerator, enum loc_network_flags flag) {
 795         enumerator->flags |= flag;
 796
 797         return 0;
 798 }
 799
 800 LOC_EXPORT int loc_database_enumerator_next_as(
 801                 struct loc_database_enumerator* enumerator, struct loc_as** as) {
 802         *as = NULL;
 803
 804         // Do not do anything if not in AS mode
 805         if (enumerator->mode != LOC_DB_ENUMERATE_ASES)
 806                 return 0;
 807
 808         struct loc_database* db = enumerator->db;
 809
 810         while (enumerator->as_index < db->as_count) {
 811                 // Fetch the next AS
 812                 int r = loc_database_fetch_as(db, as, enumerator->as_index++);
 813                 if (r)
 814                         return r;
 815
 816                 r = loc_as_match_string(*as, enumerator->string);
 817                 if (r == 1) {
 818                         DEBUG(enumerator->ctx, "AS%d (%s) matches %s\n",
 819                                 loc_as_get_number(*as), loc_as_get_name(*as), enumerator->string);
 820
 821                         return 0;
 822                 }
 823
 824                 // No match
 825                 loc_as_unref(*as);
 826                 *as = NULL;
 827         }
 828
 829         // Reset the index
 830         enumerator->as_index = 0;
 831
 832         // We have searched through all of them
 833         return 0;
 834 }
 835
 836 static int loc_database_enumerator_stack_push_node(
 837                 struct loc_database_enumerator* e, off_t offset, int i, int depth) {
 838         // Do not add empty nodes
 839         if (!offset)
 840                 return 0;
 841
 842         // Check if there is any space left on the stack
 843         if (e->network_stack_depth >= MAX_STACK_DEPTH) {
 844                 ERROR(e->ctx, "Maximum stack size reached: %d\n", e->network_stack_depth);
 845                 return -1;
 846         }
 847
 848         // Increase stack size
 849         int s = ++e->network_stack_depth;
 850
 851         DEBUG(e->ctx, "Added node %jd to stack (%d)\n", (intmax_t)offset, depth);
 852
 853         e->network_stack[s].offset = offset;
 854         e->network_stack[s].i = i;
 855         e->network_stack[s].depth = depth;
 856
 857         return 0;
 858 }
 859
 860 LOC_EXPORT int loc_database_enumerator_next_network(
 861                 struct loc_database_enumerator* enumerator, struct loc_network** network) {
 862         // Reset network
 863         *network = NULL;
 864
 865         // Do not do anything if not in network mode
 866         if (enumerator->mode != LOC_DB_ENUMERATE_NETWORKS)
 867                 return 0;
 868
 869         int r;
 870
 871         DEBUG(enumerator->ctx, "Called with a stack of %u nodes\n",
 872                 enumerator->network_stack_depth);
 873
 874         // Perform DFS
 875         while (enumerator->network_stack_depth > 0) {
 876                 DEBUG(enumerator->ctx, "Stack depth: %u\n", enumerator->network_stack_depth);
 877
 878                 // Get object from top of the stack
 879                 struct loc_node_stack* node = &enumerator->network_stack[enumerator->network_stack_depth];
 880
 881                 // Remove the node from the stack if we have already visited it
 882                 if (enumerator->networks_visited[node->offset]) {
 883                         enumerator->network_stack_depth--;
 884                         continue;
 885                 }
 886
 887                 // Mark the bits on the path correctly
 888                 in6_addr_set_bit(&enumerator->network_address,
 889                         (node->depth > 0) ? node->depth - 1 : 0, node->i);
 890
 891                 DEBUG(enumerator->ctx, "Looking at node %jd\n", (intmax_t)node->offset);
 892                 enumerator->networks_visited[node->offset]++;
 893
 894                 // Pop node from top of the stack
 895                 struct loc_database_network_node_v0* n =
 896                         enumerator->db->network_nodes_v0 + node->offset;
 897
 898                 // Add edges to stack
 899                 r = loc_database_enumerator_stack_push_node(enumerator,
 900                         be32toh(n->one), 1, node->depth + 1);
 901
 902                 if (r)
 903                         return r;
 904
 905                 r = loc_database_enumerator_stack_push_node(enumerator,
 906                         be32toh(n->zero), 0, node->depth + 1);
 907
 908                 if (r)
 909                         return r;
 910
 911                 // Check if this node is a leaf and has a network object
 912                 if (__loc_database_node_is_leaf(n)) {
 913                         off_t network_index = be32toh(n->network);
 914
 915                         DEBUG(enumerator->ctx, "Node has a network at %jd\n", (intmax_t)network_index);
 916
 917                         // Fetch the network object
 918                         r = loc_database_fetch_network(enumerator->db, network,
 919                                 &enumerator->network_address, node->depth, network_index);
 920
 921                         // Break on any errors
 922                         if (r)
 923                                 return r;
 924
 925                         // Check if we are interested in this network
 926
 927                         // Skip if the country code does not match
 928                         if (*enumerator->country_code &&
 929                                         !loc_network_match_country_code(*network, enumerator->country_code)) {
 930                                 loc_network_unref(*network);
 931                                 *network = NULL;
 932
 933                                 continue;
 934                         }
 935
 936                         // Skip if the ASN does not match
 937                         if (enumerator->asn &&
 938                                         !loc_network_match_asn(*network, enumerator->asn)) {
 939                                 loc_network_unref(*network);
 940                                 *network = NULL;
 941
 942                                 continue;
 943                         }
 944
 945                         // Skip if flags do not match
 946                         if (enumerator->flags &&
 947                                         !loc_network_match_flag(*network, enumerator->flags)) {
 948                                 loc_network_unref(*network);
 949                                 *network = NULL;
 950                         }
 951
 952                         return 0;
 953                 }
 954         }
 955
 956         // Reached the end of the search
 957
 958         // Mark all nodes as non-visited
 959         for (unsigned int i = 0; i < enumerator->db->network_nodes_count; i++)
 960                 enumerator->networks_visited[i] = 0;
 961
 962         return 0;
 963 }