From 0b14bc6d8b07468946a75b9512f3e28724edcd98 Mon Sep 17 00:00:00 2001 From: ruti <> Date: Fri, 20 Dec 2024 04:15:13 +0300 Subject: [PATCH] Optimize host search, fix ipset --- extend.c | 23 +++++++------- main.c | 35 +++++++++++++-------- mpool.c | 92 ++++++++++++++++++++++++++++++++------------------------ mpool.h | 7 +++-- 4 files changed, 90 insertions(+), 67 deletions(-) diff --git a/extend.c b/extend.c index ae764ae..e8de005 100644 --- a/extend.c +++ b/extend.c @@ -105,10 +105,17 @@ 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; } diff --git a/main.c b/main.c index 0d91c46..ebdef7d 100644 --- a/main.c +++ b/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); - return 0; + 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(); diff --git a/mpool.c b/mpool.c index ef560db..0ea3450 100644 --- a/mpool.c +++ b/mpool.c @@ -3,14 +3,12 @@ #include #include #include - +#include 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); - - return byte_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->data = str; - } + 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; + if (e != v) { + if (!hdr->static_data) + free(e->data); + free(e); + } + 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--; } diff --git a/mpool.h b/mpool.h index e89a69d..440219e 100644 --- a/mpool.h +++ b/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);