[PATCH gwproxy v13 6/8] gwproxy: Introduce --dns-server and --raw-dns

[Alviro Iskandar Setiawan <alviro.iskandar@gnuweeb.org>]

This is the initial integration of the new DNS resolver feature that
does not rely on getaddrinfo(). Introduce two new options in order to
utilize it.

1) --dns-server to be filled with DNS servers, for now it only supports
   a single host. In the futures, it may accept multiple hosts
   separated by commas.

2) --raw-dns to be filled with either 0 or 1. 0 to disable, 1 to
   enable.

While in there, plug the new DNS resolver data structure into struct
gwp_wrk and initialize it.

Note that this feature is guarded with CONFIG_NEW_DNS_RESOLVER, which
is disabled by default as it's still experimental. A future patch will
add a configure option to enable it.

Signed-off-by: Alviro Iskandar Setiawan <alviro.iskandar@gnuweeb.org>
---
 src/gwproxy/gwproxy.c | 205 +++++++++++++++++++++++++++++++++++++++---
 src/gwproxy/gwproxy.h |  25 +++++-
 2 files changed, 219 insertions(+), 11 deletions(-)

diff --git a/src/gwproxy/gwproxy.c b/src/gwproxy/gwproxy.c
index 3d277f7..3998c6a 100644
--- a/src/gwproxy/gwproxy.c
+++ b/src/gwproxy/gwproxy.c
@@ -67,6 +67,10 @@ static const struct option long_opts[] = {
 	{ "log-level",		required_argument,	NULL,	'm' },
 	{ "log-file",		required_argument,	NULL,	'f' },
 	{ "pid-file",		required_argument,	NULL,	'p' },
+#ifdef CONFIG_NEW_DNS_RESOLVER
+	{ "dns-server",		required_argument,	NULL,	'j' },
+	{ "raw-dns",		required_argument,	NULL,	'r' },
+#endif
 	{ NULL,			0,			NULL,	0 }
 };
 
@@ -77,6 +81,7 @@ static const struct gwp_cfg default_opts = {
 	.as_socks5		= false,
 	.as_http		= false,
 	.socks5_prefer_ipv6	= false,
+	.use_raw_dns		= false,
 	.protocol_timeout	= 10,
 	.socks5_auth_file	= NULL,
 	.socks5_dns_cache_secs	= 0,
@@ -94,6 +99,7 @@ static const struct gwp_cfg default_opts = {
 	.log_level		= 3,
 	.log_file		= "/dev/stdout",
 	.pid_file		= NULL,
+	.dns_servers		= "1.1.1.1"
 };
 
 __cold
@@ -127,6 +133,10 @@ static void show_help(const char *app)
 	printf("  -m, --log-level=level           Set log level (0=none, 1=error, 2=warning, 3=info, 4=debug, default: %d)\n", default_opts.log_level);
 	printf("  -f, --log-file=file             Log to the specified file (default: %s)\n", default_opts.log_file);
 	printf("  -p, --pid-file=file             Write PID to the specified file (default is no pid file)\n");
+#ifdef CONFIG_NEW_DNS_RESOLVER
+	printf("  -j, --dns-server=addr:port      DNS server address (default: system resolver)\n");
+	printf("  -r, --raw-dns=0|1               Use raw DNS for SOCKS5 (default: %d)\n", default_opts.use_raw_dns);
+#endif
 	printf("\n");
 }
 
@@ -228,6 +238,12 @@ static int parse_options(int argc, char *argv[], struct gwp_cfg *cfg)
 		case 'p':
 			cfg->pid_file = optarg;
 			break;
+		case 'j':
+			cfg->dns_servers = optarg;
+			break;
+		case 'r':
+			cfg->use_raw_dns = !!atoi(optarg);
+			break;
 		default:
 			fprintf(stderr, "Unknown option: %c\n", c);
 			show_help(argv[0]);
@@ -235,6 +251,12 @@ static int parse_options(int argc, char *argv[], struct gwp_cfg *cfg)
 		}
 	}
 
