From 251a7c9d11a676a58f22edb8afd09f20c69bd61f Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Sat, 20 Oct 2018 17:23:55 +0100
Subject: [PATCH] [Project] Add items logic in DNS resolver

---
 src/libserver/dns.c | 36 ++++++++++++++++++++++++++----------
 src/libserver/dns.h |  3 ++-
 2 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index 4f8d70648..826b4ba1e 100644
--- a/src/libserver/dns.c
+++ b/src/libserver/dns.c
@@ -49,6 +49,8 @@ struct rspamd_dns_request_ud {
 	dns_callback_type cb;
 	gpointer ud;
 	rspamd_mempool_t *pool;
+	struct rspamd_task *task;
+	struct rspamd_symcache_item *item;
 	struct rdns_request *req;
 	struct rdns_reply *reply;
 };
@@ -93,6 +95,12 @@ rspamd_dns_callback (struct rdns_reply *reply, gpointer ud)
 		 * event removing
 		 */
 		rdns_request_retain (reply->request);
+
+		if (reqdata->item) {
+			rspamd_symcache_item_async_dec_check (reqdata->task,
+					reqdata->item);
+		}
+
 		rspamd_session_remove_event (reqdata->session, rspamd_dns_fin_cb, reqdata);
 	}
 	else {
@@ -104,7 +112,7 @@ rspamd_dns_callback (struct rdns_reply *reply, gpointer ud)
 	}
 }
 
-gboolean
+struct rspamd_dns_request_ud *
 make_dns_request (struct rspamd_dns_resolver *resolver,
 	struct rspamd_async_session *session,
 	rspamd_mempool_t *pool,
@@ -119,11 +127,11 @@ make_dns_request (struct rspamd_dns_resolver *resolver,
 	g_assert (resolver != NULL);
 
 	if (resolver->r == NULL) {
-		return FALSE;
+		return NULL;
 	}
 
 	if (session && rspamd_session_blocked (session)) {
-		return FALSE;
+		return NULL;
 	}
 
 	if (pool != NULL) {
@@ -146,7 +154,9 @@ make_dns_request (struct rspamd_dns_resolver *resolver,
 
 	if (session) {
 		if (req != NULL) {
-			rspamd_session_add_event (session, NULL, (event_finalizer_t) rspamd_dns_fin_cb, reqdata,
+			rspamd_session_add_event (session,
+					(event_finalizer_t) rspamd_dns_fin_cb,
+					reqdata,
 					g_quark_from_static_string ("dns resolver"));
 		}
 	}
@@ -155,10 +165,10 @@ make_dns_request (struct rspamd_dns_resolver *resolver,
 		if (pool == NULL) {
 			g_free (reqdata);
 		}
-		return FALSE;
+		return NULL;
 	}
 
-	return TRUE;
+	return reqdata;
 }
 
 static gboolean
@@ -169,25 +179,31 @@ make_dns_request_task_common (struct rspamd_task *task,
 	const char *name,
 	gboolean forced)
 {
-	gboolean ret;
+	struct rspamd_dns_request_ud *reqdata;
 
 	if (!forced && task->dns_requests >= task->cfg->dns_max_requests) {
 		return FALSE;
 	}
 
-	ret = make_dns_request (task->resolver, task->s, task->task_pool, cb, ud,
+	reqdata = make_dns_request (task->resolver, task->s, task->task_pool, cb, ud,
 			type, name);
 
-	if (ret) {
+	if (reqdata) {
 		task->dns_requests ++;
 
+		reqdata->task = task;
+		reqdata->item = rspamd_symbols_cache_get_cur_item (task);
+		rspamd_symcache_item_async_inc (task, reqdata->item);
+
 		if (!forced && task->dns_requests >= task->cfg->dns_max_requests) {
 			msg_info_task ("<%s> stop resolving on reaching %ud requests",
 					task->message_id, task->dns_requests);
 		}
+
+		return FALSE;
 	}
 
-	return ret;
+	return TRUE;
 }
 
 gboolean
diff --git a/src/libserver/dns.h b/src/libserver/dns.h
index 52325af91..c48e4b09c 100644
--- a/src/libserver/dns.h
+++ b/src/libserver/dns.h
@@ -43,6 +43,7 @@ struct rspamd_dns_resolver {
 struct rspamd_dns_resolver * dns_resolver_init (rspamd_logger_t *logger,
 	struct event_base *ev_base, struct rspamd_config *cfg);
 
+struct rspamd_dns_request_ud;
 /**
  * Make a DNS request
  * @param resolver resolver object
@@ -54,7 +55,7 @@ struct rspamd_dns_resolver * dns_resolver_init (rspamd_logger_t *logger,
  * @param ... string or ip address based on a request type
  * @return TRUE if request was sent.
  */
-gboolean make_dns_request (struct rspamd_dns_resolver *resolver,
+struct rspamd_dns_request_ud * make_dns_request (struct rspamd_dns_resolver *resolver,
 	struct rspamd_async_session *session,
 	rspamd_mempool_t *pool,
 	dns_callback_type cb,
-- 
2.39.5