Optimize host search, fix ipset

extend.c

@@ -106,9 +106,16 @@ static int cache_add(const struct sockaddr_ina *dst, int m)
     LOG(LOG_S, "save ip: %s, m=%d\n", ADDR_STR, m);
     time_t t = time(0);
     
-    struct elem_i *val = (struct elem_i *)mem_add(params.mempool, (char *)key, len, sizeof(struct elem_i));
+    char *key_d = malloc(len);
+    if (!key_d) {
+        return -1;
+    }
+    memcpy(key_d, key, len);
+    
+    struct elem_i *val = (struct elem_i *)mem_add(params.mempool, key_d, len, sizeof(struct elem_i));
     if (!val) {
         uniperror("mem_add");
+        free(key_d);
         return -1;
     }
     val->m = m;
@@ -177,16 +184,8 @@ static bool check_host(
     if (len <= 0) {
         return 0;
     }
-    char *e = host + len;
-    for (; host < e; host++) {
-        if (mem_get(hosts, host, e - host)) {
-            return 1;
-        }
-        if (!(host = memchr(host, '.', e - host))) {
-            return 0;
-        }
-    }
-    return 0;
+    struct elem *v = mem_get(hosts, host, len);
+    return v && v->len <= len;
 }
 
 

main.c

@@ -289,18 +289,19 @@ static inline int lower_char(char *cl)
 
 struct mphdr *parse_hosts(char *buffer, size_t size)
 {
-    struct mphdr *hdr = mem_pool(1);
+    struct mphdr *hdr = mem_pool(1, CMP_HOST);
     if (!hdr) {
         return 0;
     }
     size_t num = 0;
+    bool drop = 0;
     char *end = buffer + size;
     char *e = buffer, *s = buffer;
     
     for (; e <= end; e++) {
         if (e != end && *e != ' ' && *e != '\n' && *e != '\r') {
             if (lower_char(e)) {
-                LOG(LOG_E, "invalid host: num: %zd (%.*s)\n", num + 1, (int )(e - s + 1), s);
+                drop = 1;
             }
             continue;
         }
@@ -308,13 +309,20 @@ struct mphdr *parse_hosts(char *buffer, size_t size)
             s++;
             continue;
         }
-        if (mem_add(hdr, s, e - s, sizeof(struct elem)) == 0) {
-            free(hdr);
+        if (!drop) {
+            if (!mem_add(hdr, s, e - s, sizeof(struct elem))) {
+                mem_destroy(hdr);
                 return 0;
             }
+        } 
+        else {
+            LOG(LOG_E, "invalid host: num: %zd \"%.*s\"\n", num + 1, (int )(e - s), s);
+            drop = 0;
+        }
         num++;
         s = e + 1;
     }
+    LOG(LOG_S, "hosts count: %zd\n", hdr->count);
     return hdr;
 }
 
@@ -345,7 +353,7 @@ static int parse_ip(char *out, char *str, size_t size)
 
 struct mphdr *parse_ipset(char *buffer, size_t size)
 {
-    struct mphdr *hdr = mem_pool(0);
+    struct mphdr *hdr = mem_pool(0, CMP_BITS);
     if (!hdr) {
         return 0;
     }
@@ -368,20 +376,21 @@ struct mphdr *parse_ipset(char *buffer, size_t size)
         num++;
         s = e + 1;
         
-        char ip_raw[16];
+        char *ip_raw = malloc(16);
         int bits = parse_ip(ip_raw, ip, sizeof(ip));
         if (bits <= 0) {
-            LOG(LOG_E, "invalid ip: num: %zd\n", num + 1);
+            LOG(LOG_E, "invalid ip: num: %zd\n", num);
+            free(ip_raw);
             continue;
         }
-        struct elem *elem = mem_add(hdr, ip_raw, bits / 8 + (bits % 8 ? 1 : 0), sizeof(struct elem));
+        struct elem *elem = mem_add(hdr, ip_raw, bits, sizeof(struct elem));
         if (!elem) {
-            free(hdr);
+            free(ip_raw);
+            mem_destroy(hdr);
             return 0;
         }
-        elem->len = bits;
-        elem->cmp_type = CMP_BITS;
     }
+    LOG(LOG_S, "ip count: %zd\n", hdr->count);
     return hdr;
 }
 
@@ -1038,7 +1047,7 @@ int main(int argc, char **argv)
             return -1;
         }
     }