+
+	if (cfg->use_raw_dns && !strcmp(cfg->event_loop, "io_uring")) {
+		fprintf(stderr, ERR_WRAP "Error: The raw DNS feature is currently not supported with the io_uring event loop\n" ERR_WRAP);
+		goto einval;
+	}
+
 	if (!cfg->as_socks5 && !cfg->as_http && !cfg->target) {
 		fprintf(stderr, ERR_WRAP "Error: --target is required unless --as-socks5=1 or --as-http=1\n" ERR_WRAP);
 		goto einval;
@@ -275,6 +297,138 @@ einval:
 	return -EINVAL;
 }
 
+#ifdef CONFIG_NEW_DNS_RESOLVER
+static int gwp_ctx_init_raw_dns(struct gwp_wrk *w)
+{
+	struct gwp_ctx *ctx = w->ctx;
+	int r;
+
+	w->dns = calloc(1, sizeof(*w->dns));
+	if (!w->dns)
+		return -ENOMEM;
+
+	/*
+	 * TODO(Viro_SSFS): Support multiple DNS servers by splitting
+	 *                  ctx->cfg.dns_servers by comma.
+	 *
+	 * For now, support only one DNS server.
+	 */
+	w->dns->nr = 1;
+	w->dns->resolvers = calloc(w->dns->nr, sizeof(*w->dns->resolvers));
+	if (!w->dns->resolvers) {
+		r = -ENOMEM;
+		goto err_dns;
+	}
+
+	r = gwp_dns_res_init(ctx, &w->dns->resolvers[0], ctx->cfg.dns_servers);
+	if (r < 0)
+		goto err_resolvers;
+
+	pr_dbg(&ctx->lh, "Worker %u initialized raw DNS resolver: %s (fd=%d)",
+		w->idx, ctx->cfg.dns_servers, w->dns->resolvers[0].udp_fd);
+	return 0;
+
+err_resolvers:
+	free(w->dns->resolvers);
+	w->dns->resolvers = NULL;
+err_dns:
+	free(w->dns);
+	w->dns = NULL;
+	return r;
+}
+
+static void gwp_ctx_free_raw_dns(struct gwp_wrk *w)
+{
+	size_t i;
+
+	if (!w->dns)
+		return;
+
+	if (w->dns->resolvers) {
+		for (i = 0; i < w->dns->nr; i++)
+			gwp_dns_res_free(&w->dns->resolvers[i]);
+		free(w->dns->resolvers);
+		w->dns->resolvers = NULL;
+	}
+
+	free(w->dns);
+	w->dns = NULL;
+}
+
+static int gwp_alloc_dns_packet(struct gwp_conn_pair *gcp)
+{
+	gcp->gdp = calloc(1, sizeof(*gcp->gdp));
+	return gcp->gdp ? 0 : -ENOMEM;
+}
+
+static void gwp_free_dns_packet(struct gwp_conn_pair *gcp)
+{
+	if (gcp->gdp) {
+		free(gcp->gdp->host);
+		free(gcp->gdp);
+		gcp->gdp = NULL;
+	}
+}
+
+static int gwp_raw_dns_resolve(struct gwp_wrk *w,
+			       struct gwp_conn_pair *gcp,
+			       const char *host,
+			       const char *port)
+{
+	struct gwp_dns_resolver *res;
+	struct gwp_dns_packet *gdp;
+	int r;
+
+	assert(w->dns);
+
+	res = &w->dns->resolvers[0];
+	r = gwp_alloc_dns_packet(gcp);
+	if (r < 0)
+		return r;
+
+	gdp = gcp->gdp;
+	gdp->host = strdup(host);
+	if (!gdp->host) {
+		r = -ENOMEM;
+		goto out_err;
+	}
+
+	gdp->restyp = w->ctx->cfg.socks5_prefer_ipv6 ?
+		GWP_DNS_RESTYP_PREFER_IPV6 :
+		GWP_DNS_RESTYP_PREFER_IPV4;
+
+	gdp->port = (uint16_t)atoi(port);
+	gdp->buf_len = sizeof(gdp->buf);
+	gdp->gcp = gcp;
+	r = gwp_dns_res_prep_query(res, gdp);
+	if (r < 0)
+		goto out_err;
+
+	return -EINPROGRESS;
+
+out_err:
+	gwp_free_dns_packet(gcp);
+	return r;
+}
+#else /* #ifdef CONFIG_NEW_DNS_RESOLVER */
+static int gwp_ctx_init_raw_dns(struct gwp_wrk __unused *w)
+{
+	return -ENOSYS;
+}
+
+static void gwp_ctx_free_raw_dns(struct gwp_wrk __unused *w)
+{
+}
+
+static int gwp_raw_dns_resolve(struct gwp_wrk __unused *w,
+			       struct gwp_conn_pair __unused *gcp,
+			       const char __unused *host,
+			       const char __unused *port)
+{
+	return -ENOSYS;
+}
+#endif /* #ifdef CONFIG_NEW_DNS_RESOLVER */
+
 #define FULL_ADDRSTRLEN (INET6_ADDRSTRLEN + sizeof(":65535[]") - 1)
 
 __hot
