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