From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sat, 1 Jan 2022 17:49:27 +0000 (+0000)
Subject: [Project] Rdns: Initial support of TCP IO channels
X-Git-Tag: 3.2~129^2~15
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=99151c49061b808ffe954ebcb903efcc802086f0;p=rspamd.git

[Project] Rdns: Initial support of TCP IO channels
---

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);