commit 99151c4: [Project] Rdns: Initial support of TCP IO channels

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2022-01-01 17:49:27 +0000
URL: https://github.com/rspamd/rspamd/commit/99151c49061b808ffe954ebcb903efcc802086f0

[Project] Rdns: Initial support of TCP IO channels

---
 contrib/librdns/dns_private.h |  3 +++
 contrib/librdns/resolver.c    | 19 ++++++++++++++++++-
 contrib/librdns/util.c        | 37 ++++++++++++++++++++++++++++++++++---
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/contrib/librdns/dns_private.h b/contrib/librdns/dns_private.h
index f59fa2719..4429552bc 100644
--- a/contrib/librdns/dns_private.h
+++ b/contrib/librdns/dns_private.h
@@ -34,6 +34,7 @@
 
 static const int dns_port = 53;
 static const int default_io_cnt = 8;
+static const int default_tcp_io_cnt = 2;
 
 #define UDP_PACKET_SIZE (4096 * 2)
 
@@ -51,8 +52,10 @@ struct rdns_server {
 	char *name;
 	unsigned int port;
 	unsigned int io_cnt;
+	unsigned int tcp_io_cnt;
 
 	struct rdns_io_channel **io_channels;
+	struct rdns_io_channel **tcp_io_channels;
 	void *ups_elt;
 	upstream_entry_t up;
 };
diff --git a/contrib/librdns/resolver.c b/contrib/librdns/resolver.c
index 6b1ed8211..d7d41915f 100644
--- a/contrib/librdns/resolver.c
+++ b/contrib/librdns/resolver.c
@@ -893,6 +893,18 @@ rdns_resolver_init (struct rdns_resolver *resolver)
 
 			serv->io_channels[i] = ioc;
 		}
+
+		serv->tcp_io_channels = calloc (serv->tcp_io_cnt, sizeof (struct rdns_io_channel *));
+		for (i = 0; i < serv->tcp_io_cnt; i ++) {
+			ioc = rdns_ioc_new (serv, resolver, true);
+
+			if (ioc == NULL) {
+				rdns_err ("cannot allocate memory for the resolver IO channels");
+				return false;
+			}
+
+			serv->tcp_io_channels[i] = ioc;
+		}
 	}
 
 	if (resolver->async->add_periodic) {
@@ -952,6 +964,8 @@ rdns_resolver_add_server (struct rdns_resolver *resolver,
 	}
 
 	serv->io_cnt = io_cnt;
+	/* TODO: make it configurable maybe? */
+	serv->tcp_io_cnt = default_tcp_io_cnt;
 	serv->port = port;
 
 	UPSTREAM_ADD (resolver->servers, serv, priority);
@@ -1026,7 +1040,10 @@ rdns_resolver_free (struct rdns_resolver *resolver)
 				ioc = serv->io_channels[i];
 				REF_RELEASE (ioc);
 			}
-			serv->io_cnt = 0;
+			for (i = 0; i < serv->tcp_io_cnt; i ++) {
+				ioc = serv->tcp_io_channels[i];
+				REF_RELEASE (ioc);
+			}
 			UPSTREAM_DEL (resolver->servers, serv);
 			free (serv->io_channels);
 			free (serv->name);
diff --git a/contrib/librdns/util.c b/contrib/librdns/util.c
index d69ef6cd0..3cdb88fd0 100644
--- a/contrib/librdns/util.c
+++ b/contrib/librdns/util.c
@@ -530,11 +530,42 @@ rdns_ioc_new (struct rdns_server *serv,
 		return NULL;
 	}
 
+	if (is_tcp) {
+		/* We also need to connect a TCP channel */
+		int r = connect (nioc->sock, nioc->saddr, nioc->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);
+				close(nioc->sock);
+				free(nioc);
+
+				return NULL;
+			}
+			else {
+				/* We need to wait for write readiness here */
+				nioc->async_io = resolver->async->add_write (resolver->async->data,
+						nioc->sock, nioc);
+			}
+		}
+		else {
+			nioc->flags |= RDNS_CHANNEL_CONNECTED|RDNS_CHANNEL_ACTIVE;
+		}
+
+		nioc->flags |= RDNS_CHANNEL_TCP;
+	}
+
 	nioc->srv = serv;
-	nioc->flags = RDNS_CHANNEL_ACTIVE;
 	nioc->resolver = resolver;
-	nioc->async_io = resolver->async->add_read (resolver->async->data,
-			nioc->sock, nioc);
+
+	/* If it is not NULL then we are in a delayed connection state */
+	if (nioc->async_io == NULL) {
+		nioc->flags |= RDNS_CHANNEL_ACTIVE;
+		nioc->async_io = resolver->async->add_read(resolver->async->data,
+				nioc->sock, nioc);
+	}
+
 	nioc->requests = kh_init(rdns_requests_hash);
 	REF_INIT_RETAIN (nioc, rdns_ioc_free);