From c3aec851c4e7b4e987246693aeed1333e8894da9 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Tue, 21 Jul 2015 13:33:49 +0100
Subject: [PATCH] Add constraints to dns requests ammount per task.

---
 src/libserver/cfg_file.h  |  1 +
 src/libserver/cfg_rcl.c   |  5 +++++
 src/libserver/cfg_utils.c |  1 +
 src/libserver/dns.c       | 26 ++++++++++++++++++++++++++
 src/libserver/dns.h       |  6 ++++++
 5 files changed, 39 insertions(+)

diff --git a/src/libserver/cfg_file.h b/src/libserver/cfg_file.h
index c6087b157..3ddd29586 100644
--- a/src/libserver/cfg_file.h
+++ b/src/libserver/cfg_file.h
@@ -254,6 +254,7 @@ struct rspamd_config {
 	guint32 dns_throttling_time;                    /**< time in seconds for DNS throttling					*/
 	guint32 dns_io_per_server;                      /**< number of sockets per DNS server					*/
 	GList *nameservers;                             /**< list of nameservers or NULL to parse resolv.conf	*/
+	guint32 dns_max_requests;                       /**< limit of DNS requests per task 					*/
 
 	guint upstream_max_errors;						/**< upstream max errors before shutting off			*/
 	gdouble upstream_error_time;					/**< rate of upstream errors							*/
diff --git a/src/libserver/cfg_rcl.c b/src/libserver/cfg_rcl.c
index 7c898431f..a56779f2e 100644
--- a/src/libserver/cfg_rcl.c
+++ b/src/libserver/cfg_rcl.c
@@ -1265,6 +1265,11 @@ rspamd_rcl_config_init (void)
 		rspamd_rcl_parse_struct_integer,
 		G_STRUCT_OFFSET (struct rspamd_config, dns_io_per_server),
 		RSPAMD_CL_FLAG_INT_32);
+	rspamd_rcl_add_default_handler (sub,
+		"dns_max_requests",
+		rspamd_rcl_parse_struct_integer,
+		G_STRUCT_OFFSET (struct rspamd_config, dns_max_requests),
+		RSPAMD_CL_FLAG_INT_32);
 
 	/* New DNS configuration */
 	ssub = rspamd_rcl_add_section (&sub->subsections, "dns", NULL,
diff --git a/src/libserver/cfg_utils.c b/src/libserver/cfg_utils.c
index da444cba2..2c95a5e0e 100644
--- a/src/libserver/cfg_utils.c
+++ b/src/libserver/cfg_utils.c
@@ -173,6 +173,7 @@ rspamd_config_defaults (struct rspamd_config *cfg)
 	cfg->log_extended = TRUE;
 
 	cfg->min_word_len = DEFAULT_MIN_WORD;
+	cfg->dns_max_requests = 64;
 }
 
 void
diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index 3e97160af..1821cb69a 100644
--- a/src/libserver/dns.c
+++ b/src/libserver/dns.c
@@ -117,6 +117,32 @@ make_dns_request (struct rspamd_dns_resolver *resolver,
 	return TRUE;
 }
 
+gboolean make_dns_request_task (struct rspamd_task *task,
+	dns_callback_type cb,
+	gpointer ud,
+	enum rdns_request_type type,
+	const char *name)
+{
+	gboolean ret;
+
+	if (task->dns_requests >= task->cfg->dns_max_requests) {
+		return FALSE;
+	}
+
+	ret = make_dns_request (task->resolver, task->s, task->task_pool, cb, ud,
+			type, name);
+
+	if (ret) {
+		task->dns_requests ++;
+
+		if (task->dns_requests >= task->cfg->dns_max_requests) {
+			msg_info ("<%s> stop resolving on reaching %ud requests",
+					task->message_id, task->dns_requests);
+		}
+	}
+
+	return ret;
+}
 
 struct rspamd_dns_resolver *
 dns_resolver_init (rspamd_logger_t *logger,
diff --git a/src/libserver/dns.h b/src/libserver/dns.h
index ffe80fae0..bc27f7d48 100644
--- a/src/libserver/dns.h
+++ b/src/libserver/dns.h
@@ -66,4 +66,10 @@ gboolean make_dns_request (struct rspamd_dns_resolver *resolver,
 	enum rdns_request_type type,
 	const char *name);
 
+gboolean make_dns_request_task (struct rspamd_task *task,
+	dns_callback_type cb,
+	gpointer ud,
+	enum rdns_request_type type,
+	const char *name);
+
 #endif
-- 
2.39.5