From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sat, 1 Jan 2022 22:49:48 +0000 (+0000)
Subject: [Project] Further DNS over TCP architecturing
X-Git-Tag: 3.2~129^2~14
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=02e027d0c0a5ffed16e3d387c3893ce9ff5a2521;p=rspamd.git

[Project] Further DNS over TCP architecturing
---

diff --git a/contrib/librdns/dns_private.h b/contrib/librdns/dns_private.h
index 4429552bc..75fb5a3fa 100644
--- a/contrib/librdns/dns_private.h
+++ b/contrib/librdns/dns_private.h
@@ -45,6 +45,41 @@ static const int default_tcp_io_cnt = 2;
 
 #define RESOLV_CONF "/etc/resolv.conf"
 
+struct dns_header {
+	unsigned int qid :16;
+
+#if BYTE_ORDER == BIG_ENDIAN
+	unsigned int qr:1;
+	unsigned int opcode:4;
+	unsigned int aa:1;
+	unsigned int tc:1;
+	unsigned int rd:1;
+
+	unsigned int ra:1;
+	unsigned int cd : 1;
+	unsigned int ad : 1;
+	unsigned int z : 1;
+	unsigned int rcode:4;
+#else
+	unsigned int rd :1;
+	unsigned int tc :1;
+	unsigned int aa :1;
+	unsigned int opcode :4;
+	unsigned int qr :1;
+
+	unsigned int rcode :4;
+	unsigned int z : 1;
+	unsigned int ad : 1;
+	unsigned int cd : 1;
+	unsigned int ra :1;
+#endif
+
+	unsigned int qdcount :16;
+	unsigned int ancount :16;
+	unsigned int nscount :16;
+	unsigned int arcount :16;
+};
+
 /**
  * Represents DNS server
  */
@@ -110,6 +145,33 @@ enum rdns_io_channel_flags {
 
 #define IS_CHANNEL_CONNECTED(ioc) (((ioc)->flags & RDNS_CHANNEL_CONNECTED) != 0)
 #define IS_CHANNEL_ACTIVE(ioc) (((ioc)->flags & RDNS_CHANNEL_ACTIVE) != 0)
+#define IS_CHANNEL_TCP(ioc) (((ioc)->flags & RDNS_CHANNEL_TCP) != 0)
+
+/**
+ * Used to chain output DNS requests for a TCP connection
+ */
+struct rdns_tcp_output_chain {
+	uint16_t next_write_size;
+	struct rdns_request *req;
+	uint16_t cur_write;
+	struct rdns_tcp_output_chain *prev, *next;
+};
+
+/**
+ * Specific stuff for a TCP IO chain
+ */
+struct rdns_tcp_channel {
+	uint16_t next_read_size;
+	unsigned char *cur_read_buf;
+	uint16_t cur_read;
+
+	/* Chained set of the planned writes */
+	struct rdns_tcp_output_chain *output_chain;
+	unsigned cur_output_chains;
+
+	void *async_read; /** async read event */
+	void *async_write; /** async write event */
+};
 
 KHASH_DECLARE(rdns_requests_hash, int, struct rdns_request *);
 /**
@@ -124,6 +186,15 @@ struct rdns_io_channel {
 	int flags; /**< see enum rdns_io_channel_flags */
 	void *async_io; /** async opaque ptr */
 	khash_t(rdns_requests_hash) *requests;
+	/*
+	 * For DNS replies parsing we use per-channel structure
+	 * which is used for two purposes:
+	 * 1) We read the next DNS header
+	 * 2) We find the corresponding request (if any)
+	 * 3) We read the remaining packet (associated with a request or dangling)
+	 * This structure is filled on each read-readiness for an IO channel
+	 */
+	struct rdns_tcp_channel *tcp;
 	uint64_t uses;
 	ref_entry_t ref;
 };
@@ -167,46 +238,10 @@ struct rdns_resolver {
 	ref_entry_t ref;
 };
 
-struct dns_header;
 struct dns_query;
 
 /* Internal DNS structs */
 
