From 06bc0b1642b81927f4135eb3cb5ef47e1fc3a916 Mon Sep 17 00:00:00 2001 From: Nick Peng Date: Tue, 26 Jun 2018 00:09:11 +0800 Subject: [PATCH] Add address feature --- src/Makefile | 2 +- src/conf.c | 38 +- src/conf.h | 5 +- src/dns_server.c | 58 ++- src/include/art.h | 245 ++++++++++ src/lib/art.c | 990 +++++++++++++++++++++++++++++++++++++++ src/smartdns.c | 53 +-- src/util.c | 18 +- src/util.h | 2 + systemd/smartdns.service | 5 +- 10 files changed, 1343 insertions(+), 73 deletions(-) create mode 100644 src/include/art.h create mode 100644 src/lib/art.c diff --git a/src/Makefile b/src/Makefile index 94012f7..2a8b166 100644 --- a/src/Makefile +++ b/src/Makefile @@ -1,6 +1,6 @@ BIN=smartdns -OBJS=smartdns.o fast_ping.o lib/bitops.o dns_client.o dns_server.o dns.o util.o tlog.o conf.o lib/rbtree.o +OBJS=smartdns.o fast_ping.o lib/bitops.o dns_client.o dns_server.o dns.o util.o tlog.o conf.o lib/rbtree.o lib/art.o CFLAGS=-g -O0 -Wall -Wstrict-prototypes -fno-omit-frame-pointer -Wstrict-aliasing CFLAGS +=-Iinclude CFLAGS += -DBASE_FILE_NAME=\"$(notdir $<)\" diff --git a/src/conf.c b/src/conf.c index 53682e9..6596956 100644 --- a/src/conf.c +++ b/src/conf.c @@ -18,7 +18,7 @@ int dns_conf_cachesize = DEFAULT_DNS_CACHE_SIZE; struct dns_servers dns_conf_servers[DNS_MAX_SERVERS]; int dns_conf_server_num; int dns_conf_loglevel = TLOG_ERROR; -LIST_HEAD(dns_conf_address_list); +art_tree dns_conf_address; int config_bind(char *value) { @@ -57,15 +57,29 @@ int config_server(char *value, dns_conf_server_type_t type) return 0; } +int config_address_iter_cb(void *data, const unsigned char *key, uint32_t key_len, void *value) +{ + free(value); + return 0; +} + +void config_address_destroy(void) +{ + art_iter(&dns_conf_address, config_address_iter_cb, 0); + art_tree_destroy(&dns_conf_address); +} + int config_address(char *value) { struct dns_address *address; char ip[MAX_IP_LEN]; + char domain_key[DNS_MAX_CONF_CNAME_LEN]; char *begin = NULL; char *end = NULL; int len = 0; struct sockaddr_storage addr; socklen_t addr_len = sizeof(addr); + char type = '4'; begin = strstr(value, "/"); if (begin == NULL) { @@ -87,6 +101,7 @@ int config_address(char *value) memcpy(address->domain, begin, len); address->domain[len] = 0; strncpy(ip, end + 1, MAX_IP_LEN); + reverse_string(domain_key + 1, address->domain, len); if (getaddr_by_host(ip, (struct sockaddr *)&addr, &addr_len) != 0) { goto errout; @@ -98,6 +113,7 @@ int config_address(char *value) addr_in = (struct sockaddr_in *)&addr; memcpy(address->ipv4_addr, &addr_in->sin_addr.s_addr, 4); address->addr_type = DNS_T_A; + type = '4'; } break; case AF_INET6: { struct sockaddr_in6 *addr_in6; @@ -105,16 +121,20 @@ int config_address(char *value) if (IN6_IS_ADDR_V4MAPPED(&addr_in6->sin6_addr)) { memcpy(address->ipv4_addr, addr_in6->sin6_addr.s6_addr + 12, 4); address->addr_type = DNS_T_A; + type = '4'; } else { memcpy(address->ipv6_addr, addr_in6->sin6_addr.s6_addr, 16); address->addr_type = DNS_T_AAAA; + type = '6'; } } break; default: goto errout; } - list_add_tail(&address->list, &dns_conf_address_list); + domain_key[0] = type; + len++; + art_insert(&dns_conf_address, (unsigned char *)domain_key, len, address); return 0; errout: @@ -186,6 +206,18 @@ struct config_item config_item[] = { }; int config_item_num = sizeof(config_item) / sizeof(struct config_item); +int load_conf_init(void) +{ + art_tree_init(&dns_conf_address); + + return 0; +} + +void load_exit(void) +{ + config_address_destroy(); +} + int load_conf(const char *file) { FILE *fp = NULL; @@ -196,6 +228,8 @@ int load_conf(const char *file) int line_num = 0; int i; + load_conf_init(); + fp = fopen(file, "r"); if (fp == NULL) { tlog(TLOG_ERROR, "config file %s not exist.", file); diff --git a/src/conf.h b/src/conf.h index 89cc916..0bf4b4e 100644 --- a/src/conf.h +++ b/src/conf.h @@ -2,6 +2,7 @@ #define _DNS_CONF #include "list.h" +#include "art.h" #include "dns.h" #define DNS_MAX_SERVERS 32 @@ -41,8 +42,10 @@ extern int dns_conf_verbose; extern int dns_conf_loglevel; extern char dns_conf_logfile[DNS_MAX_PATH]; extern int dns_conf_lognum; -extern struct list_head dns_conf_address_list; +extern art_tree dns_conf_address; int load_conf(const char *file); +void load_exit(void); + #endif // !_DNS_CONF \ No newline at end of file diff --git a/src/dns_server.c b/src/dns_server.c index 4bc2d4e..4053eb0 100644 --- a/src/dns_server.c +++ b/src/dns_server.c @@ -669,31 +669,59 @@ errout: return -1; } -static struct dns_address *_dns_server_get_address_by_domain(char *domain) +int _dns_server_art_iter_callback(void *data, const unsigned char *key, uint32_t key_len, void *value) { - struct dns_address *address; - char *match = NULL; - int domain_len; + struct dns_address **address; + address = data; + *address = value; + return 0; +} - list_for_each_entry(address, &dns_conf_address_list, list) - { - domain_len = strnlen(address->domain, DNS_MAX_CNAME_LEN); - match = strstr(domain, address->domain); - if (match) { - if (memcmp(address->domain, match, domain_len + 1) == 0) { - return address; - } - } +static int _dns_server_art_domain_cmp(const art_leaf *n, const unsigned char *prefix, int prefix_len) +{ + // Fail if the prefix length is too short + if (n->key_len > (uint32_t)prefix_len) { + return 1; } - return NULL; + // Compare the keys + return memcmp(n->key, prefix, n->key_len); +} + +static struct dns_address *_dns_server_get_address_by_domain(char *domain, int qtype) +{ + struct dns_address *address = NULL; + int domain_len; + char domain_key[DNS_MAX_CNAME_LEN]; + char type = '4'; + + switch(qtype) { + case DNS_T_A: + type = '4'; + break; + case DNS_T_AAAA: + type = '6'; + break; + default: + return NULL; + } + + domain_len = strlen(domain); + reverse_string(domain_key + 1, domain, domain_len); + domain_key[0] = type; + domain_len++; + if (art_iter_cmp(&dns_conf_address, (unsigned char *)domain_key, domain_len, _dns_server_art_iter_callback, _dns_server_art_domain_cmp, &address) != 0) { + return NULL; + } + + return address; } static int _dns_server_process_address(struct dns_request *request, struct dns_packet *packet) { struct dns_address *address = NULL; - address = _dns_server_get_address_by_domain(request->domain); + address = _dns_server_get_address_by_domain(request->domain, request->qtype); if (address == NULL) { goto errout; } diff --git a/src/include/art.h b/src/include/art.h new file mode 100644 index 0000000..184e643 --- /dev/null +++ b/src/include/art.h @@ -0,0 +1,245 @@ +/* +Copyright (c) 2012, Armon Dadgar +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the organization nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL ARMON DADGAR BE LIABLE FOR ANY +DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*/ +#include +#ifndef ART_H +#define ART_H + +#ifdef __cplusplus +extern "C" { +#endif + +#define NODE4 1 +#define NODE16 2 +#define NODE48 3 +#define NODE256 4 + +#define MAX_PREFIX_LEN 10 + +#if defined(__GNUC__) && !defined(__clang__) +# if __STDC_VERSION__ >= 199901L && 402 == (__GNUC__ * 100 + __GNUC_MINOR__) +/* + * GCC 4.2.2's C99 inline keyword support is pretty broken; avoid. Introduced in + * GCC 4.2.something, fixed in 4.3.0. So checking for specific major.minor of + * 4.2 is fine. + */ +# define BROKEN_GCC_C99_INLINE +# endif +#endif + +typedef int(*art_callback)(void *data, const unsigned char *key, uint32_t key_len, void *value); + +/** + * This struct is included as part + * of all the various node sizes + */ +typedef struct { + uint8_t type; + uint8_t num_children; + uint32_t partial_len; + unsigned char partial[MAX_PREFIX_LEN]; +} art_node; + +/** + * Small node with only 4 children + */ +typedef struct { + art_node n; + unsigned char keys[4]; + art_node *children[4]; +} art_node4; + +/** + * Node with 16 children + */ +typedef struct { + art_node n; + unsigned char keys[16]; + art_node *children[16]; +} art_node16; + +/** + * Node with 48 children, but + * a full 256 byte field. + */ +typedef struct { + art_node n; + unsigned char keys[256]; + art_node *children[48]; +} art_node48; + +/** + * Full node with 256 children + */ +typedef struct { + art_node n; + art_node *children[256]; +} art_node256; + +/** + * Represents a leaf. These are + * of arbitrary size, as they include the key. + */ +typedef struct { + void *value; + uint32_t key_len; + unsigned char key[]; +} art_leaf; + +/** + * Main struct, points to root. + */ +typedef struct { + art_node *root; + uint64_t size; +} art_tree; + +/** + * Initializes an ART tree + * @return 0 on success. + */ +int art_tree_init(art_tree *t); + +/** + * DEPRECATED + * Initializes an ART tree + * @return 0 on success. + */ +#define init_art_tree(...) art_tree_init(__VA_ARGS__) + +/** + * Destroys an ART tree + * @return 0 on success. + */ +int art_tree_destroy(art_tree *t); + +/** + * DEPRECATED + * Initializes an ART tree + * @return 0 on success. + */ +#define destroy_art_tree(...) art_tree_destroy(__VA_ARGS__) + +/** + * Returns the size of the ART tree. + */ +#ifdef BROKEN_GCC_C99_INLINE +# define art_size(t) ((t)->size) +#else +inline uint64_t art_size(art_tree *t) { + return t->size; +} +#endif + +/** + * Inserts a new value into the ART tree + * @arg t The tree + * @arg key The key + * @arg key_len The length of the key + * @arg value Opaque value. + * @return NULL if the item was newly inserted, otherwise + * the old value pointer is returned. + */ +void* art_insert(art_tree *t, const unsigned char *key, int key_len, void *value); + +/** + * Deletes a value from the ART tree + * @arg t The tree + * @arg key The key + * @arg key_len The length of the key + * @return NULL if the item was not found, otherwise + * the value pointer is returned. + */ +void* art_delete(art_tree *t, const unsigned char *key, int key_len); + +/** + * Searches for a value in the ART tree + * @arg t The tree + * @arg key The key + * @arg key_len The length of the key + * @return NULL if the item was not found, otherwise + * the value pointer is returned. + */ +void* art_search(const art_tree *t, const unsigned char *key, int key_len); + +/** + * Returns the minimum valued leaf + * @return The minimum leaf or NULL + */ +art_leaf* art_minimum(art_tree *t); + +/** + * Returns the maximum valued leaf + * @return The maximum leaf or NULL + */ +art_leaf* art_maximum(art_tree *t); + +/** + * Iterates through the entries pairs in the map, + * invoking a callback for each. The call back gets a + * key, value for each and returns an integer stop value. + * If the callback returns non-zero, then the iteration stops. + * @arg t The tree to iterate over + * @arg cb The callback function to invoke + * @arg data Opaque handle passed to the callback + * @return 0 on success, or the return of the callback. + */ +int art_iter(art_tree *t, art_callback cb, void *data); + +/** + * Iterates through the entries pairs in the map, + * invoking a callback for each that matches a given prefix. + * The call back gets a key, value for each and returns an integer stop value. + * If the callback returns non-zero, then the iteration stops. + * @arg t The tree to iterate over + * @arg prefix The prefix of keys to read + * @arg prefix_len The length of the prefix + * @arg cb The callback function to invoke + * @arg data Opaque handle passed to the callback + * @return 0 on success, or the return of the callback. + */ +int art_iter_prefix(art_tree *t, const unsigned char *prefix, int prefix_len, art_callback cb, void *data); + +/** + * Iterates through the entries pairs in the map, + * invoking a callback for each that matches a given prefix. + * The call back gets a key, value for each and returns an integer stop value. + * If the callback returns non-zero, then the iteration stops. + * @arg t The tree to iterate over + * @arg prefix The prefix of keys to read + * @arg prefix_len The length of the prefix + * @arg cb The callback function to invoke + * @arg data Opaque handle passed to the callback + * @return 0 on success, or the return of the callback. + */ +typedef int(*art_key_cmp_callback)(const art_leaf *n, const unsigned char *prefix, int prefix_len); +int art_iter_cmp(art_tree *t, const unsigned char *str, int str_len, art_callback cb, art_key_cmp_callback key_cmp, void *data); + +#ifdef __cplusplus +} +#endif + +#endif \ No newline at end of file diff --git a/src/lib/art.c b/src/lib/art.c new file mode 100644 index 0000000..60a51fe --- /dev/null +++ b/src/lib/art.c @@ -0,0 +1,990 @@ +/* +Copyright (c) 2012, Armon Dadgar +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the organization nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL ARMON DADGAR BE LIABLE FOR ANY +DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +*/ +#include +#include +#include +#include +#include +#include "art.h" + +#ifdef __i386__ + #include +#else +#ifdef __amd64__ + #include +#endif +#endif + +/** + * Macros to manipulate pointer tags + */ +#define IS_LEAF(x) (((uintptr_t)x & 1)) +#define SET_LEAF(x) ((void*)((uintptr_t)x | 1)) +#define LEAF_RAW(x) ((art_leaf*)((void*)((uintptr_t)x & ~1))) + +/** + * Allocates a node of the given type, + * initializes to zero and sets the type. + */ +static art_node* alloc_node(uint8_t type) { + art_node* n; + switch (type) { + case NODE4: + n = (art_node*)calloc(1, sizeof(art_node4)); + break; + case NODE16: + n = (art_node*)calloc(1, sizeof(art_node16)); + break; + case NODE48: + n = (art_node*)calloc(1, sizeof(art_node48)); + break; + case NODE256: + n = (art_node*)calloc(1, sizeof(art_node256)); + break; + default: + abort(); + } + n->type = type; + return n; +} + +/** + * Initializes an ART tree + * @return 0 on success. + */ +int art_tree_init(art_tree *t) { + t->root = NULL; + t->size = 0; + return 0; +} + +// Recursively destroys the tree +static void destroy_node(art_node *n) { + // Break if null + if (!n) return; + + // Special case leafs + if (IS_LEAF(n)) { + free(LEAF_RAW(n)); + return; + } + + // Handle each node type + int i, idx; + union { + art_node4 *p1; + art_node16 *p2; + art_node48 *p3; + art_node256 *p4; + } p; + switch (n->type) { + case NODE4: + p.p1 = (art_node4*)n; + for (i=0;inum_children;i++) { + destroy_node(p.p1->children[i]); + } + break; + + case NODE16: + p.p2 = (art_node16*)n; + for (i=0;inum_children;i++) { + destroy_node(p.p2->children[i]); + } + break; + + case NODE48: + p.p3 = (art_node48*)n; + for (i=0;i<256;i++) { + idx = ((art_node48*)n)->keys[i]; + if (!idx) continue; + destroy_node(p.p3->children[idx-1]); + } + break; + + case NODE256: + p.p4 = (art_node256*)n; + for (i=0;i<256;i++) { + if (p.p4->children[i]) + destroy_node(p.p4->children[i]); + } + break; + + default: + abort(); + } + + // Free ourself on the way up + free(n); +} + +/** + * Destroys an ART tree + * @return 0 on success. + */ +int art_tree_destroy(art_tree *t) { + destroy_node(t->root); + return 0; +} + +/** + * Returns the size of the ART tree. + */ + +#ifndef BROKEN_GCC_C99_INLINE +extern inline uint64_t art_size(art_tree *t); +#endif + +static art_node** find_child(art_node *n, unsigned char c) { + int i, mask, bitfield; + union { + art_node4 *p1; + art_node16 *p2; + art_node48 *p3; + art_node256 *p4; + } p; + switch (n->type) { + case NODE4: + p.p1 = (art_node4*)n; + for (i=0 ; i < n->num_children; i++) { + /* this cast works around a bug in gcc 5.1 when unrolling loops + * https://gcc.gnu.org/bugzilla/show_bug.cgi?id=59124 + */ + if (((unsigned char*)p.p1->keys)[i] == c) + return &p.p1->children[i]; + } + break; + + { + case NODE16: + p.p2 = (art_node16*)n; + + // support non-86 architectures + #ifdef __i386__ + // Compare the key to all 16 stored keys + __m128i cmp; + cmp = _mm_cmpeq_epi8(_mm_set1_epi8(c), + _mm_loadu_si128((__m128i*)p.p2->keys)); + + // Use a mask to ignore children that don't exist + mask = (1 << n->num_children) - 1; + bitfield = _mm_movemask_epi8(cmp) & mask; + #else + #ifdef __amd64__ + // Compare the key to all 16 stored keys + __m128i cmp; + cmp = _mm_cmpeq_epi8(_mm_set1_epi8(c), + _mm_loadu_si128((__m128i*)p.p2->keys)); + + // Use a mask to ignore children that don't exist + mask = (1 << n->num_children) - 1; + bitfield = _mm_movemask_epi8(cmp) & mask; + #else + // Compare the key to all 16 stored keys + bitfield = 0; + for (i = 0; i < 16; ++i) { + if (p.p2->keys[i] == c) + bitfield |= (1 << i); + } + + // Use a mask to ignore children that don't exist + mask = (1 << n->num_children) - 1; + bitfield &= mask; + #endif + #endif + + /* + * If we have a match (any bit set) then we can + * return the pointer match using ctz to get + * the index. + */ + if (bitfield) + return &p.p2->children[__builtin_ctz(bitfield)]; + break; + } + + case NODE48: + p.p3 = (art_node48*)n; + i = p.p3->keys[c]; + if (i) + return &p.p3->children[i-1]; + break; + + case NODE256: + p.p4 = (art_node256*)n; + if (p.p4->children[c]) + return &p.p4->children[c]; + break; + + default: + abort(); + } + return NULL; +} + +// Simple inlined if +static inline int min(int a, int b) { + return (a < b) ? a : b; +} + +/** + * Returns the number of prefix characters shared between + * the key and node. + */ +static int check_prefix(const art_node *n, const unsigned char *key, int key_len, int depth) { + int max_cmp = min(min(n->partial_len, MAX_PREFIX_LEN), key_len - depth); + int idx; + for (idx=0; idx < max_cmp; idx++) { + if (n->partial[idx] != key[depth+idx]) + return idx; + } + return idx; +} + +/** + * Checks if a leaf matches + * @return 0 on success. + */ +static int leaf_matches(const art_leaf *n, const unsigned char *key, int key_len, int depth) { + (void)depth; + // Fail if the key lengths are different + if (n->key_len != (uint32_t)key_len) return 1; + + // Compare the keys starting at the depth + return memcmp(n->key, key, key_len); +} + +/** + * Searches for a value in the ART tree + * @arg t The tree + * @arg key The key + * @arg key_len The length of the key + * @return NULL if the item was not found, otherwise + * the value pointer is returned. + */ +void* art_search(const art_tree *t, const unsigned char *key, int key_len) { + art_node **child; + art_node *n = t->root; + int prefix_len, depth = 0; + while (n) { + // Might be a leaf + if (IS_LEAF(n)) { + n = (art_node*)LEAF_RAW(n); + // Check if the expanded path matches + if (!leaf_matches((art_leaf*)n, key, key_len, depth)) { + return ((art_leaf*)n)->value; + } + return NULL; + } + + // Bail if the prefix does not match + if (n->partial_len) { + prefix_len = check_prefix(n, key, key_len, depth); + if (prefix_len != min(MAX_PREFIX_LEN, n->partial_len)) + return NULL; + depth = depth + n->partial_len; + } + + // Recursively search + child = find_child(n, key[depth]); + n = (child) ? *child : NULL; + depth++; + } + return NULL; +} + +// Find the minimum leaf under a node +static art_leaf* minimum(const art_node *n) { + // Handle base cases + if (!n) return NULL; + if (IS_LEAF(n)) return LEAF_RAW(n); + + int idx; + switch (n->type) { + case NODE4: + return minimum(((const art_node4*)n)->children[0]); + case NODE16: + return minimum(((const art_node16*)n)->children[0]); + case NODE48: + idx=0; + while (!((const art_node48*)n)->keys[idx]) idx++; + idx = ((const art_node48*)n)->keys[idx] - 1; + return minimum(((const art_node48*)n)->children[idx]); + case NODE256: + idx=0; + while (!((const art_node256*)n)->children[idx]) idx++; + return minimum(((const art_node256*)n)->children[idx]); + default: + abort(); + } +} + +// Find the maximum leaf under a node +static art_leaf* maximum(const art_node *n) { + // Handle base cases + if (!n) return NULL; + if (IS_LEAF(n)) return LEAF_RAW(n); + + int idx; + switch (n->type) { + case NODE4: + return maximum(((const art_node4*)n)->children[n->num_children-1]); + case NODE16: + return maximum(((const art_node16*)n)->children[n->num_children-1]); + case NODE48: + idx=255; + while (!((const art_node48*)n)->keys[idx]) idx--; + idx = ((const art_node48*)n)->keys[idx] - 1; + return maximum(((const art_node48*)n)->children[idx]); + case NODE256: + idx=255; + while (!((const art_node256*)n)->children[idx]) idx--; + return maximum(((const art_node256*)n)->children[idx]); + default: + abort(); + } +} + +/** + * Returns the minimum valued leaf + */ +art_leaf* art_minimum(art_tree *t) { + return minimum((art_node*)t->root); +} + +/** + * Returns the maximum valued leaf + */ +art_leaf* art_maximum(art_tree *t) { + return maximum((art_node*)t->root); +} + +static art_leaf* make_leaf(const unsigned char *key, int key_len, void *value) { + art_leaf *l = (art_leaf*)calloc(1, sizeof(art_leaf)+key_len); + l->value = value; + l->key_len = key_len; + memcpy(l->key, key, key_len); + return l; +} + +static int longest_common_prefix(art_leaf *l1, art_leaf *l2, int depth) { + int max_cmp = min(l1->key_len, l2->key_len) - depth; + int idx; + for (idx=0; idx < max_cmp; idx++) { + if (l1->key[depth+idx] != l2->key[depth+idx]) + return idx; + } + return idx; +} + +static void copy_header(art_node *dest, art_node *src) { + dest->num_children = src->num_children; + dest->partial_len = src->partial_len; + memcpy(dest->partial, src->partial, min(MAX_PREFIX_LEN, src->partial_len)); +} + +static void add_child256(art_node256 *n, art_node **ref, unsigned char c, void *child) { + (void)ref; + n->n.num_children++; + n->children[c] = (art_node*)child; +} + +static void add_child48(art_node48 *n, art_node **ref, unsigned char c, void *child) { + if (n->n.num_children < 48) { + int pos = 0; + while (n->children[pos]) pos++; + n->children[pos] = (art_node*)child; + n->keys[c] = pos + 1; + n->n.num_children++; + } else { + art_node256 *new_node = (art_node256*)alloc_node(NODE256); + for (int i=0;i<256;i++) { + if (n->keys[i]) { + new_node->children[i] = n->children[n->keys[i] - 1]; + } + } + copy_header((art_node*)new_node, (art_node*)n); + *ref = (art_node*)new_node; + free(n); + add_child256(new_node, ref, c, child); + } +} + +static void add_child16(art_node16 *n, art_node **ref, unsigned char c, void *child) { + if (n->n.num_children < 16) { + unsigned mask = (1 << n->n.num_children) - 1; + + // support non-x86 architectures + #ifdef __i386__ + __m128i cmp; + + // Compare the key to all 16 stored keys + cmp = _mm_cmplt_epi8(_mm_set1_epi8(c), + _mm_loadu_si128((__m128i*)n->keys)); + + // Use a mask to ignore children that don't exist + unsigned bitfield = _mm_movemask_epi8(cmp) & mask; + #else + #ifdef __amd64__ + __m128i cmp; + + // Compare the key to all 16 stored keys + cmp = _mm_cmplt_epi8(_mm_set1_epi8(c), + _mm_loadu_si128((__m128i*)n->keys)); + + // Use a mask to ignore children that don't exist + unsigned bitfield = _mm_movemask_epi8(cmp) & mask; + #else + // Compare the key to all 16 stored keys + unsigned bitfield = 0; + for (short i = 0; i < 16; ++i) { + if (c < n->keys[i]) + bitfield |= (1 << i); + } + + // Use a mask to ignore children that don't exist + bitfield &= mask; + #endif + #endif + + // Check if less than any + unsigned idx; + if (bitfield) { + idx = __builtin_ctz(bitfield); + memmove(n->keys+idx+1,n->keys+idx,n->n.num_children-idx); + memmove(n->children+idx+1,n->children+idx, + (n->n.num_children-idx)*sizeof(void*)); + } else + idx = n->n.num_children; + + // Set the child + n->keys[idx] = c; + n->children[idx] = (art_node*)child; + n->n.num_children++; + + } else { + art_node48 *new_node = (art_node48*)alloc_node(NODE48); + + // Copy the child pointers and populate the key map + memcpy(new_node->children, n->children, + sizeof(void*)*n->n.num_children); + for (int i=0;in.num_children;i++) { + new_node->keys[n->keys[i]] = i + 1; + } + copy_header((art_node*)new_node, (art_node*)n); + *ref = (art_node*)new_node; + free(n); + add_child48(new_node, ref, c, child); + } +} + +static void add_child4(art_node4 *n, art_node **ref, unsigned char c, void *child) { + if (n->n.num_children < 4) { + int idx; + for (idx=0; idx < n->n.num_children; idx++) { + if (c < n->keys[idx]) break; + } + + // Shift to make room + memmove(n->keys+idx+1, n->keys+idx, n->n.num_children - idx); + memmove(n->children+idx+1, n->children+idx, + (n->n.num_children - idx)*sizeof(void*)); + + // Insert element + n->keys[idx] = c; + n->children[idx] = (art_node*)child; + n->n.num_children++; + + } else { + art_node16 *new_node = (art_node16*)alloc_node(NODE16); + + // Copy the child pointers and the key map + memcpy(new_node->children, n->children, + sizeof(void*)*n->n.num_children); + memcpy(new_node->keys, n->keys, + sizeof(unsigned char)*n->n.num_children); + copy_header((art_node*)new_node, (art_node*)n); + *ref = (art_node*)new_node; + free(n); + add_child16(new_node, ref, c, child); + } +} + +static void add_child(art_node *n, art_node **ref, unsigned char c, void *child) { + switch (n->type) { + case NODE4: + return add_child4((art_node4*)n, ref, c, child); + case NODE16: + return add_child16((art_node16*)n, ref, c, child); + case NODE48: + return add_child48((art_node48*)n, ref, c, child); + case NODE256: + return add_child256((art_node256*)n, ref, c, child); + default: + abort(); + } +} + +/** + * Calculates the index at which the prefixes mismatch + */ +static int prefix_mismatch(const art_node *n, const unsigned char *key, int key_len, int depth) { + int max_cmp = min(min(MAX_PREFIX_LEN, n->partial_len), key_len - depth); + int idx; + for (idx=0; idx < max_cmp; idx++) { + if (n->partial[idx] != key[depth+idx]) + return idx; + } + + // If the prefix is short we can avoid finding a leaf + if (n->partial_len > MAX_PREFIX_LEN) { + // Prefix is longer than what we've checked, find a leaf + art_leaf *l = minimum(n); + max_cmp = min(l->key_len, key_len)- depth; + for (; idx < max_cmp; idx++) { + if (l->key[idx+depth] != key[depth+idx]) + return idx; + } + } + return idx; +} + +static void* recursive_insert(art_node *n, art_node **ref, const unsigned char *key, int key_len, void *value, int depth, int *old) { + // If we are at a NULL node, inject a leaf + if (!n) { + *ref = (art_node*)SET_LEAF(make_leaf(key, key_len, value)); + return NULL; + } + + // If we are at a leaf, we need to replace it with a node + if (IS_LEAF(n)) { + art_leaf *l = LEAF_RAW(n); + + // Check if we are updating an existing value + if (!leaf_matches(l, key, key_len, depth)) { + *old = 1; + void *old_val = l->value; + l->value = value; + return old_val; + } + + // New value, we must split the leaf into a node4 + art_node4 *new_node = (art_node4*)alloc_node(NODE4); + + // Create a new leaf + art_leaf *l2 = make_leaf(key, key_len, value); + + // Determine longest prefix + int longest_prefix = longest_common_prefix(l, l2, depth); + new_node->n.partial_len = longest_prefix; + memcpy(new_node->n.partial, key+depth, min(MAX_PREFIX_LEN, longest_prefix)); + // Add the leafs to the new node4 + *ref = (art_node*)new_node; + add_child4(new_node, ref, l->key[depth+longest_prefix], SET_LEAF(l)); + add_child4(new_node, ref, l2->key[depth+longest_prefix], SET_LEAF(l2)); + return NULL; + } + + // Check if given node has a prefix + if (n->partial_len) { + // Determine if the prefixes differ, since we need to split + int prefix_diff = prefix_mismatch(n, key, key_len, depth); + if ((uint32_t)prefix_diff >= n->partial_len) { + depth += n->partial_len; + goto RECURSE_SEARCH; + } + + // Create a new node + art_node4 *new_node = (art_node4*)alloc_node(NODE4); + *ref = (art_node*)new_node; + new_node->n.partial_len = prefix_diff; + memcpy(new_node->n.partial, n->partial, min(MAX_PREFIX_LEN, prefix_diff)); + + // Adjust the prefix of the old node + if (n->partial_len <= MAX_PREFIX_LEN) { + add_child4(new_node, ref, n->partial[prefix_diff], n); + n->partial_len -= (prefix_diff+1); + memmove(n->partial, n->partial+prefix_diff+1, + min(MAX_PREFIX_LEN, n->partial_len)); + } else { + n->partial_len -= (prefix_diff+1); + art_leaf *l = minimum(n); + add_child4(new_node, ref, l->key[depth+prefix_diff], n); + memcpy(n->partial, l->key+depth+prefix_diff+1, + min(MAX_PREFIX_LEN, n->partial_len)); + } + + // Insert the new leaf + art_leaf *l = make_leaf(key, key_len, value); + add_child4(new_node, ref, key[depth+prefix_diff], SET_LEAF(l)); + return NULL; + } + +RECURSE_SEARCH:; + + // Find a child to recurse to + art_node **child = find_child(n, key[depth]); + if (child) { + return recursive_insert(*child, child, key, key_len, value, depth+1, old); + } + + // No child, node goes within us + art_leaf *l = make_leaf(key, key_len, value); + add_child(n, ref, key[depth], SET_LEAF(l)); + return NULL; +} + +/** + * Inserts a new value into the ART tree + * @arg t The tree + * @arg key The key + * @arg key_len The length of the key + * @arg value Opaque value. + * @return NULL if the item was newly inserted, otherwise + * the old value pointer is returned. + */ +void* art_insert(art_tree *t, const unsigned char *key, int key_len, void *value) { + int old_val = 0; + void *old = recursive_insert(t->root, &t->root, key, key_len, value, 0, &old_val); + if (!old_val) t->size++; + return old; +} + +static void remove_child256(art_node256 *n, art_node **ref, unsigned char c) { + n->children[c] = NULL; + n->n.num_children--; + + // Resize to a node48 on underflow, not immediately to prevent + // trashing if we sit on the 48/49 boundary + if (n->n.num_children == 37) { + art_node48 *new_node = (art_node48*)alloc_node(NODE48); + *ref = (art_node*)new_node; + copy_header((art_node*)new_node, (art_node*)n); + + int pos = 0; + for (int i=0;i<256;i++) { + if (n->children[i]) { + new_node->children[pos] = n->children[i]; + new_node->keys[i] = pos + 1; + pos++; + } + } + free(n); + } +} + +static void remove_child48(art_node48 *n, art_node **ref, unsigned char c) { + int pos = n->keys[c]; + n->keys[c] = 0; + n->children[pos-1] = NULL; + n->n.num_children--; + + if (n->n.num_children == 12) { + art_node16 *new_node = (art_node16*)alloc_node(NODE16); + *ref = (art_node*)new_node; + copy_header((art_node*)new_node, (art_node*)n); + + int child = 0; + for (int i=0;i<256;i++) { + pos = n->keys[i]; + if (pos) { + new_node->keys[child] = i; + new_node->children[child] = n->children[pos - 1]; + child++; + } + } + free(n); + } +} + +static void remove_child16(art_node16 *n, art_node **ref, art_node **l) { + int pos = l - n->children; + memmove(n->keys+pos, n->keys+pos+1, n->n.num_children - 1 - pos); + memmove(n->children+pos, n->children+pos+1, (n->n.num_children - 1 - pos)*sizeof(void*)); + n->n.num_children--; + + if (n->n.num_children == 3) { + art_node4 *new_node = (art_node4*)alloc_node(NODE4); + *ref = (art_node*)new_node; + copy_header((art_node*)new_node, (art_node*)n); + memcpy(new_node->keys, n->keys, 4); + memcpy(new_node->children, n->children, 4*sizeof(void*)); + free(n); + } +} + +static void remove_child4(art_node4 *n, art_node **ref, art_node **l) { + int pos = l - n->children; + memmove(n->keys+pos, n->keys+pos+1, n->n.num_children - 1 - pos); + memmove(n->children+pos, n->children+pos+1, (n->n.num_children - 1 - pos)*sizeof(void*)); + n->n.num_children--; + + // Remove nodes with only a single child + if (n->n.num_children == 1) { + art_node *child = n->children[0]; + if (!IS_LEAF(child)) { + // Concatenate the prefixes + int prefix = n->n.partial_len; + if (prefix < MAX_PREFIX_LEN) { + n->n.partial[prefix] = n->keys[0]; + prefix++; + } + if (prefix < MAX_PREFIX_LEN) { + int sub_prefix = min(child->partial_len, MAX_PREFIX_LEN - prefix); + memcpy(n->n.partial+prefix, child->partial, sub_prefix); + prefix += sub_prefix; + } + + // Store the prefix in the child + memcpy(child->partial, n->n.partial, min(prefix, MAX_PREFIX_LEN)); + child->partial_len += n->n.partial_len + 1; + } + *ref = child; + free(n); + } +} + +static void remove_child(art_node *n, art_node **ref, unsigned char c, art_node **l) { + switch (n->type) { + case NODE4: + return remove_child4((art_node4*)n, ref, l); + case NODE16: + return remove_child16((art_node16*)n, ref, l); + case NODE48: + return remove_child48((art_node48*)n, ref, c); + case NODE256: + return remove_child256((art_node256*)n, ref, c); + default: + abort(); + } +} + +static art_leaf* recursive_delete(art_node *n, art_node **ref, const unsigned char *key, int key_len, int depth) { + // Search terminated + if (!n) return NULL; + + // Handle hitting a leaf node + if (IS_LEAF(n)) { + art_leaf *l = LEAF_RAW(n); + if (!leaf_matches(l, key, key_len, depth)) { + *ref = NULL; + return l; + } + return NULL; + } + + // Bail if the prefix does not match + if (n->partial_len) { + int prefix_len = check_prefix(n, key, key_len, depth); + if (prefix_len != min(MAX_PREFIX_LEN, n->partial_len)) { + return NULL; + } + depth = depth + n->partial_len; + } + + // Find child node + art_node **child = find_child(n, key[depth]); + if (!child) return NULL; + + // If the child is leaf, delete from this node + if (IS_LEAF(*child)) { + art_leaf *l = LEAF_RAW(*child); + if (!leaf_matches(l, key, key_len, depth)) { + remove_child(n, ref, key[depth], child); + return l; + } + return NULL; + + // Recurse + } else { + return recursive_delete(*child, child, key, key_len, depth+1); + } +} + +/** + * Deletes a value from the ART tree + * @arg t The tree + * @arg key The key + * @arg key_len The length of the key + * @return NULL if the item was not found, otherwise + * the value pointer is returned. + */ +void* art_delete(art_tree *t, const unsigned char *key, int key_len) { + art_leaf *l = recursive_delete(t->root, &t->root, key, key_len, 0); + if (l) { + t->size--; + void *old = l->value; + free(l); + return old; + } + return NULL; +} + +// Recursively iterates over the tree +static int recursive_iter(art_node *n, art_callback cb, void *data) { + // Handle base cases + if (!n) return 0; + if (IS_LEAF(n)) { + art_leaf *l = LEAF_RAW(n); + return cb(data, (const unsigned char*)l->key, l->key_len, l->value); + } + + int idx, res; + switch (n->type) { + case NODE4: + for (int i=0; i < n->num_children; i++) { + res = recursive_iter(((art_node4*)n)->children[i], cb, data); + if (res) return res; + } + break; + + case NODE16: + for (int i=0; i < n->num_children; i++) { + res = recursive_iter(((art_node16*)n)->children[i], cb, data); + if (res) return res; + } + break; + + case NODE48: + for (int i=0; i < 256; i++) { + idx = ((art_node48*)n)->keys[i]; + if (!idx) continue; + + res = recursive_iter(((art_node48*)n)->children[idx-1], cb, data); + if (res) return res; + } + break; + + case NODE256: + for (int i=0; i < 256; i++) { + if (!((art_node256*)n)->children[i]) continue; + res = recursive_iter(((art_node256*)n)->children[i], cb, data); + if (res) return res; + } + break; + + default: + abort(); + } + return 0; +} + +/** + * Iterates through the entries pairs in the map, + * invoking a callback for each. The call back gets a + * key, value for each and returns an integer stop value. + * If the callback returns non-zero, then the iteration stops. + * @arg t The tree to iterate over + * @arg cb The callback function to invoke + * @arg data Opaque handle passed to the callback + * @return 0 on success, or the return of the callback. + */ +int art_iter(art_tree *t, art_callback cb, void *data) { + return recursive_iter(t->root, cb, data); +} + +int art_iter_cmp(art_tree *t, const unsigned char *key, int key_len, art_callback cb, art_key_cmp_callback key_cmp, void *data) +{ + art_node **child; + art_node *n = t->root; + int prefix_len, depth = 0; + while (n) { + // Might be a leaf + if (IS_LEAF(n)) { + n = (art_node*)LEAF_RAW(n); + // Check if the expanded path matches + if (!key_cmp((art_leaf*)n, key, key_len)) { + art_leaf *l = (art_leaf*)n; + return cb(data, (const unsigned char*)l->key, l->key_len, l->value); + } + return 0; + } + + // If the depth matches the prefix, we need to handle this node + if (depth == key_len) { + art_leaf *l = minimum(n); + if (!key_cmp(l, key, key_len)) + return recursive_iter(n, cb, data); + return 0; + } + + // Bail if the prefix does not match + if (n->partial_len) { + prefix_len = prefix_mismatch(n, key, key_len, depth); + + // Guard if the mis-match is longer than the MAX_PREFIX_LEN + if ((uint32_t)prefix_len > n->partial_len) { + prefix_len = n->partial_len; + } + + // If there is no match, search is terminated + if (!prefix_len) { + return 0; + + // If we've matched the prefix, iterate on this node + } else if (depth + prefix_len == key_len) { + return recursive_iter(n, cb, data); + } + + // if there is a full match, go deeper + depth = depth + n->partial_len; + } + + // Recursively search + child = find_child(n, key[depth]); + n = (child) ? *child : NULL; + depth++; + } + return 0; +} + +/** + * Checks if a leaf prefix matches + * @return 0 on success. + */ +static int leaf_prefix_matches(const art_leaf *n, const unsigned char *prefix, int prefix_len) { + // Fail if the key length is too short + if (n->key_len < (uint32_t)prefix_len) return 1; + + // Compare the keys + return memcmp(n->key, prefix, prefix_len); +} + +/** + * Iterates through the entries pairs in the map, + * invoking a callback for each that matches a given prefix. + * The call back gets a key, value for each and returns an integer stop value. + * If the callback returns non-zero, then the iteration stops. + * @arg t The tree to iterate over + * @arg prefix The prefix of keys to read + * @arg prefix_len The length of the prefix + * @arg cb The callback function to invoke + * @arg data Opaque handle passed to the callback + * @return 0 on success, or the return of the callback. + */ +int art_iter_prefix(art_tree *t, const unsigned char *key, int key_len, art_callback cb, void *data) { + return art_iter_cmp(t, key, key_len, cb, leaf_prefix_matches, data); +} \ No newline at end of file diff --git a/src/smartdns.c b/src/smartdns.c index 8c23591..d9e7015 100644 --- a/src/smartdns.c +++ b/src/smartdns.c @@ -16,6 +16,7 @@ * along with this program. If not, see . */ +#include "art.h" #include "atomic.h" #include "conf.h" #include "dns_client.h" @@ -230,6 +231,7 @@ void smartdns_exit(void) dns_client_exit(); fast_ping_exit(); tlog_exit(); + load_exit(); } void sig_handle(int sig) @@ -246,57 +248,6 @@ void sig_handle(int sig) _exit(0); } -#if 0 -struct data_rbtree { - struct rb_node list; - int value; -}; - -int rbtree_test() -{ - struct rb_root root = RB_ROOT; - struct rb_node *n; - int i; - - for (i = 0; i < 10; i++) { - struct data_rbtree *r = malloc(sizeof(struct data_rbtree)); - struct rb_node **new = &root.rb_node, *parent = NULL; - r->value = i; - - while (*new) { - parent = *new; - if (i < rb_entry(parent, struct data_rbtree, list)->value) - new = &parent->rb_left; - else - new = &parent->rb_right; - } - - rb_link_node(&r->list, parent, new); - rb_insert_color(&r->list, &root); - } - - n = root.rb_node; - int num = 5; - while (n) { - struct data_rbtree *r = rb_entry(n, struct data_rbtree, list); - if (r->value > num) { - n = n->rb_left; - } else if (r->value < num) { - n = n->rb_right; - } else { - printf("n = %d\n", r->value); - break; - } - } - - struct rb_node *node; - for (node = rb_first(&root); node; node = rb_next(node)) - printf("V = %d\n", rb_entry(node, struct data_rbtree, list)->value); - - return 0; -} -#endif - int main(int argc, char *argv[]) { int ret; diff --git a/src/util.c b/src/util.c index 6702a8b..aa19869 100644 --- a/src/util.c +++ b/src/util.c @@ -77,7 +77,6 @@ errout: freeaddrinfo(result); } return -1; - } int parse_ip(const char *value, char *ip, int *port) @@ -144,7 +143,7 @@ int set_fd_nonblock(int fd, int nonblock) return -1; } - flags = (nonblock) ? (flags | O_NONBLOCK) : (flags & ~O_NONBLOCK); + flags = (nonblock) ? (flags | O_NONBLOCK) : (flags & ~O_NONBLOCK); ret = fcntl(fd, F_SETFL, flags); if (ret == -1) { return -1; @@ -152,3 +151,18 @@ int set_fd_nonblock(int fd, int nonblock) return 0; } + +char *reverse_string(char *output, char *input, int len) +{ + char *begin = output; + len--; + while (len >= 0) { + *output = *(input + len); + output++; + len--; + } + + *output = 0; + + return begin; +} diff --git a/src/util.h b/src/util.h index 9c99bbd..6cccd23 100644 --- a/src/util.h +++ b/src/util.h @@ -18,4 +18,6 @@ int parse_ip(const char *value, char *ip, int *port); int set_fd_nonblock(int fd, int nonblock); +char *reverse_string(char *output, char *input, int len); + #endif \ No newline at end of file diff --git a/systemd/smartdns.service b/systemd/smartdns.service index 6b2d686..ffa50c3 100644 --- a/systemd/smartdns.service +++ b/systemd/smartdns.service @@ -3,12 +3,15 @@ Description=smart dns server After=network.target [Service] +Type=forking PIDFile=/var/run/smartdns.pid EnvironmentFile=/etc/default/smartdns ExecStart=/usr/sbin/smartdns $SMART_DNS_OPTS KillMode=process Restart=always -RestartSec=0 +RestartSec=2 +StartLimitBurst=0 +StartLimitIntervalSec=60 [Install] WantedBy=multi-user.target