[PATCH gwproxy v11 4/6] gwproxy: Refactor code base to add experimental raw DNS backend

[Ahmad Gani <reyuki@gnuweeb.org>]

The raw DNS backend is marked as experimental and disabled by default.
It can be enabled at build time with the configure option --use-raw-dns,
and at runtime with the gwproxy option -r 1.

When experimental raw DNS is enabled in gwproxy, the program benefits from:
- disabling DNS worker threads
- disabling eventfd that used to communicate between threads
- disabling condition variables and mutex locks
- disabling reference counting (no need to prevent UAF)

Use fallback mechanism for raw DNS
If *_PREFER_* is used in the restyp, the raw DNS backend will attempt to
retry DNS query with different address family if it failed.

Add new cmdline arguments:
- raw-dns: enables the experimental raw DNS feature
- dns-server: sets the default DNS server
- session-map-cap: sets the default capacity for the session_map and stack.arr arrays

The struct gwp_dns_ctx is now exported due to the need for accessing it
outside dns.c.

Signed-off-by: Ahmad Gani <reyuki@gnuweeb.org>
---
 src/gwproxy/dns.c      | 269 ++++++++++++++++++++++++++++++++++-------
 src/gwproxy/dns.h      |  57 ++++++++-
 src/gwproxy/ev/epoll.c | 142 ++++++++++++++++++++--
 src/gwproxy/gwproxy.c  | 199 +++++++++++++++++++++++++++++-
 src/gwproxy/gwproxy.h  |  88 ++++++++++++++
 5 files changed, 688 insertions(+), 67 deletions(-)

diff --git a/src/gwproxy/dns.c b/src/gwproxy/dns.c
index bd8c90fb1449..c881b034611f 100644
--- a/src/gwproxy/dns.c
+++ b/src/gwproxy/dns.c
@@ -23,8 +23,7 @@
 
 #include <pthread.h>
 #include <sys/eventfd.h>