@@ -451,6 +605,7 @@ static int gwp_ctx_init_thread(struct gwp_wrk *w,
 			       const struct gwp_sockaddr *bind_addr)
 {
 	struct gwp_ctx *ctx = w->ctx;
+	struct gwp_cfg *cfg = &ctx->cfg;
 	int r;
 
 	r = gwp_ctx_init_thread_sock(w, bind_addr);
@@ -459,13 +614,29 @@ static int gwp_ctx_init_thread(struct gwp_wrk *w,
 		return r;
 	}
 
+
+	if (cfg->use_raw_dns) {
+		r = gwp_ctx_init_raw_dns(w);
+		if (r < 0) {
+			pr_err(&ctx->lh, "Failed to initialize raw DNS: %s", strerror(-r));
+			goto out_err;
+		}
+	}
+
 	r = gwp_ctx_init_thread_event(w);
 	if (r < 0) {
 		pr_err(&ctx->lh, "gwp_ctx_init_thread_event: %s\n", strerror(-r));
-		gwp_ctx_free_thread_sock(w);
+		goto out_err_raw_dns;
 	}
 
 	return r;
+
+out_err_raw_dns:
+	if (cfg->use_raw_dns)
+		gwp_ctx_free_raw_dns(w);
+out_err:
+	gwp_ctx_free_thread_sock(w);
+	return r;
 }
 
 static void free_conn(struct gwp_conn *conn);
@@ -529,6 +700,11 @@ static void gwp_ctx_signal_all_workers(struct gwp_ctx *ctx)
 __cold
 static void gwp_ctx_free_thread(struct gwp_wrk *w)
 {
+	struct gwp_cfg *cfg = &w->ctx->cfg;
+
+	if (cfg->use_raw_dns)
+		gwp_ctx_free_raw_dns(w);
+
 	gwp_ctx_free_thread_sock_pairs(w);
 	gwp_ctx_free_thread_sock(w);
 	gwp_ctx_free_thread_event(w);
@@ -695,7 +871,7 @@ static int gwp_ctx_init_dns(struct gwp_ctx *ctx)
 	};
 	int r;
 
