--ipset

extend.c

@@ -116,6 +116,7 @@ static int cache_add(const struct sockaddr_ina *dst, int m)
     return 0;
 }
 
+static bool check_l34(struct desync_params *dp, int st, const struct sockaddr_in6 *dst);
 
 int connect_hook(struct poolhd *pool, struct eval *val, 
         const struct sockaddr_ina *dst, int next)
@@ -188,7 +189,26 @@ static bool check_host(
     return 0;
 }
 
+
+static bool check_ip(
+        struct mphdr *ipset, const struct sockaddr_in6 *addr)
+{
+    const struct sockaddr_ina *dst = (const struct sockaddr_ina *)addr;
     
+    int len = sizeof(dst->in.sin_addr);
+    char *data = (char *)&dst->in.sin_addr;
+    
+    if (dst->sa.sa_family == AF_INET6) {
+        len = sizeof(dst->in6.sin6_addr);
+        data = (char *)&dst->in6.sin6_addr;
+    }
+    if (mem_get(ipset, data, len * 8)) {
+        return 1;
+    }
+    return 0;
+}
+
+
 static bool check_proto_tcp(int proto, const char *buffer, ssize_t n)
 {
     if (!(proto & ~IS_IPV4)) {
@@ -206,12 +226,12 @@ static bool check_proto_tcp(int proto, const char *buffer, ssize_t n)
 }
 
 
-static bool check_l34(int proto, const uint16_t *pf, int st, const struct sockaddr_in6 *dst)
+static bool check_l34(struct desync_params *dp, int st, const struct sockaddr_in6 *dst)
 {
-    if ((proto & IS_UDP) && (st != SOCK_DGRAM)) {
+    if ((dp->proto & IS_UDP) && (st != SOCK_DGRAM)) {
         return 0;
     }
-    if (proto & IS_IPV4) {
+    if (dp->proto & IS_IPV4) {
         static const char *pat = "\0\0\0\0\0\0\0\0\0\0\xff\xff";
         
         if (dst->sin6_family != AF_INET 
@@ -219,8 +239,11 @@ static bool check_l34(int proto, const uint16_t *pf, int st, const struct sockad
             return 0;
         }
     }
-    if (pf[0] && 
+    if (dp->pf[0] && 
-            (dst->sin6_port < pf[0] || dst->sin6_port > pf[1])) {
+            (dst->sin6_port < dp->pf[0] || dst->sin6_port > dp->pf[1])) {
+        return 0;
+    }
+    if (dp->ipset && !check_ip(dp->ipset, dst)) {
         return 0;
     }
     return 1;
@@ -340,8 +363,8 @@ static int setup_conn(struct eval *client, const char *buffer, ssize_t n)
     if (!m) for (; m < params.dp_count; m++) {
         struct desync_params *dp = &params.dp[m];
         if (!dp->detect 
-                && (check_l34(dp->proto, dp->pf, SOCK_STREAM, &client->pair->in6)
+                && check_l34(dp, SOCK_STREAM, &client->pair->in6)
-                    && check_proto_tcp(dp->proto, buffer, n)) 
+                && check_proto_tcp(dp->proto, buffer, n) 
                 && (!dp->hosts || check_host(dp->hosts, buffer, n))) {
             break;
         }
@@ -544,7 +567,7 @@ ssize_t udp_hook(struct eval *val,
         for (; m < params.dp_count; m++) {
             struct desync_params *dp = &params.dp[m];
             if (!dp->detect 
-                    && check_l34(dp->proto, dp->pf, SOCK_DGRAM, &dst->in6)) {
+                    && check_l34(dp, SOCK_DGRAM, &dst->in6)) {
                 break;
             }
         }

main.c

@@ -84,6 +84,7 @@ const static char help_text[] = {
     #endif
     "    -K, --proto <t,h,u,i>     Protocol whitelist: tls,http,udp,ipv4\n"
     "    -H, --hosts <file|:str>   Hosts whitelist, filename or :string\n"
+    "    -j, --ipset <file|:str>   IP whitelist\n"
     "    -V, --pf <port[-portr]>   Ports range whitelist\n"
     "    -R, --round <num[-numr]>  Number of request to which desync will be applied\n"
     "    -s, --split <pos_t>       Position format: offset[:repeats:skip][+flag1[flag2]]\n"
@@ -167,6 +168,7 @@ const struct option options[] = {
     {"drop-sack",     0, 0, 'Y'},
     {"protect-path",  1, 0, 'P'}, //
     #endif
+    {"ipset",         1, 0, 'j'},
     {0}
 };
     
@@ -317,6 +319,73 @@ struct mphdr *parse_hosts(char *buffer, size_t size)
 }
 
 
+static int parse_ip(char *out, char *str, size_t size)
+{
+    long bits = 0;
+    char *sep = memchr(str, '/', size);
+    if (sep) {
+        bits = strtol(sep + 1, 0, 10);
+        if (bits <= 0) {
+            return 0;
+        }
+        *sep = 0;
+    }
+    int len = 4;
+    
+    if (inet_pton(AF_INET, str, out) <= 0) {
+        if (inet_pton(AF_INET6, str, out) <= 0) {
+            return 0;
+        }
+        else len = 16;
+    }
+    if (!bits || bits > len * 8) bits = len * 8;
+    return (int )bits;
+}
+
+
+struct mphdr *parse_ipset(char *buffer, size_t size)
+{
+    struct mphdr *hdr = mem_pool(0);
+    if (!hdr) {
+        return 0;
+    }
+    size_t num = 0;
+    char *end = buffer + size;
+    char *e = buffer, *s = buffer;
+    
+    for (; e <= end; e++) {
+        if (e != end && *e != ' ' && *e != '\n' && *e != '\r') {
+            continue;
+        }
+        if (s == e) {
+            s++;
+            continue;
+        }
+        char ip[e - s + 1];
+        ip[e - s] = 0;
+        memcpy(ip, s, e - s);
+        
+        num++;
+        s = e + 1;
+        
+        char ip_raw[16];
+        int bits = parse_ip(ip_raw, ip, sizeof(ip));
+        if (bits <= 0) {
+            LOG(LOG_E, "invalid ip: num: %zd\n", num + 1);
+            continue;
+        }
+        struct elem *elem = mem_add(hdr, ip_raw, bits / 8 + (bits % 8 ? 1 : 0), sizeof(struct elem));
+        if (!elem) {
+            free(hdr);
+            return 0;
+        }
+        elem->len = bits;
+        elem->cmp_type = CMP_BITS;
+    }
+    return hdr;
+}
+
+
 int get_addr(const char *str, struct sockaddr_ina *addr)
 {
     struct addrinfo hints = {0}, *res = 0;
@@ -475,6 +544,10 @@ void clear_params(void)
                 mem_destroy(s.hosts);
                 s.hosts = 0;
             }
+            if (s.ipset != 0) {
+                mem_destroy(s.ipset);
+                s.hosts = 0;
+            }
         }
         free(params.dp);
         params.dp = 0;
@@ -615,7 +688,7 @@ int main(int argc, char **argv)
             break;
             
         case 'A':
-            if (!(dp->hosts || dp->proto || dp->pf[0] || dp->detect)) {
+            if (!(dp->hosts || dp->proto || dp->pf[0] || dp->detect || dp->ipset)) {
                 all_limited = 0;
             }
             dp = add((void *)&params.dp, &params.dp_count,
@@ -715,6 +788,25 @@ int main(int argc, char **argv)
             }
             break;
             
+        case 'j':;
+            if (dp->ipset) {
+                continue;
+            }
+            ssize_t size;
+            char *data = ftob(optarg, &size);
+            if (!data) {
+                uniperror("read/parse");
+                invalid = 1;
+                continue;
+            }
+            dp->ipset = parse_ipset(data, size);
+            if (!dp->ipset) {
+                perror("parse_ipset");
+                invalid = 1;
+            }
+            free(data);
+            break;
+            
         case 's':
         case 'd':
         case 'o':

mpool.c

@@ -1,17 +1,47 @@
 #include "mpool.h"
 
+#include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
 
 
-static inline int scmp(const struct elem *p, const struct elem *q)
+static int bit_cmp(const struct elem *p, const struct elem *q)
 {
-    if (p->len != q ->len) {
+    if (p->len < q->len) {
+        return -1;
+    }
+    int df = q->len % 8, bytes = q->len / 8;
+    int cmp = memcmp(p->data, q->data, bytes);
+    
+    if (cmp || !df) {
+        return cmp;
+    }
+    uint8_t c1 = p->data[bytes] >> (8 - df);
+    uint8_t c2 = q->data[bytes] >> (8 - df);
+    if (c1 != c2) {
+        if (c1 < c2) return -1;
+        else return 1;
+    }
+    return 0;
+}
+
+static int byte_cmp(const struct elem *p, const struct elem *q)
+{
+    if (p->len != q->len) {
         return p->len < q->len ? -1 : 1;
     }
     return memcmp(p->data, q->data, p->len);
 }
-    
+
+
+static int scmp(const struct elem *p, const struct elem *q)
+{
+    if (q->cmp_type == CMP_BITS)
+        return bit_cmp(p, q);
+        
+    return byte_cmp(p, q);
+}
+
 KAVL_INIT(my, struct elem, head, scmp)
 
 

mpool.h

@@ -5,9 +5,13 @@
 #include <time.h>
 #include "kavl.h"
 
+#define CMP_BYTES 0
+#define CMP_BITS 1
+
 struct elem {
     int len;
     char *data;
+    char cmp_type;
     KAVL_HEAD(struct elem) head;
 };
 

params.h

@@ -88,6 +88,7 @@ struct desync_params {
     int proto;
     int detect;
     struct mphdr *hosts;
+    struct mphdr *ipset;
     uint16_t pf[2];
     int rounds[2];