-
-struct gwp_dns_ctx;
+#include <gwproxy/dnsparser.h>
 
 struct gwp_dns_wrk {
 	struct gwp_dns_ctx	*ctx;
@@ -32,20 +31,6 @@ struct gwp_dns_wrk {
 	pthread_t		thread;
 };
 
-struct gwp_dns_ctx {
-	volatile bool		should_stop;
-	pthread_mutex_t		lock;
-	pthread_cond_t		cond;
-	uint32_t		nr_sleeping;
-	uint32_t		nr_entries;
-	struct gwp_dns_entry	*head;
-	struct gwp_dns_entry	*tail;
-	struct gwp_dns_wrk	*workers;
-	struct gwp_dns_cache	*cache;
-	time_t			last_scan;
-	struct gwp_dns_cfg	cfg;
-};
-
 static void put_all_entries(struct gwp_dns_entry *head)
 {
 	struct gwp_dns_entry *e, *next;
@@ -182,6 +167,174 @@ int gwp_dns_resolve(struct gwp_dns_ctx *ctx, const char *name,
 	return found ? 0 : -EHOSTUNREACH;
 }
 
+#ifdef CONFIG_RAW_DNS
+
+static void _gwp_dns_entry_free(struct gwp_dns_entry *e)
+{
+	assert(e);
+	free(e->name);
+	free(e);
+}
+
+void gwp_dns_raw_entry_free(struct gwp_dns_ctx *ctx, struct gwp_dns_entry *e)
+{
+	struct gwp_dns_entry *new_e;
+
+	assert(e);
+
+	new_e = ctx->entries[--ctx->nr_entries];
+	assert(ctx->nr_entries == new_e->idx);
+	new_e->idx = e->idx;
+	ctx->entries[e->idx] = new_e;
+	ctx->entries[ctx->nr_entries] = NULL;
+
+	_gwp_dns_entry_free(e);
+}
+
+static void free_all_queued_entries(struct gwp_dns_ctx *ctx)
+{
+	uint32_t i;
+	for (i = 0; i < ctx->nr_entries; i++) {
+		struct gwp_dns_entry *e = ctx->entries[i];
+		_gwp_dns_entry_free(e);
+	}
+
+	free(ctx->entries);
+}
+
+static bool realloc_entries(struct gwp_dns_ctx *ctx)
+{
+	struct gwp_dns_entry **tmp;
+	int new_cap;
+
+	new_cap = ctx->entry_cap * 2;
+	tmp = realloc(ctx->entries, new_cap * sizeof(*tmp));
+	if (!tmp)
+		return 1;
+
+	ctx->entries = tmp;
+	ctx->entry_cap = new_cap;
+
+	return 0;
+}
+
+static int attempt_fallback(struct gwp_dns_ctx *ctx, struct gwp_dns_entry *e)
+{
+	int af, r;
+
+	/* Fallback to other family if this one yield no results */
+	switch (ctx->cfg.restyp) {
+	case GWP_DNS_RESTYP_PREFER_IPV4:
+		af = AF_INET6;
+		break;
+	case GWP_DNS_RESTYP_PREFER_IPV6:
+		af = AF_INET;
+		break;
+	default:
+		assert(0);
+		return -EINVAL;
+	}
+
+	r = (int)gwdns_build_query(e->txid,
+				e->name, af, e->payload, sizeof(e->payload));
+	if (r > 0) {
+		e->payloadlen = r;
+		r = -EAGAIN;
+	}
+
+	return r;
+}
+
+int gwp_dns_process(uint8_t buff[UDP_MSG_LIMIT], int bufflen, struct gwp_dns_ctx *ctx, struct gwp_dns_entry *e)
+{
+	struct gwdns_addrinfo_node *ai;
+	ssize_t r;
+
+	r = gwdns_parse_query(e->txid, e->service, buff, bufflen, &ai);
+	if (r) {
+		if (r == -ENODATA)
+			r = attempt_fallback(ctx, e);
+		goto exit;
+	}
+
+	e->addr = ai->ai_addr;
+
+	gwdns_free_parsed_query(ai);
+exit:
+	return (int)r;
+}
+
+struct gwp_dns_entry *gwp_raw_dns_queue(uint16_t txid, struct gwp_dns_ctx *ctx,
+				    const char *name, const char *service)
+{
+	struct gwp_dns_entry *e;
+	size_t nl, sl;
+	ssize_t r;
+	int af;
+
+	e = malloc(sizeof(*e));
+	if (!e)
+		return NULL;
+
+	if (ctx->nr_entries == ctx->entry_cap && realloc_entries(ctx))
+		return NULL;
+
+	switch (ctx->cfg.restyp) {
+	case GWP_DNS_RESTYP_PREFER_IPV4:
+	case GWP_DNS_RESTYP_IPV4_ONLY:
+	case GWP_DNS_RESTYP_DEFAULT:
+		af = AF_INET;
+		break;
+	case GWP_DNS_RESTYP_PREFER_IPV6:
+	case GWP_DNS_RESTYP_IPV6_ONLY:
+		af = AF_INET6;
+		break;
+	default:
+		assert(0);
+		goto out_free_e;
+	}
+
+	r = gwdns_build_query(txid,
+				name, af, e->payload, sizeof(e->payload));
+	if (r < 0)
+		goto out_free_e;
+	e->payloadlen = (int)r;
+
+	/*
+	 * Merge name and service into a single allocated string to
+	 * avoid multiple allocations.
+	 *
+	 * The format is: "<name>\0<service>\0" where both name and
+	 * service are null-terminated strings.
+	 */
+	nl = strlen(name);
+	sl = service ? strlen(service) : 0;
+	e->name = malloc(nl + 1 + sl + 1);
+	if (!e->name)
+		goto out_free_e;
+
+	e->service = e->name + nl + 1;
+	memcpy(e->name, name, nl + 1);
+	if (service)
+		memcpy(e->service, service, sl + 1);
+	else
+		e->service[0] = '\0';
+
+	e->res = 0;
+	e->idx = ctx->nr_entries++;
+	ctx->entries[e->idx] = e;
+
+	return e;
+out_free_e:
+	free(e);
+	return NULL;
+}
+#else
+static void free_all_queued_entries(__maybe_unused struct gwp_dns_ctx *ctx)
+{
+}
+#endif /* #ifdef CONFIG_RAW_DNS */
+
 static void gwp_dns_entry_free(struct gwp_dns_entry *e)
 {
 	if (!e)
@@ -713,39 +866,57 @@ int gwp_dns_ctx_init(struct gwp_dns_ctx **ctx_p, const struct gwp_dns_cfg *cfg)
 		return -ENOMEM;
 
 	ctx->cfg = *cfg;
-	r = pthread_mutex_init(&ctx->lock, NULL);
-	if (r) {
-		r = -r;
-		goto out_free_ctx;
-	}
 
-	r = pthread_cond_init(&ctx->cond, NULL);
-	if (r) {
-		r = -r;
-		goto out_destroy_mutex;
-	}
+	if (cfg->use_raw_dns) {
+#ifdef CONFIG_RAW_DNS
+		r = convert_str_to_ssaddr(cfg->ns_addr_str, &ctx->ns_addr, 53);
+		if (r)
+			goto out_free_ctx;
+		ctx->ns_addrlen = ctx->ns_addr.sa.sa_family == AF_INET
+				? sizeof(ctx->ns_addr.i4)
+				: sizeof(ctx->ns_addr.i6);
+		ctx->entry_cap = DEFAULT_ENTRIES_CAP;
+		ctx->entries = malloc(ctx->entry_cap * sizeof(*ctx->entries));
+		if (!ctx->entries) {
+			__sys_close(r);
+			r = -ENOMEM;
+			goto out_free_ctx;
+		}
+#endif
+	} else {
+		r = pthread_mutex_init(&ctx->lock, NULL);
+		if (r) {
+			r = -r;
+			goto out_free_ctx;
+		}
 
-	r = init_cache(ctx);
-	if (r)
-		goto out_destroy_cond;
+		r = pthread_cond_init(&ctx->cond, NULL);
+		if (r) {
+			r = -r;
+			goto out_destroy_mutex;
+		}
 
-	ctx->nr_sleeping = 0;
-	ctx->nr_entries = 0;
-	ctx->workers = NULL;
-	ctx->head = NULL;
-	ctx->tail = NULL;
-	ctx->should_stop = false;
-	ctx->last_scan = time(NULL);
-	r = init_workers(ctx);
-	if (r)
-		goto out_free_cache;
+		ctx->nr_sleeping = 0;
+		ctx->workers = NULL;
+		ctx->head = NULL;
+		ctx->tail = NULL;
+		r = init_workers(ctx);
+		if (r)
+			goto out_destroy_cond;
+		r = init_cache(ctx);
+		if (r)
+			goto out_destroy_cond;
+
+		ctx->should_stop = false;
+		ctx->last_scan = time(NULL);
+	}
 
+	ctx->nr_entries = 0;
 	*ctx_p = ctx;
 	return 0;
-out_free_cache:
-	free_cache(ctx->cache);
 out_destroy_cond:
-	pthread_cond_destroy(&ctx->cond);
+	if (!cfg->use_raw_dns)
+		pthread_cond_destroy(&ctx->cond);
 out_destroy_mutex:
 	pthread_mutex_destroy(&ctx->lock);
 out_free_ctx:
@@ -762,11 +933,15 @@ static void put_all_queued_entries(struct gwp_dns_ctx *ctx)
 
 void gwp_dns_ctx_free(struct gwp_dns_ctx *ctx)
 {
-	free_workers(ctx);
-	pthread_mutex_destroy(&ctx->lock);
-	pthread_cond_destroy(&ctx->cond);
-	put_all_queued_entries(ctx);
-	free_cache(ctx->cache);
+	if (ctx->cfg.use_raw_dns) {
+		free_all_queued_entries(ctx);
+	} else {
+		free_workers(ctx);
+		pthread_mutex_destroy(&ctx->lock);
+		pthread_cond_destroy(&ctx->cond);
+		put_all_queued_entries(ctx);
+		free_cache(ctx->cache);
+	}
 	free(ctx);
 }
 
diff --git a/src/gwproxy/dns.h b/src/gwproxy/dns.h
index 420a80e5cf6a..049505997e37 100644
--- a/src/gwproxy/dns.h
+++ b/src/gwproxy/dns.h
@@ -9,12 +9,20 @@
 #include <stdatomic.h>
 #include <stdbool.h>
 #include <netinet/in.h>
+#include <gwproxy/common.h>
 #include <gwproxy/net.h>
+#include <gwproxy/dnsparser.h>
 #include <gwproxy/syscall.h>
 
-struct gwp_dns_wrk;
-
 struct gwp_dns_entry {
+#ifdef CONFIG_RAW_DNS
+	uint32_t		idx;
+	int			payloadlen;
+	union {
+		uint16_t	txid;
+		uint8_t		payload[UDP_MSG_LIMIT];
+	};
+#endif
 	char			*name;
 	char			*service;
 	_Atomic(int)		refcnt;
@@ -32,13 +40,37 @@ enum {
 	GWP_DNS_RESTYP_PREFER_IPV6	= 4,
 };
 
+#define DEFAULT_ENTRIES_CAP 255
+
 struct gwp_dns_cfg {
 	int		cache_expiry;	/* In seconds. <= 0 to disable cache. */
 	uint32_t	nr_workers;
 	uint32_t	restyp;
+	bool		use_raw_dns;
+#ifdef CONFIG_RAW_DNS
+	const char	*ns_addr_str;
+#endif
 };
 
-struct gwp_dns_ctx;
+struct gwp_dns_ctx {
+#ifdef CONFIG_RAW_DNS
+	uint32_t		entry_cap;
+	struct gwp_dns_entry	**entries;
+	struct gwp_sockaddr	ns_addr;
+	socklen_t		ns_addrlen;
+#endif
+	volatile bool		should_stop;
+	pthread_mutex_t		lock;
+	pthread_cond_t		cond;
+	uint32_t		nr_sleeping;
+	uint32_t		nr_entries;
+	struct gwp_dns_entry	*head;
+	struct gwp_dns_entry	*tail;
+	struct gwp_dns_wrk	*workers;
+	struct gwp_dns_cache	*cache;
+	time_t			last_scan;
+	struct gwp_dns_cfg	cfg;
+};
 
 /**
  * Initialize the DNS context. Stores the context in `*ctx_p`. When
@@ -88,6 +120,25 @@ struct gwp_dns_entry *gwp_dns_queue(struct gwp_dns_ctx *ctx,
  */
 bool gwp_dns_entry_put(struct gwp_dns_entry *entry);
 
+#ifdef CONFIG_RAW_DNS
+struct gwp_dns_entry *gwp_raw_dns_queue(uint16_t txid, struct gwp_dns_ctx *ctx,
+				    const char *name, const char *service);
+
+void gwp_dns_raw_entry_free(struct gwp_dns_ctx *ctx, struct gwp_dns_entry *e);
+
+int gwp_dns_process(uint8_t buff[UDP_MSG_LIMIT], int bufflen, struct gwp_dns_ctx *ctx, struct gwp_dns_entry *e);
+#else
+static inline struct gwp_dns_entry *gwp_raw_dns_queue(__maybe_unused uint16_t txid, __maybe_unused struct gwp_dns_ctx *ctx,
+				    __maybe_unused const char *name, __maybe_unused const char *service)
+{
+	return NULL;
+}
+
+static inline void gwp_dns_raw_entry_free(__maybe_unused struct gwp_dns_ctx *ctx, __maybe_unused struct gwp_dns_entry *e)
+{
+}
+#endif
+
 /**
  * Lookup a DNS entry in the cache. If the entry is found, it fills the
  * `addr` structure with the resolved address and returns 0. If the entry is
diff --git a/src/gwproxy/ev/epoll.c b/src/gwproxy/ev/epoll.c
index d46568a6a2b1..5926bb5e090a 100644
--- a/src/gwproxy/ev/epoll.c
+++ b/src/gwproxy/ev/epoll.c
@@ -39,6 +39,16 @@ int gwp_ctx_init_thread_epoll(struct gwp_wrk *w)
 		goto out_close_ep_fd;
 	}
 
+	if (w->ctx->cfg.use_raw_dns) {
+#ifdef CONFIG_RAW_DNS
+		ev.events = EPOLLIN;
+		ev.data.u64 = EV_BIT_DNS_QUERY;
+		r = __sys_epoll_ctl(ep_fd, EPOLL_CTL_ADD, w->dns_resolver.udp_fd, &ev);
+		if (unlikely(r))
+			return (int)r;
+#endif
+	}
+
 	w->evsz = 512;
 	events = calloc(w->evsz, sizeof(*events));
 	if (!events) {
@@ -778,6 +788,89 @@ static int handle_connect(struct gwp_wrk *w, struct gwp_conn_pair *gcp)
 	return 0;
 }
 
+#ifdef CONFIG_RAW_DNS
+static int send_raw_dns_query(struct gwp_wrk *w, struct gwp_conn_pair *gcp)
+{
+	struct gwp_dns_entry *gde = gcp->gde;
+	struct gwp_dns_ctx *dctx;
+	ssize_t r;
+
+	dctx = w->ctx->dns;
+	r = __sys_sendto(
+		w->dns_resolver.udp_fd, gde->payload, gde->payloadlen, MSG_NOSIGNAL,
+		&dctx->ns_addr.sa, dctx->ns_addrlen
+	);
+	if (r < 0)
+		return (int)r;
+	pr_dbg(&w->ctx->lh, "standard DNS query for %s has been sent", gde->name);
+
+	return 0;
+}
+static int read_raw_dns_answer(struct gwp_wrk *w, struct gwp_conn_pair **pgcp,
+				struct gwp_dns_entry **pgde)
+{
+	uint8_t buff[UDP_MSG_LIMIT];
+	struct gwp_conn_pair *gcp;
+	struct gwp_dns_entry *gde;
+	struct gwp_dns_ctx *dctx;
+	uint16_t txid;
+	int r;
+
+	dctx = w->ctx->dns;
+	r = (int)__sys_recvfrom(
+		w->dns_resolver.udp_fd, buff, sizeof(buff), 0,
+		&dctx->ns_addr.sa, &dctx->ns_addrlen
+	);
+	if (r < 0)
+		return (int)r;
+	if (r < 38)
+		return -EINVAL;
+
+	memcpy(&txid, buff, 2);
+	if (txid >= w->dns_resolver.sess_map_cap) {
+		pr_err(&w->ctx->lh, "txid is bigger than sess_map_cap: %hu\n", txid);
+		return -EINVAL;
+	}
+
+	gcp = w->dns_resolver.sess_map[txid];
+	if (!gcp)
+		return -ENOENT;
+
+	gde = gcp->gde;
+	pr_dbg(&w->ctx->lh, "proxy session for DNS query %s was found!", gde->name);
+	r = gwp_dns_process(buff, r, dctx, gde);
+	if (r == -EAGAIN) {
+		pr_dbg(&w->ctx->lh, "DNS Fallback");
+		r = send_raw_dns_query(w, gcp);
+		if (r)
+			return r;
+
+		return -EAGAIN;
+	} else if (r)
+		gde->res = r;
+
+	r = return_txid(w, gde);
+	assert(!r);
+	if (r)
+		return r;
+
+	*pgcp = gcp;
+	*pgde = gde;
+	return 0;
+}
+#else
+static int send_raw_dns_query(__maybe_unused struct gwp_wrk *w, __maybe_unused struct gwp_conn_pair *gcp)
+{
+	return 0;
+}
+static int read_raw_dns_answer(__maybe_unused struct gwp_wrk *w,
+			__maybe_unused struct gwp_conn_pair **pgcp,
+			__maybe_unused struct gwp_dns_entry **pgde)
+{
+	return 0;
+}
+#endif
+
 static int arm_poll_for_dns_query(struct gwp_wrk *w,
 					struct gwp_conn_pair *gcp)
 {
@@ -786,16 +879,22 @@ static int arm_poll_for_dns_query(struct gwp_wrk *w,
 	int r;
 
 	assert(gde);
-	assert(gde->ev_fd >= 0);
 
 	ev.events = EPOLLIN;
 	ev.data.u64 = 0;
 	ev.data.ptr = gcp;
 	ev.data.u64 |= EV_BIT_DNS_QUERY;
 
-	r = __sys_epoll_ctl(w->ep_fd, EPOLL_CTL_ADD, gde->ev_fd, &ev);
-	if (unlikely(r))
-		return r;
+	if (w->ctx->cfg.use_raw_dns) {
+		r = send_raw_dns_query(w, gcp);
+		if (r)
+			return r;
+	} else {
+		assert(gde->ev_fd >= 0);
+		r = __sys_epoll_ctl(w->ep_fd, EPOLL_CTL_ADD, gde->ev_fd, &ev);
+		if (unlikely(r))
+			return r;
+	}
 
 	return 0;
 }
@@ -822,18 +921,30 @@ static void log_dns_query(struct gwp_wrk *w, struct gwp_conn_pair *gcp,
 __hot
 static int handle_ev_dns_query(struct gwp_wrk *w, struct gwp_conn_pair *gcp)
 {
-	struct gwp_dns_entry *gde = gcp->gde;
-	int r, ct = gcp->conn_state;
+	struct gwp_dns_entry *gde = NULL;
+	int r, ct;
 
-	assert(gde);
-	assert(gde->ev_fd >= 0);
+	if (w->ctx->cfg.use_raw_dns) {
+		r = read_raw_dns_answer(w, &gcp, &gde);
+		if (r) {
+			if (r == -EAGAIN)
+				return 0;
+			else
+				return r;
+		}
+	} else {
+		gde = gcp->gde;
+		assert(gde);
+		assert(gde->ev_fd >= 0);
+		r = __sys_epoll_ctl(w->ep_fd, EPOLL_CTL_DEL, gde->ev_fd, NULL);
+		if (unlikely(r))
+			return r;
+	}
+
+	ct = gcp->conn_state;
 	assert(ct == CONN_STATE_SOCKS5_DNS_QUERY ||
 	       ct == CONN_STATE_HTTP_DNS_QUERY);
 
-	r = __sys_epoll_ctl(w->ep_fd, EPOLL_CTL_DEL, gde->ev_fd, NULL);
-	if (unlikely(r))
-		return r;
-
 	log_dns_query(w, gcp, gde);
 	if (likely(!gde->res)) {
 		gcp->target_addr = gde->addr;
@@ -845,7 +956,12 @@ static int handle_ev_dns_query(struct gwp_wrk *w, struct gwp_conn_pair *gcp)
 			r = -EIO;
 	}
 
-	gwp_dns_entry_put(gde);
+	if (w->ctx->cfg.use_raw_dns) {
+		pr_dbg(&w->ctx->lh, "removing proxy session from the %s's txid map", gde->name);
+		gwp_dns_raw_entry_free(w->ctx->dns, gde);
+	} else {
+		gwp_dns_entry_put(gde);
+	}
 	gcp->gde = NULL;
 	return r;
 }
diff --git a/src/gwproxy/gwproxy.c b/src/gwproxy/gwproxy.c
index 3d277f7ee6f8..e14780f23945 100644
--- a/src/gwproxy/gwproxy.c
+++ b/src/gwproxy/gwproxy.c
@@ -44,6 +44,11 @@
 
 static const struct option long_opts[] = {
 	{ "help",		no_argument,		NULL,	'h' },
+	{ "raw-dns",		required_argument,	NULL,	'r' },
+#ifdef CONFIG_RAW_DNS
+	{ "dns-server",		required_argument,	NULL,	'j' },
+	{ "session-map-cap",	required_argument,	NULL,	's' },
+#endif
 	{ "event-loop",		required_argument,	NULL,	'e' },
 	{ "bind",		required_argument,	NULL,	'b' },
 	{ "target",		required_argument,	NULL,	't' },
@@ -74,6 +79,7 @@ static const struct gwp_cfg default_opts = {
 	.event_loop		= "epoll",
 	.bind			= "[::]:1080",
 	.target			= NULL,
+	.use_raw_dns		= false,
 	.as_socks5		= false,
 	.as_http		= false,
 	.socks5_prefer_ipv6	= false,
@@ -94,6 +100,10 @@ static const struct gwp_cfg default_opts = {
 	.log_level		= 3,
 	.log_file		= "/dev/stdout",
 	.pid_file		= NULL,
+#ifdef CONFIG_RAW_DNS
+	.ns_addr_str		= "1.1.1.1",
+	.sess_map_cap		= 16
+#endif
 };
 
 __cold
@@ -104,6 +114,10 @@ static void show_help(const char *app)
 	printf("  -h, --help                      Show this help message and exit\n");
 	printf("  -e, --event-loop=name           Specify the event loop to use (default: %s)\n", default_opts.event_loop);
 	printf("                                  Available values: epoll, io_uring\n");
+	printf("  -r, --raw-dns=0|1               Use experimental raw DNS as the backend (default: %d)\n", default_opts.use_raw_dns);
+#ifdef CONFIG_RAW_DNS
+	printf("  -j, --dns-server=addr           DNS server address (default: %s)\n", default_opts.ns_addr_str);
+#endif
 	printf("  -b, --bind=addr:port            Bind to the specified address (default: %s)\n", default_opts.bind);
 	printf("  -t, --target=addr_port          Target address to connect to\n");
 	printf("  -S, --as-socks5=0|1             Run as a SOCKS5 proxy (default: %d)\n", default_opts.as_socks5);
@@ -159,6 +173,9 @@ static int parse_options(int argc, char *argv[], struct gwp_cfg *cfg)
 		case 'h':
 			show_help(argv[0]);
 			exit(0);
+		case 'r':
+			cfg->use_raw_dns = !!atoi(optarg);
+			break;
 		case 'e':
 			cfg->event_loop = optarg;
 			break;
@@ -228,6 +245,14 @@ static int parse_options(int argc, char *argv[], struct gwp_cfg *cfg)
 		case 'p':
 			cfg->pid_file = optarg;
 			break;
+#ifdef CONFIG_RAW_DNS
+		case 'j':
+			cfg->ns_addr_str = optarg;
+			break;
+		case 's':
+			cfg->sess_map_cap = atoi(optarg);
+			break;
+#endif
 		default:
 			fprintf(stderr, "Unknown option: %c\n", c);
 			show_help(argv[0]);
@@ -446,6 +471,136 @@ static void gwp_ctx_free_thread_event(struct gwp_wrk *w)
 	}
 }
 
+#ifdef CONFIG_RAW_DNS
+static inline int realloc_stack(struct dns_resolver *resolv)
+{
+	struct gwp_conn_pair **cp;
+	size_t newcap, oldcap;
+	uint16_t *stack;
+	uint16_t j, idx;
+
+	oldcap = resolv->sess_map_cap;
+	newcap = oldcap * 2;
+	if (newcap > 65536)
+		return -ENOSPC;
+
+	cp = realloc(resolv->sess_map, sizeof(*cp) * newcap);
+	if (!cp)
+		return -ENOMEM;
+
+	resolv->sess_map = cp;
+	resolv->sess_map_cap = newcap;
+	memset(&resolv->sess_map[oldcap], 0, oldcap);
+
+	stack = realloc(resolv->stack.arr, sizeof(*stack) * newcap);
+	if (!stack)
+		return -ENOMEM;
+	resolv->stack.arr = stack;
+	resolv->stack.top = newcap;
+
+	j = (uint16_t)newcap - 1;
+	for (idx = 0; j > resolv->stack.top; idx++,j--)
+		stack[idx] = j;
+
+	return 0;
+}
+
+static inline int pop_txid(struct dns_resolver *resolv)
+{
+	int r;
+	if (resolv->stack.top == 0) {
+		r = realloc_stack(resolv);
+		if (r)
+			return -EAGAIN;
+	}
+
+	return resolv->stack.arr[--resolv->stack.top];
+}
+
+static struct gwp_dns_entry *generate_gde(struct gwp_wrk *w, struct gwp_conn_pair *gcp, const char *host, const char *port)
+{
+	struct gwp_dns_entry *gde;
+	struct gwp_dns_ctx *dns = w->ctx->dns;
+	int txid;
+
+	txid = pop_txid(&w->dns_resolver);
+	if (txid < 0)
+		return  NULL;
+
+	gde = gwp_raw_dns_queue((uint16_t)txid, dns, host, port);
+	if (!gde) {
+		return_txid(w, gde);
+		return  NULL;
+	}
+	pr_dbg(&w->ctx->lh, "constructing standard DNS query packet for %s", host);
+	w->dns_resolver.sess_map[txid] = gcp;
+
+	return gde;
+}
+
+static int gwp_ctx_init_dns_resolver(struct gwp_wrk *w)
+{
+	struct dns_resolver *resolv;
+	struct gwp_dns_ctx *dns;
+	struct gwp_cfg *cfg;
+	struct gwp_ctx *ctx;
+	int udp_fd, r;
+	void *ptr;
+
+	ctx = w->ctx;
+	cfg = &ctx->cfg;
+	dns = ctx->dns;
+	udp_fd = __sys_socket(dns->ns_addr.sa.sa_family, SOCK_DGRAM | SOCK_NONBLOCK, 0);
+	if (udp_fd < 0) {
+		pr_err(&w->ctx->lh, "Failed to create udp_fd: %s\n", strerror(-udp_fd));
+		return udp_fd;
+	}
+
+	resolv = &w->dns_resolver;
+	resolv->sess_map = NULL;
+	resolv->stack.arr = NULL;
+	r = reset_stack(w);
+	if (r) {
+		if (resolv->stack.arr)
+			free(resolv->stack.arr);
+		if (resolv->sess_map)
+			free(resolv->sess_map);
+		__sys_close(udp_fd);
+
+		return r;
+	}
+
+	memset(resolv->sess_map, 0, cfg->sess_map_cap * sizeof(ptr));
+	resolv->udp_fd = udp_fd;
+
+	return 0;
+}
+
+static void gwp_ctx_free_raw_dns(struct gwp_wrk *w)
+{
+	struct dns_resolver *resolv = &w->dns_resolver;
+	__sys_close(resolv->udp_fd);
+	free(resolv->sess_map);
+	free(resolv->stack.arr);
+}
+#else
+static struct gwp_dns_entry *generate_gde(__maybe_unused struct gwp_wrk *w,
+					  __maybe_unused struct gwp_conn_pair *gcp,
+					  __maybe_unused const char *host,
+					  __maybe_unused const char *port)
+{
+	return NULL;
+}
+static int gwp_ctx_init_dns_resolver(__maybe_unused struct gwp_wrk *w)
+{
+	return 0;
+}
+
+static void gwp_ctx_free_raw_dns(__maybe_unused struct gwp_wrk *w)
+{
+}
+#endif /* #ifdef CONFIG_RAW_DNS */
+
 __cold
 static int gwp_ctx_init_thread(struct gwp_wrk *w,
 			       const struct gwp_sockaddr *bind_addr)
@@ -459,6 +614,14 @@ static int gwp_ctx_init_thread(struct gwp_wrk *w,
 		return r;
 	}
 
+	if (ctx->cfg.use_raw_dns) {
+		r = gwp_ctx_init_dns_resolver(w);
+		if (r) {
+			pr_err(&ctx->lh, "gwp_ctx_init_thread_event: %s\n", strerror(-r));
+			gwp_ctx_free_thread_sock(w);
+		}
+	}
+
 	r = gwp_ctx_init_thread_event(w);
 	if (r < 0) {
 		pr_err(&ctx->lh, "gwp_ctx_init_thread_event: %s\n", strerror(-r));
@@ -529,6 +692,8 @@ static void gwp_ctx_signal_all_workers(struct gwp_ctx *ctx)
 __cold
 static void gwp_ctx_free_thread(struct gwp_wrk *w)
 {
+	if (w->ctx->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);
@@ -689,12 +854,23 @@ static int gwp_ctx_init_dns(struct gwp_ctx *ctx)
 {
 	struct gwp_cfg *cfg = &ctx->cfg;
 	const struct gwp_dns_cfg dns_cfg = {
+		.use_raw_dns = cfg->use_raw_dns,
 		.cache_expiry = cfg->socks5_dns_cache_secs,
 		.restyp = cfg->socks5_prefer_ipv6 ? GWP_DNS_RESTYP_PREFER_IPV6 : 0,
-		.nr_workers = cfg->nr_dns_workers
+		.nr_workers = cfg->nr_dns_workers,
+#ifdef CONFIG_RAW_DNS
+		.ns_addr_str = cfg->ns_addr_str
+#endif
 	};
 	int r;
 
+	if (cfg->use_raw_dns) {
+#ifndef CONFIG_RAW_DNS
+		pr_err(&ctx->lh, "raw DNS backend is not enabled in this build");
+		return -ENOSYS;
+#endif
+	}
+
 	if (!cfg->as_socks5 && !cfg->as_http) {
 		ctx->dns = NULL;
 		return 0;
@@ -733,6 +909,10 @@ static int gwp_ctx_parse_ev(struct gwp_ctx *ctx)
 		ctx->ev_used = GWP_EV_EPOLL;
 		pr_dbg(&ctx->lh, "Using event loop: epoll");
 	} else if (!strcmp(ev, "io_uring") || !strcmp(ev, "iou")) {
+		if (ctx->cfg.use_raw_dns) {
+			pr_err(&ctx->lh, "raw DNS backend is not supported for io_uring yet");
+			return -ENOSYS;
+		}
 		ctx->ev_used = GWP_EV_IO_URING;
 		pr_dbg(&ctx->lh, "Using event loop: io_uring");
 	} else {
@@ -990,8 +1170,15 @@ 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)
-		gwp_dns_entry_put(gcp->gde);
+	if (gcp->gde) {
+		pr_dbg(&w->ctx->lh, "client disconnected before query for %s resolved", gcp->gde->name);
+		if (w->ctx->cfg.use_raw_dns) {
+			return_txid(w, gcp->gde);
+			gwp_dns_raw_entry_free(w->ctx->dns, gcp->gde);
+		} else {
+			gwp_dns_entry_put(gcp->gde);
+		}
+	}
 
 	switch (gcp->prot_type) {
 	case GWP_PROT_TYPE_SOCKS5:
@@ -1192,7 +1379,11 @@ static int queue_dns_resolution(struct gwp_wrk *w, struct gwp_conn_pair *gcp,
 	struct gwp_dns_ctx *dns = w->ctx->dns;
 	struct gwp_dns_entry *gde;
 
-	gde = gwp_dns_queue(dns, host, port);
+	if (w->ctx->cfg.use_raw_dns)
+		gde = generate_gde(w, gcp, host, port);
+	else
+		gde = gwp_dns_queue(dns, host, port);
+
 	if (unlikely(!gde)) {
 		pr_err(&w->ctx->lh, "Failed to allocate DNS entry for %s:%s", host, port);
 		return -ENOMEM;
diff --git a/src/gwproxy/gwproxy.h b/src/gwproxy/gwproxy.h
index 095c0ce700c3..9930934b893e 100644
--- a/src/gwproxy/gwproxy.h
+++ b/src/gwproxy/gwproxy.h
@@ -9,6 +9,7 @@
 #define _GNU_SOURCE
 #endif
 
+#include <gwproxy/common.h>
 #include <gwproxy/syscall.h>
 #include <gwproxy/socks5.h>
 #include <gwproxy/dns.h>
@@ -24,6 +25,7 @@ struct gwp_cfg {
 	const char	*event_loop;
 	const char	*bind;
 	const char	*target;
+	bool		use_raw_dns;
 	bool		as_socks5;
 	bool		as_http;
 	bool		socks5_prefer_ipv6;
@@ -44,6 +46,10 @@ struct gwp_cfg {
 	int		log_level;
 	const char	*log_file;
 	const char	*pid_file;
+#ifdef CONFIG_RAW_DNS
+	const char	*ns_addr_str;
+	int		sess_map_cap;
+#endif
 };
 
 struct gwp_ctx;
@@ -195,8 +201,24 @@ struct iou {
 };
 #endif
 
+struct stack_u16 {
+	uint16_t top;
+	uint16_t *arr;
+};
+
+struct dns_resolver {
+	struct stack_u16	stack;
+	struct gwp_conn_pair	**sess_map;
+	uint16_t		sess_map_cap;
+	int			udp_fd;
+};
+
 struct gwp_wrk {
 	int			tcp_fd;
+#ifdef CONFIG_RAW_DNS
+	/* For mapping DNS queries to the corresponding proxy session */
+	struct dns_resolver	dns_resolver;
+#endif
 	struct gwp_conn_slot	conn_slot;
 
 	union {
@@ -251,6 +273,72 @@ int gwp_create_timer(int fd, int sec, int nsec);
 void gwp_setup_cli_sock_options(struct gwp_wrk *w, int fd);
 const char *ip_to_str(const struct gwp_sockaddr *gs);
 
+#ifdef CONFIG_RAW_DNS
+static inline int reset_stack(struct gwp_wrk *w)
+{
+	struct dns_resolver *resolv;
+	struct gwp_cfg *cfg;
+	uint16_t *p1;
+	int i, d;
+	void *p2;
+
+	cfg = &w->ctx->cfg;
+	d = cfg->sess_map_cap;
+	resolv = &w->dns_resolver;
+	p1 = realloc(resolv->stack.arr, d * sizeof(*resolv->stack.arr));
+	if (!p1)
+		return -ENOMEM;
+	resolv->stack.arr = p1;
+	resolv->stack.top = d;
+
+	i = d;
+	resolv->stack.top = i--;
+	for (; i >= 0; i--)
+		p1[i] = i;
+
+	p2 = realloc(resolv->sess_map, d * sizeof(*resolv->sess_map));
+	if (!p2)
+		return -ENOMEM;
+	resolv->sess_map = p2;
+	memset(p2, 0, d * sizeof(*resolv->sess_map));
+	resolv->sess_map_cap = d;
+
+	return 0;
+}
+
+static inline int return_txid(struct gwp_wrk *w, struct gwp_dns_entry *gde)
+{
+	struct dns_resolver *resolv;
+	struct gwp_cfg *cfg;
+	uint16_t idx;
+
+	resolv = &w->dns_resolver;
+	cfg = &w->ctx->cfg;
+	/* the program shouldn't return more than its pop */
+	if (resolv->stack.top >= resolv->sess_map_cap) {
+		assert(0);
+		return -EINVAL;
+	}
+
+	idx = resolv->stack.top++;
+	resolv->stack.arr[idx] = gde->txid;
+	resolv->sess_map[idx] = NULL;
+	/* shrinks it when there is no active query
+	 * and the allocated size is larger than the default size.
+	 */
+	if (idx == resolv->sess_map_cap &&
+		resolv->sess_map_cap > cfg->sess_map_cap)
+		reset_stack(w);
+
+	return 0;
+}
+#else
+static inline int return_txid(__maybe_unused struct gwp_wrk *w, __maybe_unused struct gwp_dns_entry *gde)
+{
+	return 0;
+}
+#endif
+
 static inline void gwp_conn_buf_advance(struct gwp_conn *conn, size_t len)
 {
 	assert(len <= conn->len);
-- 
Ahmad Gani