From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sun, 2 Jan 2022 21:29:28 +0000 (+0000)
Subject: [Project] Rdns: Add more functions for TCP based requests
X-Git-Tag: 3.2~129^2~11
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=be5153fa1b8cc23b03e616fc0288b9f1994754d8;p=rspamd.git

[Project] Rdns: Add more functions for TCP based requests
---

diff --git a/contrib/librdns/dns_private.h b/contrib/librdns/dns_private.h
index 75fb5a3fa..38b6c8cc1 100644
--- a/contrib/librdns/dns_private.h
+++ b/contrib/librdns/dns_private.h
@@ -34,7 +34,7 @@
 
 static const int dns_port = 53;
 static const int default_io_cnt = 8;
-static const int default_tcp_io_cnt = 2;
+static const int default_tcp_io_cnt = 1;
 
 #define UDP_PACKET_SIZE (4096 * 2)
 
@@ -174,10 +174,12 @@ struct rdns_tcp_channel {
 };
 
 KHASH_DECLARE(rdns_requests_hash, int, struct rdns_request *);
+#define RDNS_IO_CHANNEL_TAG UINT64_C(0xe190a5ba12f094c8)
 /**
  * IO channel for a specific DNS server
  */
 struct rdns_io_channel {
+	uint64_t struct_magic; /**< tag for this structure */
 	struct rdns_server *srv;
 	struct rdns_resolver *resolver;
 	struct sockaddr *saddr;
diff --git a/contrib/librdns/rdns.h b/contrib/librdns/rdns.h
index 266ccabc2..2e00951ea 100644
--- a/contrib/librdns/rdns.h
+++ b/contrib/librdns/rdns.h
@@ -480,7 +480,7 @@ bool rdns_format_dns_name (struct rdns_resolver *resolver,
 
 void rdns_process_read (int fd, void *arg);
 void rdns_process_timer (void *arg);
-void rdns_process_retransmit (int fd, void *arg);
+void rdns_process_write (int fd, void *arg);
 
 #ifdef  __cplusplus
 }
diff --git a/contrib/librdns/rdns_ev.h b/contrib/librdns/rdns_ev.h
index 3827bcde2..35f532ae3 100644
--- a/contrib/librdns/rdns_ev.h
+++ b/contrib/librdns/rdns_ev.h
@@ -88,7 +88,7 @@ rdns_libev_read_event (struct ev_loop *loop, ev_io *ev, int revents)
 static void
 rdns_libev_write_event (struct ev_loop *loop, ev_io *ev, int revents)
 {
-	rdns_process_retransmit (ev->fd, ev->data);
+	rdns_process_write(ev->fd, ev->data);
 }
 
 static void
diff --git a/contrib/librdns/rdns_event.h b/contrib/librdns/rdns_event.h
index b3fc64aed..027181a0b 100644
--- a/contrib/librdns/rdns_event.h
+++ b/contrib/librdns/rdns_event.h
@@ -85,7 +85,7 @@ rdns_libevent_read_event (int fd, short what, void *ud)
 static void
 rdns_libevent_write_event (int fd, short what, void *ud)
 {
-	rdns_process_retransmit (fd, ud);
+	rdns_process_write (fd, ud);
 }
 
 static void
diff --git a/contrib/librdns/resolver.c b/contrib/librdns/resolver.c
index c27d9afcc..520e85588 100644
--- a/contrib/librdns/resolver.c
+++ b/contrib/librdns/resolver.c
@@ -562,10 +562,9 @@ rdns_process_ioc_refresh (void *arg)
 	}
 }
 
-void
-rdns_process_retransmit (int fd, void *arg)
+static void
+rdns_process_udp_retransmit (int fd, struct rdns_request *req)
 {
-	struct rdns_request *req = (struct rdns_request *)arg;
 	struct rdns_resolver *resolver;
 	struct rdns_reply *rep;
 	int r;
@@ -589,7 +588,7 @@ rdns_process_retransmit (int fd, void *arg)
 	if (r == 0) {
 		/* Retransmit one more time */
 		req->async_event = req->async->add_write (req->async->data,
-						fd, req);
+				fd, req);
 		req->state = RDNS_REQUEST_WAIT_SEND;
 	}
 	else if (r == -1) {
@@ -608,11 +607,50 @@ rdns_process_retransmit (int fd, void *arg)
 	}
 	else {
 		req->async_event = req->async->add_timer (req->async->data,
-			req->timeout, req);
+				req->timeout, req);
 		req->state = RDNS_REQUEST_WAIT_REPLY;
 	}
 }
 
+static void
+rdns_process_tcp_write (int fd, struct rdns_io_channel *ioc)
+{
+
+}
+
+void
+rdns_process_write (int fd, void *arg)
+{
+	/*
+	 * We first need to dispatch *arg to understand what has caused the write
+	 * readiness event.
+	 * The one possibility is that it was a UDP retransmit request, so our
+	 * arg will be struct rdns_request *
+	 * Another possibility is that write event was triggered by some TCP related
+	 * stuff. In this case the only possibility is that our arg is struct rdns_io_channel *
+	 * To distinguish these two cases (due to flaws in the rdns architecture in the first
+	 * place) we compare the first 8 bytes with RDNS_IO_CHANNEL_TAG
+	 */
+	uint64_t tag;
+
+	memcpy (&tag, arg, sizeof(tag));
+
+	if (tag == RDNS_IO_CHANNEL_TAG) {
+		struct rdns_io_channel *ioc = (struct rdns_io_channel *) arg;
+
+		if (IS_CHANNEL_CONNECTED(ioc)) {
+			rdns_process_tcp_write(fd, ioc);
+		}
+		else {
+			rdns_process_tcp_connect(fd, ioc);
+		}
+	}
+	else {
+		struct rdns_request *req = (struct rdns_request *) arg;
+		rdns_process_udp_retransmit(fd, req);
+	}
+}
+
 struct rdns_server *
 rdns_select_request_upstream (struct rdns_resolver *resolver,
 							  struct rdns_request *req,
@@ -1116,6 +1154,7 @@ rdns_resolver_free (struct rdns_resolver *resolver)
 			}
 			UPSTREAM_DEL (resolver->servers, serv);
 			free (serv->io_channels);
+			free (serv->tcp_io_channels);
 			free (serv->name);
 			free (serv);
 		}
diff --git a/contrib/librdns/util.c b/contrib/librdns/util.c
index 2fc9ec6c3..e33ab709c 100644
--- a/contrib/librdns/util.c
+++ b/contrib/librdns/util.c
@@ -540,6 +540,7 @@ rdns_ioc_new (struct rdns_server *serv,
 		return NULL;
 	}
 
+	nioc->struct_magic = RDNS_IO_CHANNEL_TAG;
 	nioc->sock = rdns_make_client_socket (serv->name, serv->port,
 			is_tcp ? SOCK_STREAM : SOCK_DGRAM, &nioc->saddr, &nioc->slen);
 	if (nioc->sock == -1) {