diff --git a/src/main.c b/src/main.c index bca310a..9e71338 100644 --- a/src/main.c +++ b/src/main.c @@ -1,10 +1,8 @@ #include #include -#include #include #include #include -#include #include #include #include @@ -15,6 +13,9 @@ #include #include +#include +#include + #define MAX_MESSAGE_SIZE 0x200 #define DNS_FLAG_TC 0x200 @@ -105,6 +106,8 @@ typedef struct client_ctx { } client_ctx_t; static struct nft_ctx *nft_ctx; +static const char *nft_nat_table = "dotp"; +static char *nft_fake_set; char *malloc_sprintf(const char *fmt, ...) { char *buffer = NULL; @@ -170,16 +173,28 @@ static void nat_expire(EV_P_ ev_timer *w, int revents) { ev_timer_stop(EV_A, w); - char *cmd = malloc_sprintf("delete rule ip nat postrouting handle %d", - nat->src_handle); + char *cmd = malloc_sprintf("delete rule ip %s postrouting handle %d", + nft_nat_table, nat->src_handle); nft_run_cmd_from_buffer(nft_ctx, cmd); free(cmd); - cmd = malloc_sprintf("delete rule ip nat prerouting handle %d", + cmd = malloc_sprintf("delete rule ip %s prerouting handle %d", nft_nat_table, nat->dst_handle); nft_run_cmd_from_buffer(nft_ctx, cmd); free(cmd); + if (nft_fake_set) { + char real_ip[16]; + + uint32_t af_real = htonl(nat->real); + inet_ntop(AF_INET, &af_real, real_ip, 16); + + char *cmd = + malloc_sprintf("delete element %s { %s }", nft_fake_set, real_ip); + nft_run_cmd_from_buffer(nft_ctx, cmd); + free(cmd); + } + uint32_t mask = pool->size - 1; ip_nat_t **p_other = &pool->real[nat->real & mask]; @@ -265,9 +280,17 @@ static ip_nat_t *find_nat(EV_P_ ip_pool_t *pool, uint32_t real_addr) { inet_ntop(AF_INET, &af_fake, fake_ip, 16); inet_ntop(AF_INET, &af_real, real_ip, 16); + if (nft_fake_set) { + char *cmd = + malloc_sprintf("add element %s { %s }", nft_fake_set, real_ip); + nft_run_cmd_from_buffer(nft_ctx, cmd); + free(cmd); + } + nft_ctx_buffer_output(nft_ctx); - char *cmd = malloc_sprintf( - "add rule ip nat prerouting ip daddr %s dnat to %s", fake_ip, real_ip); + char *cmd = + malloc_sprintf("add rule ip %s prerouting ip daddr %s dnat to %s", + nft_nat_table, fake_ip, real_ip); nft_run_cmd_from_buffer(nft_ctx, cmd); char *echo_fmt = malloc_sprintf("%s # handle %%d", cmd); sscanf(nft_ctx_get_output_buffer(nft_ctx), echo_fmt, &nat->dst_handle); @@ -276,8 +299,8 @@ static ip_nat_t *find_nat(EV_P_ ip_pool_t *pool, uint32_t real_addr) { nft_ctx_unbuffer_output(nft_ctx); nft_ctx_buffer_output(nft_ctx); - cmd = malloc_sprintf("add rule ip nat postrouting ip saddr %s snat to %s", - real_ip, fake_ip); + cmd = malloc_sprintf("add rule ip %s postrouting ip saddr %s snat to %s", + nft_nat_table, real_ip, fake_ip); nft_run_cmd_from_buffer(nft_ctx, cmd); echo_fmt = malloc_sprintf("%s # handle %%d", cmd); sscanf(nft_ctx_get_output_buffer(nft_ctx), echo_fmt, &nat->src_handle); @@ -579,8 +602,27 @@ static void terminate(EV_P_ ev_signal *w, int revents) { ev_break(EV_A, EVBREAK_ALL); } +static const char *parse_nft_string(const char *s) { + const char *p = s; + if (isalpha(*p)) { + p++; + while (isalnum(*p) || *p == '/' || *p == '-' || *p == '_' || *p == '.') { + p++; + } + } else if (*p == '\"') { + p++; + while (*p != '\"') { + p++; + } + p++; + } + return p; +} + #define OPT_UPSTREAM_HOST 0x101 #define OPT_UPSTREAM_PORT 0x102 +#define OPT_FAKE_SET 0x103 +#define OPT_DAEMONIZE 0x104 int main(int argc, char *const *argv) { domain_set_t *domain_set = domain_set_new("", 0); @@ -603,6 +645,8 @@ int main(int argc, char *const *argv) { {"prefix", required_argument, NULL, 'x'}, {"upstream-host", required_argument, NULL, OPT_UPSTREAM_HOST}, {"upstream-port", required_argument, NULL, OPT_UPSTREAM_PORT}, + {"fake-set", required_argument, NULL, OPT_FAKE_SET}, + {"daemonize", no_argument, NULL, OPT_DAEMONIZE}, {NULL, 0, NULL, 0}}; int option_index; @@ -666,6 +710,39 @@ int main(int argc, char *const *argv) { goto fail; } break; + case OPT_FAKE_SET: { + const char *s = optarg; + size_t t[3]; + t[0] = parse_nft_string(s) - s; + if (s[t[0]] != '#' || t[0] == 0) { + goto fail; + } + + t[1] = parse_nft_string(s + t[0] + 1) - s; + if (s[t[1]] != '#' || t[1] == t[0] + 1) { + goto fail; + } + + t[2] = parse_nft_string(s + t[1] + 1) - s; + if (s[t[2]] != '\0' || t[2] == t[1] + 1) { + goto fail; + } + + char *fake_set = malloc(t[2] + 1); + + memcpy(fake_set, s, t[2] + 1); + fake_set[t[0]] = ' '; + fake_set[t[1]] = ' '; + nft_fake_set = fake_set; + + break; + } + case OPT_DAEMONIZE: + if (daemon(0, 0) == -1) { + domain_set_fini(domain_set); + return 1; + } + break; fail: opterr = 1; fprintf(stderr, "Failed to parse option: %s", optarg); @@ -676,14 +753,16 @@ int main(int argc, char *const *argv) { fprintf(stderr, "Usage: %s -h LISTEN_HOST -p LISTEN_PORT\n" " -x FAKE_IP_PREFIX [-d DOMAIN]\n" - " --upstream-host UPSTREAM_HOST --upstream-port " - "UPSTREAM_PORT\n", + " --upstream-host UPSTREAM_HOST\n" + " --upstream-port UPSTREAM_PORT\n" + " [ --fake-set FAKE_SET ]\n" + " [ --daemonize ]\n", argv[0]); break; } } - if (!(listen_host_set && prefix_set && upstream_host_set)) { + if (!opterr && !(listen_host_set && prefix_set && upstream_host_set)) { fprintf(stderr, "LISTEN_ADDR, FAKE_IP_PREFIX, UPSTREAM_HOST must be set.\n"); opterr = 1; @@ -704,17 +783,19 @@ int main(int argc, char *const *argv) { nft_ctx_output_set_flags(nft_ctx, NFT_CTX_OUTPUT_HANDLE | NFT_CTX_OUTPUT_ECHO); - char *cmd = malloc_sprintf("add table ip nat"); + char *cmd = malloc_sprintf("add table ip %s", nft_nat_table); nft_run_cmd_from_buffer(nft_ctx, cmd); free(cmd); - cmd = malloc_sprintf("add chain ip nat prerouting" - "{ type nat hook prerouting priority dstnat; }"); + cmd = malloc_sprintf("add chain ip %s prerouting" + "{ type nat hook prerouting priority dstnat; }", + nft_nat_table); nft_run_cmd_from_buffer(nft_ctx, cmd); free(cmd); - cmd = malloc_sprintf("add chain ip nat postrouting" - "{ type nat hook postrouting priority srcnat; }"); + cmd = malloc_sprintf("add chain ip %s postrouting" + "{ type nat hook postrouting priority srcnat; }", + nft_nat_table); nft_run_cmd_from_buffer(nft_ctx, cmd); free(cmd); @@ -752,10 +833,12 @@ int main(int argc, char *const *argv) { ip_pool_fini(&ip_pool); - cmd = malloc_sprintf("delete table ip nat"); + cmd = malloc_sprintf("delete table ip %s", nft_nat_table); nft_run_cmd_from_buffer(nft_ctx, cmd); free(cmd); + free(nft_fake_set); + nft_ctx_free(nft_ctx); domain_set_fini(domain_set);