-    params.mempool = mem_pool(0);
+    params.mempool = mem_pool(0, CMP_BYTES);
     if (!params.mempool) {
         uniperror("mem_pool");
         clear_params();

mpool.c

@@ -3,14 +3,12 @@
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
-
+#include <assert.h>
 
 static int bit_cmp(const struct elem *p, const struct elem *q)
 {
-    if (p->len < q->len) {
-        return -1;
-    }
-    int df = q->len % 8, bytes = q->len / 8;
+    int len = q->len < p->len ? q->len : p->len;
+    int df = len % 8, bytes = len / 8;
     int cmp = memcmp(p->data, q->data, bytes);
     
     if (cmp || !df) {
@@ -33,23 +31,44 @@ static int byte_cmp(const struct elem *p, const struct elem *q)
     return memcmp(p->data, q->data, p->len);
 }
 
+static int host_cmp(const struct elem *p, const struct elem *q)
+{
+    int len = q->len < p->len ? q->len : p->len;
+    char *pd = p->data + p->len, *qd = q->data + q->len;
+    
+    while (len-- > 0) {
+        if (*--pd != *--qd) {
+            return *pd < *qd ? -1 : 1;
+        }
+    }
+    if (p->len == q->len 
+            || (p->len > q->len ? pd[-1] : qd[-1]) == '.')
+        return 0;
+    
+    return 1;
+}
 
 static int scmp(const struct elem *p, const struct elem *q)
 {
-    if (q->cmp_type == CMP_BITS)
+    switch (p->cmp_type) {
+    case CMP_BITS:
         return bit_cmp(p, q);
-        
+    case CMP_HOST:
+        return host_cmp(p, q);
+    default:
         return byte_cmp(p, q);
     }
+}
 
 KAVL_INIT(my, struct elem, head, scmp)
 
 
-struct mphdr *mem_pool(bool is_static)
+struct mphdr *mem_pool(bool is_static, unsigned char cmp_type)
 {
     struct mphdr *hdr = calloc(sizeof(struct mphdr), 1);
     if (hdr) {
         hdr->static_data = is_static;
+        hdr->cmp_type = cmp_type;
     }
     return hdr;
 }
@@ -57,12 +76,11 @@ struct mphdr *mem_pool(bool is_static)
 
 struct elem *mem_get(const struct mphdr *hdr, const char *str, int len)
 {
-    struct {
-        int len;
-        const char *data;
-    } temp = { .len = len, .data = str };
-    
-    return kavl_find(my, hdr->root, (struct elem *)&temp, 0);
+    struct elem temp = { 
+        .cmp_type = hdr->cmp_type,
+        .len = len, .data = (char *)str 
+    };
+    return kavl_find(my, hdr->root, &temp, 0);
 }
 
 
@@ -73,38 +91,31 @@ struct elem *mem_add(struct mphdr *hdr, char *str, int len, size_t struct_size)
         return 0;
     }
     e->len = len;
-    while (1) {
-        if (!hdr->static_data) {
-            e->data = malloc(len);
-            if (!e->data) {
-                break;
-            }
-            memcpy(e->data, str, len);
-        }
-        else {
+    e->cmp_type = hdr->cmp_type;
     e->data = str;
-        }
+    
     v = kavl_insert(my, &hdr->root, e, 0);
+    while (e != v && e->len < v->len) {
+        mem_delete(hdr, v->data, v->len);
+        v = kavl_insert(my, &hdr->root, e, 0);
+    }
     if (e != v) {
         if (!hdr->static_data)
             free(e->data);
-            break;
-        }
-        return v;
-    }
         free(e);
-    return 0;
+    }
+    else hdr->count++;
+    return v;
 }
 
 
 void mem_delete(struct mphdr *hdr, const char *str, int len)
 {
-    struct {
-        int len;
-        const char *data;
-    } temp = { .len = len, .data = str };
-    
-    struct elem *e = kavl_erase(my, &hdr->root, (struct elem *)&temp, 0);
+    struct elem temp = { 
+        .cmp_type = hdr->cmp_type,
+        .len = len, .data = (char *)str 
+    };
+    struct elem *e = kavl_erase(my, &hdr->root, &temp, 0);
     if (!e) {
         return;
     }
@@ -113,6 +124,7 @@ void mem_delete(struct mphdr *hdr, const char *str, int len)
         e->data = 0;
     }
     free(e);
+    hdr->count--;
 }
 
 

mpool.h

@@ -7,11 +7,12 @@
 
 #define CMP_BYTES 0
 #define CMP_BITS 1
+#define CMP_HOST 2
 
 struct elem {
     int len;
     char *data;
-    char cmp_type;
+    unsigned char cmp_type;
     KAVL_HEAD(struct elem) head;
 };
 
@@ -23,10 +24,12 @@ struct elem_i {
 
 struct mphdr {
     bool static_data;
+    unsigned char cmp_type;
+    size_t count;
     struct elem *root;
 };
 
-struct mphdr *mem_pool(bool is_static);
+struct mphdr *mem_pool(bool is_static, unsigned char cmp_type);
 
 struct elem *mem_get(const struct mphdr *hdr, const char *str, int len);