-	if (!cfg->as_socks5 && !cfg->as_http) {
+	if ((!cfg->as_socks5 && !cfg->as_http) || cfg->use_raw_dns) {
 		ctx->dns = NULL;
 		return 0;
 	}
@@ -966,6 +1142,7 @@ __hot
 int gwp_free_conn_pair(struct gwp_wrk *w, struct gwp_conn_pair *gcp)
 {
 	struct gwp_conn_slot *gcs = &w->conn_slot;
+	struct gwp_cfg *cfg = &w->ctx->cfg;
 	struct gwp_conn_pair *tmp;
 	uint32_t i = gcp->idx;
 
@@ -990,7 +1167,9 @@ int gwp_free_conn_pair(struct gwp_wrk *w, struct gwp_conn_pair *gcp)
 	if (gcp->timer_fd >= 0)
 		__sys_close(gcp->timer_fd);
 
-	if (gcp->gde)
+	if (cfg->use_raw_dns && gcp->gdp)
+		free(gcp->gdp);
+	else if (gcp->gde)
 		gwp_dns_entry_put(gcp->gde);
 
 	switch (gcp->prot_type) {
@@ -1207,16 +1386,21 @@ static int prepare_target_addr_domain(struct gwp_wrk *w,
 				      const char *host, const char *port)
 {
 	struct gwp_ctx *ctx = w->ctx;
+	struct gwp_cfg *cfg = &ctx->cfg;
 	int r;
 
-	r = gwp_dns_cache_lookup(ctx->dns, host, port, &gcp->target_addr);
-	if (!r) {
-		pr_dbg(&ctx->lh, "Found %s:%s in DNS cache %s", host, port,
-			ip_to_str(&gcp->target_addr));
-		return 0;
-	}
+	if (cfg->use_raw_dns) {
+		return gwp_raw_dns_resolve(w, gcp, host, port);
+	} else {
+		r = gwp_dns_cache_lookup(ctx->dns, host, port, &gcp->target_addr);
+		if (!r) {
+			pr_dbg(&ctx->lh, "Found %s:%s in DNS cache %s", host, port,
+				ip_to_str(&gcp->target_addr));
+			return 0;
+		}
 
-	return queue_dns_resolution(w, gcp, host, port);
+		return queue_dns_resolution(w, gcp, host, port);
+	}
 }
 
 static int socks5_prepare_target_addr_domain(struct gwp_wrk *w,
@@ -1546,6 +1730,7 @@ static void *gwp_ctx_thread_entry(void *arg)
 		r = -EINVAL;
 		break;
 	}
+
 	ctx->stop = true;
 	gwp_ctx_signal_all_workers(ctx);
 	pr_info(&ctx->lh, "Worker %u stopped", w->idx);
diff --git a/src/gwproxy/gwproxy.h b/src/gwproxy/gwproxy.h
index 445846f..82bd8d8 100644
--- a/src/gwproxy/gwproxy.h
+++ b/src/gwproxy/gwproxy.h
@@ -18,6 +18,10 @@
 #include <liburing.h>
 #endif
 
+#ifdef CONFIG_NEW_DNS_RESOLVER
+#include <gwproxy/dns_resolver.h>
+#endif
+
 #include <gwproxy/http1.h>
 
 struct gwp_cfg {
@@ -27,6 +31,7 @@ struct gwp_cfg {
 	bool		as_socks5;
 	bool		as_http;
 	bool		socks5_prefer_ipv6;
+	bool		use_raw_dns;
 	int		protocol_timeout;
 	const char	*socks5_auth_file;
 	int		socks5_dns_cache_secs;
@@ -44,6 +49,7 @@ struct gwp_cfg {
 	int		log_level;
 	const char	*log_file;
 	const char	*pid_file;
+	const char	*dns_servers;
 };
 
 struct gwp_ctx;
@@ -59,6 +65,7 @@ enum {
 	EV_BIT_SOCKS5_AUTH_FILE		= (8ull << 48ull),
 
 	EV_BIT_HTTP_CONN		= (18ull << 48ull),
+	EV_BIT_RAW_DNS_QUERY		= (19ull << 48ull),
 
 	/*
 	 * This ev_bit is used for user_data masking during protocol
@@ -153,6 +160,8 @@ struct gwp_http_conn {
 	struct gwnet_http_req_hdr	req_hdr;
 };
 
+struct gwp_dns_packet;
+
 struct gwp_conn_pair {
 	struct gwp_conn		target;
 	struct gwp_conn		client;
@@ -172,7 +181,10 @@ struct gwp_conn_pair {
 		struct gwp_socks5_conn	*s5_conn;
 		struct gwp_http_conn	*http_conn;
 	};
-	struct gwp_dns_entry	*gde;
+	union {
+		struct gwp_dns_entry	*gde;
+		struct gwp_dns_packet	*gdp;
+	};
 	struct gwp_sockaddr	client_addr;
 	struct gwp_sockaddr	target_addr;
 };
@@ -192,6 +204,13 @@ struct iou {
 };
 #endif
 
+struct gwp_dns_resolver;
+
+struct gwp_wrk_dns {
+	struct gwp_dns_resolver		*resolvers;
+	uint32_t			nr;
+};
+
 struct gwp_wrk {
 	int			tcp_fd;
 	struct gwp_conn_slot	conn_slot;
@@ -218,6 +237,10 @@ struct gwp_wrk {
 	struct gwp_ctx		*ctx;
 	uint32_t		idx;
 	pthread_t		thread;
+
+#ifdef CONFIG_NEW_DNS_RESOLVER
+	struct gwp_wrk_dns	*dns;
+#endif
 };
 
 enum {
-- 
Alviro Iskandar Setiawan