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                 return -EINVAL;
 472
 473         DEBUG(db->ctx, "Fetching network at position %jd\n", (intmax_t)pos);
 474
 475         int r;
 476         switch (db->version) {
 477                 case 0:
 478                         r = loc_network_new_from_database_v0(db->ctx, network,
 479                                 address, prefix, db->networks_v0 + pos);
 480                         break;
 481
 482                 default:
 483                         return -1;
 484         }
 485
 486         if (r == 0) {
 487                 char* string = loc_network_str(*network);
 488                 DEBUG(db->ctx, "Got network %s\n", string);
 489                 free(string);
 490         }
 491
 492         return r;
 493 }
 494
 495 static int __loc_database_node_is_leaf(const struct loc_database_network_node_v0* node) {
 496         return (node->network != htobe32(0xffffffff));
 497 }
 498
 499 static int __loc_database_lookup_handle_leaf(struct loc_database* db, const struct in6_addr* address,
 500                 struct loc_network** network, struct in6_addr* network_address, unsigned int prefix,
 501                 const struct loc_database_network_node_v0* node) {
 502         off_t network_index = be32toh(node->network);
 503
 504         DEBUG(db->ctx, "Handling leaf node at %jd (%jd)\n", (intmax_t)(node - db->network_nodes_v0), (intmax_t)network_index);
 505
 506         // Fetch the network
 507         int r = loc_database_fetch_network(db, network,
 508                 network_address, prefix, network_index);
 509         if (r) {
 510                 ERROR(db->ctx, "Could not fetch network %jd from database\n", (intmax_t)network_index);
 511                 return r;
 512         }
 513
 514         // Check if the given IP address is inside the network
 515         r = loc_network_match_address(*network, address);
 516         if (r) {
 517                 DEBUG(db->ctx, "Searched address is not part of the network\n");
 518
 519                 loc_network_unref(*network);
 520                 *network = NULL;
 521                 return 1;
 522         }
 523
 524         // A network was found and the IP address matches
 525         return 0;
 526 }
 527
 528 // Searches for an exact match along the path
 529 static int __loc_database_lookup(struct loc_database* db, const struct in6_addr* address,
 530                 struct loc_network** network, struct in6_addr* network_address,
 531                 const struct loc_database_network_node_v0* node, unsigned int level) {
 532         int r;
 533         off_t node_index;
 534
 535         // Follow the path
 536         int bit = in6_addr_get_bit(address, level);
 537         in6_addr_set_bit(network_address, level, bit);
 538
 539         if (bit == 0)
 540                 node_index = be32toh(node->zero);
 541         else
 542                 node_index = be32toh(node->one);
 543
 544         // If the node index is zero, the tree ends here
 545         // and we cannot descend any further
 546         if (node_index > 0) {
 547                 // Check boundaries
 548                 if ((size_t)node_index >= db->network_nodes_count)
 549                         return -EINVAL;
 550
 551                 // Move on to the next node
 552                 r = __loc_database_lookup(db, address, network, network_address,
 553                         db->network_nodes_v0 + node_index, level + 1);
 554
 555                 // End here if a result was found
 556                 if (r == 0)
 557                         return r;
 558
 559                 // Raise any errors
 560                 else if (r < 0)
 561                         return r;
 562
 563                 DEBUG(db->ctx, "No match found below level %u\n", level);
 564         } else {
 565                 DEBUG(db->ctx, "Tree ended at level %u\n", level);
 566         }
 567
 568         // If this node has a leaf, we will check if it matches
 569         if (__loc_database_node_is_leaf(node)) {
 570                 r = __loc_database_lookup_handle_leaf(db, address, network, network_address, level, node);
 571                 if (r <= 0)
 572                         return r;
 573         }
 574
 575         return 1;
 576 }
 577
 578 LOC_EXPORT int loc_database_lookup(struct loc_database* db,
 579                 struct in6_addr* address, struct loc_network** network) {
 580         struct in6_addr network_address;
 581         memset(&network_address, 0, sizeof(network_address));
 582
 583         *network = NULL;
 584
 585         // Save start time
 586         clock_t start = clock();
 587
 588         int r = __loc_database_lookup(db, address, network, &network_address,
 589                 db->network_nodes_v0, 0);
 590
 591         clock_t end = clock();
 592
 593         // Log how fast this has been
 594         DEBUG(db->ctx, "Executed network search in %.4fms\n",
 595                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 596
 597         return r;
 598 }
 599
 600 LOC_EXPORT int loc_database_lookup_from_string(struct loc_database* db,
 601                 const char* string, struct loc_network** network) {
 602         struct in6_addr address;
 603
 604         int r = loc_parse_address(db->ctx, string, &address);
 605         if (r)
 606                 return r;
 607
 608         return loc_database_lookup(db, &address, network);
 609 }
 610
 611 // Returns the country at position pos
 612 static int loc_database_fetch_country(struct loc_database* db,
 613                 struct loc_country** country, off_t pos) {
 614         if ((size_t)pos >= db->countries_count)
 615                 return -EINVAL;
 616
 617         DEBUG(db->ctx, "Fetching country at position %jd\n", (intmax_t)pos);
 618
 619         int r;
 620         switch (db->version) {
 621                 case 0:
 622                         r = loc_country_new_from_database_v0(db->ctx, db->pool, country, db->countries_v0 + pos);
 623                         break;
 624
 625                 default:
 626                         return -1;
 627         }
 628
 629         if (r == 0) {
 630                 DEBUG(db->ctx, "Got country %s\n", loc_country_get_code(*country));
 631         }
 632
 633         return r;
 634 }
 635
 636 // Performs a binary search to find the country in the list
 637 LOC_EXPORT int loc_database_get_country(struct loc_database* db,
 638                 struct loc_country** country, const char* code) {
 639         off_t lo = 0;
 640         off_t hi = db->countries_count - 1;
 641
 642         // Save start time
 643         clock_t start = clock();
 644
 645         while (lo <= hi) {
 646                 off_t i = (lo + hi) / 2;
 647
 648                 // Fetch country in the middle between lo and hi
 649                 int r = loc_database_fetch_country(db, country, i);
 650                 if (r)
 651                         return r;
 652
 653                 // Check if this is a match
 654                 const char* cc = loc_country_get_code(*country);
 655                 int result = strcmp(code, cc);
 656
 657                 if (result == 0) {
 658                         clock_t end = clock();
 659
 660                         // Log how fast this has been
 661                         DEBUG(db->ctx, "Found country %s in %.4fms\n", cc,
 662                                 (double)(end - start) / CLOCKS_PER_SEC * 1000);
 663
 664                         return 0;
 665                 }
 666
 667                 // If it wasn't, we release the country and
 668                 // adjust our search pointers
 669                 loc_country_unref(*country);
 670
 671                 if (result > 0) {
 672                         lo = i + 1;
 673                 } else
 674                         hi = i - 1;
 675         }
 676
 677         // Nothing found
 678         *country = NULL;
 679
 680         return 1;
 681 }
 682
 683 // Enumerator
 684
 685 LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator,
 686                 struct loc_database* db, enum loc_database_enumerator_mode mode) {
 687         struct loc_database_enumerator* e = calloc(1, sizeof(*e));
 688         if (!e)
 689                 return -ENOMEM;
 690
 691         // Reference context
 692         e->ctx = loc_ref(db->ctx);
 693         e->db = loc_database_ref(db);
 694         e->mode = mode;
 695         e->refcount = 1;
 696
 697         // Initialise graph search
 698         //e->network_stack[++e->network_stack_depth] = 0;
 699         e->network_stack_depth = 1;
 700         e->networks_visited = calloc(db->network_nodes_count, sizeof(*e->networks_visited));
 701
 702         DEBUG(e->ctx, "Database enumerator object allocated at %p\n", e);
 703
 704         *enumerator = e;
 705         return 0;
 706 }
 707
 708 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_ref(struct loc_database_enumerator* enumerator) {
 709         enumerator->refcount++;
 710
 711         return enumerator;
 712 }
 713
 714 static void loc_database_enumerator_free(struct loc_database_enumerator* enumerator) {
 715         DEBUG(enumerator->ctx, "Releasing database enumerator %p\n", enumerator);
 716
 717         // Release all references
 718         loc_database_unref(enumerator->db);
 719         loc_unref(enumerator->ctx);
 720
 721         if (enumerator->string)
 722                 free(enumerator->string);
 723
 724         // Free network search
 725         free(enumerator->networks_visited);
 726
 727         free(enumerator);
 728 }
 729
 730 LOC_EXPORT struct loc_database_enumerator* loc_database_enumerator_unref(struct loc_database_enumerator* enumerator) {
 731         if (!enumerator)
 732                 return NULL;
 733
 734         if (--enumerator->refcount > 0)
 735                 return enumerator;
 736
 737         loc_database_enumerator_free(enumerator);
 738         return NULL;
 739 }
 740
 741 LOC_EXPORT int loc_database_enumerator_set_string(struct loc_database_enumerator* enumerator, const char* string) {
 742         enumerator->string = strdup(string);
 743
 744         // Make the string lowercase
 745         for (char *p = enumerator->string; *p; p++)
 746                 *p = tolower(*p);
 747
 748         return 0;
 749 }
 750
 751 LOC_EXPORT int loc_database_enumerator_set_country_code(struct loc_database_enumerator* enumerator, const char* country_code) {
 752         // Set empty country code
 753         if (!country_code || !*country_code) {
 754                 *enumerator->country_code = '\0';
 755                 return 0;
 756         }
 757
 758         // Treat A1, A2, A3 as special country codes,
 759         // but perform search for flags instead
 760         if (strcmp(country_code, "A1") == 0) {
 761                 return loc_database_enumerator_set_flag(enumerator,
 762                         LOC_NETWORK_FLAG_ANONYMOUS_PROXY);
 763         } else if (strcmp(country_code, "A2") == 0) {
 764                 return loc_database_enumerator_set_flag(enumerator,
 765                         LOC_NETWORK_FLAG_SATELLITE_PROVIDER);
 766         } else if (strcmp(country_code, "A3") == 0) {
 767                 return loc_database_enumerator_set_flag(enumerator,
 768                         LOC_NETWORK_FLAG_ANYCAST);
 769         }
 770
 771         // Country codes must be two characters
 772         if (!loc_country_code_is_valid(country_code))
 773                 return -EINVAL;
 774
 775         for (unsigned int i = 0; i < 3; i++) {
 776                 enumerator->country_code[i] = country_code[i];
 777         }
 778
 779         return 0;
 780 }
 781
 782 LOC_EXPORT int loc_database_enumerator_set_asn(
 783                 struct loc_database_enumerator* enumerator, unsigned int asn) {
 784         enumerator->asn = asn;
 785
 786         return 0;
 787 }
 788
 789 LOC_EXPORT int loc_database_enumerator_set_flag(
 790                 struct loc_database_enumerator* enumerator, enum loc_network_flags flag) {
 791         enumerator->flags |= flag;
 792
 793         return 0;
 794 }
 795
 796 LOC_EXPORT int loc_database_enumerator_next_as(
 797                 struct loc_database_enumerator* enumerator, struct loc_as** as) {
 798         *as = NULL;
 799
 800         // Do not do anything if not in AS mode
 801         if (enumerator->mode != LOC_DB_ENUMERATE_ASES)
 802                 return 0;
 803
 804         struct loc_database* db = enumerator->db;
 805
 806         while (enumerator->as_index < db->as_count) {
 807                 // Fetch the next AS
 808                 int r = loc_database_fetch_as(db, as, enumerator->as_index++);
 809                 if (r)
 810                         return r;
 811
 812                 r = loc_as_match_string(*as, enumerator->string);
 813                 if (r == 1) {
 814                         DEBUG(enumerator->ctx, "AS%d (%s) matches %s\n",
 815                                 loc_as_get_number(*as), loc_as_get_name(*as), enumerator->string);
 816
 817                         return 0;
 818                 }
 819
 820                 // No match
 821                 loc_as_unref(*as);
 822                 *as = NULL;
 823         }
 824
 825         // Reset the index
 826         enumerator->as_index = 0;
 827
 828         // We have searched through all of them
 829         return 0;
 830 }
 831
 832 static int loc_database_enumerator_stack_push_node(
 833                 struct loc_database_enumerator* e, off_t offset, int i, int depth) {
 834         // Do not add empty nodes
 835         if (!offset)
 836                 return 0;
 837
 838         // Check if there is any space left on the stack
 839         if (e->network_stack_depth >= MAX_STACK_DEPTH) {
 840                 ERROR(e->ctx, "Maximum stack size reached: %d\n", e->network_stack_depth);
 841                 return -1;
 842         }
 843
 844         // Increase stack size
 845         int s = ++e->network_stack_depth;
 846
 847         DEBUG(e->ctx, "Added node %jd to stack (%d)\n", (intmax_t)offset, depth);
 848
 849         e->network_stack[s].offset = offset;
 850         e->network_stack[s].i = i;
 851         e->network_stack[s].depth = depth;
 852
 853         return 0;
 854 }
 855
 856 LOC_EXPORT int loc_database_enumerator_next_network(
 857                 struct loc_database_enumerator* enumerator, struct loc_network** network) {
 858         // Reset network
 859         *network = NULL;
 860
 861         // Do not do anything if not in network mode
 862         if (enumerator->mode != LOC_DB_ENUMERATE_NETWORKS)
 863                 return 0;
 864
 865         int r;
 866
 867         DEBUG(enumerator->ctx, "Called with a stack of %u nodes\n",
 868                 enumerator->network_stack_depth);
 869
 870         // Perform DFS
 871         while (enumerator->network_stack_depth > 0) {
 872                 DEBUG(enumerator->ctx, "Stack depth: %u\n", enumerator->network_stack_depth);
 873
 874                 // Get object from top of the stack
 875                 struct loc_node_stack* node = &enumerator->network_stack[enumerator->network_stack_depth];
 876
 877                 // Remove the node from the stack if we have already visited it
 878                 if (enumerator->networks_visited[node->offset]) {
 879                         enumerator->network_stack_depth--;
 880                         continue;
 881                 }
 882
 883                 // Mark the bits on the path correctly
 884                 in6_addr_set_bit(&enumerator->network_address,
 885                         (node->depth > 0) ? node->depth - 1 : 0, node->i);
 886
 887                 DEBUG(enumerator->ctx, "Looking at node %jd\n", (intmax_t)node->offset);
 888                 enumerator->networks_visited[node->offset]++;
 889
 890                 // Pop node from top of the stack
 891                 struct loc_database_network_node_v0* n =
 892                         enumerator->db->network_nodes_v0 + node->offset;
 893
 894                 // Add edges to stack
 895                 r = loc_database_enumerator_stack_push_node(enumerator,
 896                         be32toh(n->one), 1, node->depth + 1);
 897
 898                 if (r)
 899                         return r;
 900
 901                 r = loc_database_enumerator_stack_push_node(enumerator,
 902                         be32toh(n->zero), 0, node->depth + 1);
 903
 904                 if (r)
 905                         return r;
 906
 907                 // Check if this node is a leaf and has a network object
 908                 if (__loc_database_node_is_leaf(n)) {
 909                         off_t network_index = be32toh(n->network);
 910
 911                         DEBUG(enumerator->ctx, "Node has a network at %jd\n", (intmax_t)network_index);
 912
 913                         // Fetch the network object
 914                         r = loc_database_fetch_network(enumerator->db, network,
 915                                 &enumerator->network_address, node->depth, network_index);
 916
 917                         // Break on any errors
 918                         if (r)
 919                                 return r;
 920
 921                         // Check if we are interested in this network
 922
 923                         // Skip if the country code does not match
 924                         if (*enumerator->country_code &&
 925                                         !loc_network_match_country_code(*network, enumerator->country_code)) {
 926                                 loc_network_unref(*network);
 927                                 *network = NULL;
 928
 929                                 continue;
 930                         }
 931
 932                         // Skip if the ASN does not match
 933                         if (enumerator->asn &&
 934                                         !loc_network_match_asn(*network, enumerator->asn)) {
 935                                 loc_network_unref(*network);
 936                                 *network = NULL;
 937
 938                                 continue;
 939                         }
 940
 941                         // Skip if flags do not match
 942                         if (enumerator->flags &&
 943                                         !loc_network_match_flag(*network, enumerator->flags)) {
 944                                 loc_network_unref(*network);
 945                                 *network = NULL;
 946                         }
 947
 948                         return 0;
 949                 }
 950         }
 951
 952         // Reached the end of the search
 953
 954         // Mark all nodes as non-visited
 955         for (unsigned int i = 0; i < enumerator->db->network_nodes_count; i++)
 956                 enumerator->networks_visited[i] = 0;
 957
 958         return 0;
 959 }