From fbfe86bd55f54e47260374b5cd1a6b5c69b795cb Mon Sep 17 00:00:00 2001 From: "mingang.he" Date: Mon, 30 Jul 2018 21:32:48 +0800 Subject: [PATCH] Support option for IP blacklist --- Makefile | 2 +- blacklist-ip-example | 2 + main.c | 44 ++++++++++++++++- string-set.c | 111 +++++++++++++++++++++++++++++++++++++++++++ string-set.h | 20 ++++++++ 5 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 blacklist-ip-example create mode 100644 string-set.c create mode 100644 string-set.h diff --git a/Makefile b/Makefile index bea8cf8..00a90d5 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ all: graftcp graftcp-local/graftcp-local -graftcp: main.o util.o +graftcp: main.o util.o string-set.o cc $^ -o $@ graftcp-local/graftcp-local: graftcp-local/*.go diff --git a/blacklist-ip-example b/blacklist-ip-example new file mode 100644 index 0000000..5395b75 --- /dev/null +++ b/blacklist-ip-example @@ -0,0 +1,2 @@ +111.111.123.222 +192.168.0.1 diff --git a/main.c b/main.c index 8c9bae0..88ea847 100644 --- a/main.c +++ b/main.c @@ -2,12 +2,40 @@ #include #include "graftcp.h" +#include "string-set.h" char *LOCAL_ADDR = "127.0.0.1"; uint16_t LOCAL_PORT = 2233; struct sockaddr_in PROXY_SA; char *LOCAL_PIPE_PAHT = "/tmp/graftcplocal.fifo"; int LOCAL_PIPE_FD; +struct str_set *BLACKLIST_IP = NULL; + +static void load_blackip_file(char *path) +{ + FILE *f; + char *line = NULL; + size_t len = 0; + ssize_t read; + + f = fopen(path, "r"); + if (f == NULL) { + perror("fopen"); + exit(1); + } + if (BLACKLIST_IP == NULL) { + BLACKLIST_IP = str_set_new(); + } + while ((read = getline(&line, &len, f)) != -1) { + /* 7 is the shortest ip: (x.x.x.x) */ + if (read < 7) + continue; + line[read - 1] = '\0'; + str_set_put(BLACKLIST_IP, line); + line = NULL; + } + fclose(f); +} void socket_pre_handle(struct proc_info *pinfp) { @@ -39,13 +67,19 @@ void connect_pre_handle(struct proc_info *pinfp) unsigned short dest_ip_port = SOCKPORT(dest_sa); struct in_addr dest_ip_addr; + char *dest_ip_addr_str; dest_ip_addr.s_addr = SOCKADDR(dest_sa); + dest_ip_addr_str = inet_ntoa(dest_ip_addr); + if (BLACKLIST_IP) { + if (is_str_set_member(BLACKLIST_IP, dest_ip_addr_str)) + return; + } putdata(pinfp->pid, addr, (char *)&PROXY_SA, sizeof(PROXY_SA)); char buf[1024] = {0}; - strcpy(buf, inet_ntoa(dest_ip_addr)); + strcpy(buf, dest_ip_addr_str); strcat(buf, ":"); sprintf(&buf[strlen(buf)], "%d:%d\n", ntohs(dest_ip_port), pinfp->pid); if (write(LOCAL_PIPE_FD, buf, strlen(buf)) <= 0) { @@ -253,6 +287,8 @@ static void usage(char **argv) " Which port is graftcp-local listening? Default: 2233\n" " -f --local-fifo=\n" " Path of fifo to communicate with graftcp-local. Default: /tmp/graftcplocal.fifo\n" + " -b --blackip-file=\n" + " The IP in black-ip-file will connect direct.\n" "\n", argv[0]); } @@ -264,10 +300,11 @@ int main(int argc, char **argv) {"local-addr", required_argument, 0, 'a'}, {"local-port", required_argument, 0, 'p'}, {"local-fifo", required_argument, 0, 'f'}, + {"blackip-file", required_argument, 0, 'b'}, {0, 0, 0, 0} }; - while ((opt = getopt_long(argc, argv, "+ha:p:f:", long_opts, &index)) != -1) { + while ((opt = getopt_long(argc, argv, "+ha:p:f:b:", long_opts, &index)) != -1) { switch (opt) { case 'a': LOCAL_ADDR = optarg; @@ -278,6 +315,9 @@ int main(int argc, char **argv) case 'f': LOCAL_PIPE_PAHT = optarg; break; + case 'b': + load_blackip_file(optarg); + break; case 0: case 'h': default: diff --git a/string-set.c b/string-set.c new file mode 100644 index 0000000..43382e2 --- /dev/null +++ b/string-set.c @@ -0,0 +1,111 @@ +#include +#include +#include + +#include "string-set.h" + +/* see http://www.cse.yorku.ca/~oz/hash.html */ +static unsigned long str_hash(const char *str) +{ + unsigned long hash = 5381; + int c; + + while (c = *str++) + hash = ((hash << 5) + hash) + c; /* hash * 33 + c */ + + return hash; +} + +struct str_set *str_set_new() +{ + struct str_set *set; + int i; + + set = calloc(1, sizeof(*set) + 509 * sizeof(set->buckets[0])); + set->size = 509; + set->buckets = (struct member **)(set + 1); + for (i = 0; i < set->size; i++) + set->buckets[i] = NULL; + set->length = 0; + return set; +} + +void str_set_put(struct str_set *set, const char *elem) +{ + int i; + struct member *p; + + assert(set); + assert(elem); + i = str_hash(elem) % set->size; + for (p = set->buckets[i]; p; p = p->link) { + if (strcmp(elem, p->element) == 0) + break; + } + if (p == NULL) { + p = calloc(1, sizeof(*p)); + p->element = elem; + p->link = set->buckets[i]; + set->buckets[i] = p; + set->length++; + } else { + p->element = elem; + } +} + +char *str_set_remove(struct str_set *set, const char *elem) +{ + int i; + struct member **pp; + + assert(set); + assert(elem); + i = str_hash(elem) % set->size; + for (pp = &set->buckets[i]; *pp; pp = &(*pp)->link) + if (strcmp(elem, (*pp)->element) == 0) { + struct member *p = *pp; + *pp = p->link; + elem = p->element; + free(p); + set->length--; + return (char *)elem; + } + return NULL; +} + + +int str_set_length(struct str_set *set) +{ + assert(set); + return set->length; +} + +int is_str_set_member(struct str_set *set, const void *elem) +{ + int i; + struct member *p; + + assert(set); + assert(elem); + i = str_hash(elem) % set->size; + for (p = set->buckets[i]; p; p = p->link) { + if (strcmp(elem, p->element) == 0) + break; + } + return p != NULL; +} + +void str_set_free(struct str_set **set) +{ + assert(set && *set); + if ((*set)->length > 0) { + int i; + struct member *p, *q; + for (i = 0; i < (*set)->size; i++) + for (p = (*set)->buckets[i]; p; p = q) { + q = p->link; + free(p); + } + } + free(*set); +} diff --git a/string-set.h b/string-set.h new file mode 100644 index 0000000..f568e98 --- /dev/null +++ b/string-set.h @@ -0,0 +1,20 @@ +#ifndef _STRING_SET_H +#define _STRING_SET_H + +struct str_set { + int length; + int size; + struct member { + struct member *link; + const char *element; + } **buckets; +}; + +struct str_set *str_set_new(); +void str_set_put(struct str_set *set, const char *elem); +char *str_set_remove(struct str_set *set, const char *elem); +int str_set_length(struct str_set *set); +int is_str_set_member(struct str_set *set, const void *elem); +void str_set_free(struct str_set **set); + +#endif