location-downloader: Add man page
[people/ms/libloc.git] / 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 }