-struct dns_header {
-	unsigned int qid :16;
-
-#if BYTE_ORDER == BIG_ENDIAN
-	unsigned int qr:1;
-	unsigned int opcode:4;
-	unsigned int aa:1;
-	unsigned int tc:1;
-	unsigned int rd:1;
-
-	unsigned int ra:1;
-	unsigned int cd : 1;
-	unsigned int ad : 1;
-	unsigned int z : 1;
-	unsigned int rcode:4;
-#else
-	unsigned int rd :1;
-	unsigned int tc :1;
-	unsigned int aa :1;
-	unsigned int opcode :4;
-	unsigned int qr :1;
-
-	unsigned int rcode :4;
-	unsigned int z : 1;
-	unsigned int ad : 1;
-	unsigned int cd : 1;
-	unsigned int ra :1;
-#endif
-
-	unsigned int qdcount :16;
-	unsigned int ancount :16;
-	unsigned int nscount :16;
-	unsigned int arcount :16;
-};
-
 enum dns_section {
 	DNS_S_QD = 0x01,
 #define DNS_S_QUESTION          DNS_S_QD
diff --git a/contrib/librdns/resolver.c b/contrib/librdns/resolver.c
index d7d41915f..6ef6d4558 100644
--- a/contrib/librdns/resolver.c
+++ b/contrib/librdns/resolver.c
@@ -286,10 +286,43 @@ rdns_parse_reply (uint8_t *in, int r, struct rdns_request *req,
 	return true;
 }
 
-void
-rdns_process_read (int fd, void *arg)
+static void
+rdns_process_tcp_read (int fd, struct rdns_io_channel *ioc)
+{
+
+}
+
+static void
+rdns_process_tcp_connect (int fd, struct rdns_io_channel *ioc)
+{
+	struct rdns_resolver *resolver = ioc->resolver;
+	struct rdns_server *serv = ioc->srv;
+	int r = connect (ioc->sock, ioc->saddr, ioc->slen);
+
+	if (r == -1) {
+		if (errno != EAGAIN && errno != EINTR && errno != EINPROGRESS) {
+			rdns_err ("cannot connect a TCP socket: %s for server %s",
+					strerror(errno), serv->name);
+			resolver->async->del_write (resolver->async->data, ioc->async_io);
+		}
+		else {
+			/* We need to wait again for write readiness here */
+			ioc->async_io = resolver->async->add_write (resolver->async->data,
+					ioc->sock, ioc);
+		}
+	}
+	else {
+		/* Always be ready to read from a TCP socket */
+		resolver->async->del_write (resolver->async->data, ioc->async_io);
+		ioc->flags |= RDNS_CHANNEL_CONNECTED|RDNS_CHANNEL_ACTIVE;
+		ioc->tcp->async_read = resolver->async->add_read(resolver->async->data,
+				ioc->sock, ioc);
+	}
+}
+
+static void
+rdns_process_udp_read (int fd, struct rdns_io_channel *ioc)
 {
-	struct rdns_io_channel *ioc = arg;
 	struct rdns_resolver *resolver;
 	struct rdns_request *req = NULL;
 	ssize_t r;
@@ -310,7 +343,7 @@ rdns_process_read (int fd, void *arg)
 				sizeof (in), resolver->curve_plugin->data, &req,
 				ioc->saddr, ioc->slen);
 		if (req == NULL &&
-				r > (int)(sizeof (struct dns_header) + sizeof (struct dns_query))) {
+			r > (int)(sizeof (struct dns_header) + sizeof (struct dns_query))) {
 			req = rdns_find_dns_request (in, ioc);
 		}
 	}
@@ -336,6 +369,25 @@ rdns_process_read (int fd, void *arg)
 	}
 }
 
+void
+rdns_process_read (int fd, void *arg)
+{
+	struct rdns_io_channel *ioc = (struct rdns_io_channel *)arg;
+
+
+	if (IS_CHANNEL_TCP(ioc)) {
+		if (IS_CHANNEL_CONNECTED(ioc)) {
+			rdns_process_tcp_read (fd, ioc);
+		}
+		else {
+			rdns_process_tcp_connect (fd, ioc);
+		}
+	}
+	else {
+		rdns_process_udp_read (fd, ioc);
+	}
+}
+
 void
 rdns_process_timer (void *arg)
 {
diff --git a/contrib/librdns/util.c b/contrib/librdns/util.c
index 3cdb88fd0..9b74b7466 100644
--- a/contrib/librdns/util.c
+++ b/contrib/librdns/util.c
@@ -515,7 +515,16 @@ rdns_ioc_new (struct rdns_server *serv,
 			  struct rdns_resolver *resolver,
 			  bool is_tcp)
 {
-	struct rdns_io_channel *nioc = calloc (1, sizeof (struct rdns_io_channel));
+	struct rdns_io_channel *nioc;
+
+	if (is_tcp) {
+		nioc = calloc (1, sizeof (struct rdns_io_channel)
+				+ sizeof (struct rdns_tcp_channel));
+	}
+	else {
+		nioc = calloc (1, sizeof (struct rdns_io_channel));
+	}
+
 	if (nioc == NULL) {
 		rdns_err ("calloc fails to allocate rdns_io_channel");
 		return NULL;
@@ -531,7 +540,8 @@ rdns_ioc_new (struct rdns_server *serv,
 	}
 
 	if (is_tcp) {
-		/* We also need to connect a TCP channel */
+		/* We also need to connect a TCP channel and set a TCP buffer */
+		nioc->tcp = (struct rdns_tcp_channel *)(((unsigned char *)nioc) + sizeof(*nioc));
 		int r = connect (nioc->sock, nioc->saddr, nioc->slen);
 
 		if (r == -1) {
@@ -550,7 +560,10 @@ rdns_ioc_new (struct rdns_server *serv,
 			}
 		}
 		else {
+			/* Always be ready to read from a TCP socket */
 			nioc->flags |= RDNS_CHANNEL_CONNECTED|RDNS_CHANNEL_ACTIVE;
+			nioc->tcp->async_read = resolver->async->add_read(resolver->async->data,
+					nioc->sock, nioc);
 		}
 
 		nioc->flags |= RDNS_CHANNEL_TCP;
@@ -560,7 +573,7 @@ rdns_ioc_new (struct rdns_server *serv,
 	nioc->resolver = resolver;
 
 	/* If it is not NULL then we are in a delayed connection state */
-	if (nioc->async_io == NULL) {
+	if (!is_tcp) {
 		nioc->flags |= RDNS_CHANNEL_ACTIVE;
 		nioc->async_io = resolver->async->add_read(resolver->async->data,
 				nioc->sock, nioc);