X-Git-Url: http://git.ipfire.org/?p=people%2Fms%2Flibloc.git;a=blobdiff_plain;f=src%2Fdatabase.c;h=862fe75dffdf65a1dedd6be990af1f755c5a85ca;hp=fe5e6dc6e0ebcb86a56b5629dc982ff5bc879ff1;hb=191830dadd6e8528e3a62c9c2e2e22e72dd30248;hpb=273948cfa02a1bacf906f1b4dc296f81b7d8441e diff --git a/src/database.c b/src/database.c index fe5e6dc..862fe75 100644 --- a/src/database.c +++ b/src/database.c @@ -31,6 +31,7 @@ #include #include +#include #include #include #include @@ -59,19 +60,41 @@ struct loc_database { struct loc_database_network_v0* networks_v0; size_t networks_count; + // Countries + struct loc_database_country_v0* countries_v0; + size_t countries_count; + struct loc_stringpool* pool; }; +#define MAX_STACK_DEPTH 256 + +struct loc_node_stack { + off_t offset; + int i; // Is this node 0 or 1? + int depth; +}; + struct loc_database_enumerator { struct loc_ctx* ctx; struct loc_database* db; + enum loc_database_enumerator_mode mode; int refcount; // Search string char* string; + char country_code[3]; + uint32_t asn; + enum loc_network_flags flags; // Index of the AS we are looking at unsigned int as_index; + + // Network state + struct in6_addr network_address; + struct loc_node_stack network_stack[MAX_STACK_DEPTH]; + int network_stack_depth; + unsigned int* networks_visited; }; static int loc_database_read_magic(struct loc_database* db, FILE* f) { @@ -172,6 +195,30 @@ static int loc_database_read_networks_section_v0(struct loc_database* db, return 0; } +static int loc_database_read_countries_section_v0(struct loc_database* db, + FILE* f, const struct loc_database_header_v0* header) { + off_t countries_offset = be32toh(header->countries_offset); + size_t countries_length = be32toh(header->countries_length); + + DEBUG(db->ctx, "Reading countries section from %jd (%zu bytes)\n", + countries_offset, countries_length); + + if (countries_length > 0) { + db->countries_v0 = mmap(NULL, countries_length, PROT_READ, + MAP_SHARED, fileno(f), countries_offset); + + if (db->countries_v0 == MAP_FAILED) + return -errno; + } + + db->countries_count = countries_length / sizeof(*db->countries_v0); + + INFO(db->ctx, "Read %zu countries from the database\n", + db->countries_count); + + return 0; +} + static int loc_database_read_header_v0(struct loc_database* db, FILE* f) { struct loc_database_header_v0 header; @@ -213,6 +260,11 @@ static int loc_database_read_header_v0(struct loc_database* db, FILE* f) { if (r) return r; + // countries + r = loc_database_read_countries_section_v0(db, f, &header); + if (r) + return r; + return 0; } @@ -242,8 +294,8 @@ static int loc_database_read(struct loc_database* db, FILE* f) { clock_t end = clock(); - INFO(db->ctx, "Opened database in %.8fs\n", - (double)(end - start) / CLOCKS_PER_SEC); + INFO(db->ctx, "Opened database in %.4fms\n", + (double)(end - start) / CLOCKS_PER_SEC * 1000); return 0; } @@ -386,8 +438,8 @@ LOC_EXPORT int loc_database_get_as(struct loc_database* db, struct loc_as** as, clock_t end = clock(); // Log how fast this has been - DEBUG(db->ctx, "Found AS%u in %.8fs\n", as_number, - (double)(end - start) / CLOCKS_PER_SEC); + DEBUG(db->ctx, "Found AS%u in %.4fms\n", as_number, + (double)(end - start) / CLOCKS_PER_SEC * 1000); return 0; } @@ -535,8 +587,8 @@ LOC_EXPORT int loc_database_lookup(struct loc_database* db, clock_t end = clock(); // Log how fast this has been - DEBUG(db->ctx, "Executed network search in %.8fs\n", - (double)(end - start) / CLOCKS_PER_SEC); + DEBUG(db->ctx, "Executed network search in %.4fms\n", + (double)(end - start) / CLOCKS_PER_SEC * 1000); return r; } @@ -552,9 +604,82 @@ LOC_EXPORT int loc_database_lookup_from_string(struct loc_database* db, return loc_database_lookup(db, &address, network); } +// Returns the country at position pos +static int loc_database_fetch_country(struct loc_database* db, + struct loc_country** country, off_t pos) { + if ((size_t)pos >= db->countries_count) + return -EINVAL; + + DEBUG(db->ctx, "Fetching country at position %jd\n", pos); + + int r; + switch (db->version) { + case 0: + r = loc_country_new_from_database_v0(db->ctx, db->pool, country, db->countries_v0 + pos); + break; + + default: + return -1; + } + + if (r == 0) { + DEBUG(db->ctx, "Got country %s\n", loc_country_get_code(*country)); + } + + return r; +} + +// Performs a binary search to find the country in the list +LOC_EXPORT int loc_database_get_country(struct loc_database* db, + struct loc_country** country, const char* code) { + off_t lo = 0; + off_t hi = db->countries_count - 1; + + // Save start time + clock_t start = clock(); + + while (lo <= hi) { + off_t i = (lo + hi) / 2; + + // Fetch country in the middle between lo and hi + int r = loc_database_fetch_country(db, country, i); + if (r) + return r; + + // Check if this is a match + const char* cc = loc_country_get_code(*country); + int result = strcmp(code, cc); + + if (result == 0) { + clock_t end = clock(); + + // Log how fast this has been + DEBUG(db->ctx, "Found country %s in %.4fms\n", cc, + (double)(end - start) / CLOCKS_PER_SEC * 1000); + + return 0; + } + + // If it wasn't, we release the country and + // adjust our search pointers + loc_country_unref(*country); + + if (result > 0) { + lo = i + 1; + } else + hi = i - 1; + } + + // Nothing found + *country = NULL; + + return 1; +} + // Enumerator -LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator, struct loc_database* db) { +LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enumerator, + struct loc_database* db, enum loc_database_enumerator_mode mode) { struct loc_database_enumerator* e = calloc(1, sizeof(*e)); if (!e) return -ENOMEM; @@ -562,8 +687,14 @@ LOC_EXPORT int loc_database_enumerator_new(struct loc_database_enumerator** enum // Reference context e->ctx = loc_ref(db->ctx); e->db = loc_database_ref(db); + e->mode = mode; e->refcount = 1; + // Initialise graph search + //e->network_stack[++e->network_stack_depth] = 0; + e->network_stack_depth = 1; + e->networks_visited = calloc(db->network_nodes_count, sizeof(*e->networks_visited)); + DEBUG(e->ctx, "Database enumerator object allocated at %p\n", e); *enumerator = e; @@ -586,6 +717,9 @@ static void loc_database_enumerator_free(struct loc_database_enumerator* enumera if (enumerator->string) free(enumerator->string); + // Free network search + free(enumerator->networks_visited); + free(enumerator); } @@ -610,31 +744,212 @@ LOC_EXPORT int loc_database_enumerator_set_string(struct loc_database_enumerator return 0; } -LOC_EXPORT struct loc_as* loc_database_enumerator_next_as(struct loc_database_enumerator* enumerator) { +LOC_EXPORT int loc_database_enumerator_set_country_code(struct loc_database_enumerator* enumerator, const char* country_code) { + // Set empty country code + if (!country_code || !*country_code) { + *enumerator->country_code = '\0'; + return 0; + } + + // Treat A1, A2, A3 as special country codes, + // but perform search for flags instead + if (strcmp(country_code, "A1") == 0) { + return loc_database_enumerator_set_flag(enumerator, + LOC_NETWORK_FLAG_ANONYMOUS_PROXY); + } else if (strcmp(country_code, "A2") == 0) { + return loc_database_enumerator_set_flag(enumerator, + LOC_NETWORK_FLAG_SATELLITE_PROVIDER); + } else if (strcmp(country_code, "A3") == 0) { + return loc_database_enumerator_set_flag(enumerator, + LOC_NETWORK_FLAG_ANYCAST); + } + + // Country codes must be two characters + if (!loc_country_code_is_valid(country_code)) + return -EINVAL; + + for (unsigned int i = 0; i < 3; i++) { + enumerator->country_code[i] = country_code[i]; + } + + return 0; +} + +LOC_EXPORT int loc_database_enumerator_set_asn( + struct loc_database_enumerator* enumerator, unsigned int asn) { + enumerator->asn = asn; + + return 0; +} + +LOC_EXPORT int loc_database_enumerator_set_flag( + struct loc_database_enumerator* enumerator, enum loc_network_flags flag) { + enumerator->flags |= flag; + + return 0; +} + +LOC_EXPORT int loc_database_enumerator_next_as( + struct loc_database_enumerator* enumerator, struct loc_as** as) { + *as = NULL; + + // Do not do anything if not in AS mode + if (enumerator->mode != LOC_DB_ENUMERATE_ASES) + return 0; + struct loc_database* db = enumerator->db; - struct loc_as* as; while (enumerator->as_index < db->as_count) { // Fetch the next AS - int r = loc_database_fetch_as(db, &as, enumerator->as_index++); + int r = loc_database_fetch_as(db, as, enumerator->as_index++); if (r) - return NULL; + return r; - r = loc_as_match_string(as, enumerator->string); + r = loc_as_match_string(*as, enumerator->string); if (r == 1) { DEBUG(enumerator->ctx, "AS%d (%s) matches %s\n", - loc_as_get_number(as), loc_as_get_name(as), enumerator->string); + loc_as_get_number(*as), loc_as_get_name(*as), enumerator->string); - return as; + return 0; } // No match - loc_as_unref(as); + loc_as_unref(*as); + *as = NULL; } // Reset the index enumerator->as_index = 0; // We have searched through all of them - return NULL; + return 0; +} + +static int loc_database_enumerator_stack_push_node( + struct loc_database_enumerator* e, off_t offset, int i, int depth) { + // Do not add empty nodes + if (!offset) + return 0; + + // Check if there is any space left on the stack + if (e->network_stack_depth >= MAX_STACK_DEPTH) { + ERROR(e->ctx, "Maximum stack size reached: %d\n", e->network_stack_depth); + return -1; + } + + // Increase stack size + int s = ++e->network_stack_depth; + + DEBUG(e->ctx, "Added node %jd to stack (%d)\n", offset, depth); + + e->network_stack[s].offset = offset; + e->network_stack[s].i = i; + e->network_stack[s].depth = depth; + + return 0; +} + +LOC_EXPORT int loc_database_enumerator_next_network( + struct loc_database_enumerator* enumerator, struct loc_network** network) { + // Reset network + *network = NULL; + + // Do not do anything if not in network mode + if (enumerator->mode != LOC_DB_ENUMERATE_NETWORKS) + return 0; + + int r; + + DEBUG(enumerator->ctx, "Called with a stack of %u nodes\n", + enumerator->network_stack_depth); + + // Perform DFS + while (enumerator->network_stack_depth > 0) { + DEBUG(enumerator->ctx, "Stack depth: %u\n", enumerator->network_stack_depth); + + // Get object from top of the stack + struct loc_node_stack* node = &enumerator->network_stack[enumerator->network_stack_depth]; + + // Remove the node from the stack if we have already visited it + if (enumerator->networks_visited[node->offset]) { + enumerator->network_stack_depth--; + continue; + } + + // Mark the bits on the path correctly + in6_addr_set_bit(&enumerator->network_address, + (node->depth > 0) ? node->depth - 1 : 0, node->i); + + DEBUG(enumerator->ctx, "Looking at node %jd\n", node->offset); + enumerator->networks_visited[node->offset]++; + + // Pop node from top of the stack + struct loc_database_network_node_v0* n = + enumerator->db->network_nodes_v0 + node->offset; + + // Add edges to stack + r = loc_database_enumerator_stack_push_node(enumerator, + be32toh(n->one), 1, node->depth + 1); + + if (r) + return r; + + r = loc_database_enumerator_stack_push_node(enumerator, + be32toh(n->zero), 0, node->depth + 1); + + if (r) + return r; + + // Check if this node is a leaf and has a network object + if (__loc_database_node_is_leaf(n)) { + off_t network_index = be32toh(n->network); + + DEBUG(enumerator->ctx, "Node has a network at %jd\n", network_index); + + // Fetch the network object + r = loc_database_fetch_network(enumerator->db, network, + &enumerator->network_address, node->depth, network_index); + + // Break on any errors + if (r) + return r; + + // Check if we are interested in this network + + // Skip if the country code does not match + if (enumerator->country_code && + !loc_network_match_country_code(*network, enumerator->country_code)) { + loc_network_unref(*network); + *network = NULL; + + continue; + } + + // Skip if the ASN does not match + if (enumerator->asn && + !loc_network_match_asn(*network, enumerator->asn)) { + loc_network_unref(*network); + *network = NULL; + + continue; + } + + // Skip if flags do not match + if (enumerator->flags && + !loc_network_match_flag(*network, enumerator->flags)) { + loc_network_unref(*network); + *network = NULL; + } + + return 0; + } + } + + // Reached the end of the search + + // Mark all nodes as non-visited + for (unsigned int i = 0; i < enumerator->db->network_nodes_count; i++) + enumerator->networks_visited[i] = 0; + + return 0; }