[Minor] DNS: Add servfail cache

2 files changed, 145 insertions, 8 deletions

diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index 08454c656..710f96b3a 100644
--- a/src/libserver/dns.c
+++ b/src/libserver/dns.c
@@ -57,6 +57,36 @@ struct rspamd_dns_request_ud {
 	struct rdns_reply *reply;
 };
 
+struct rspamd_dns_fail_cache_entry {
+	const char *name;
+	gint32 namelen;
+	enum rdns_request_type type;
+};
+
+static guint
+rspamd_dns_fail_hash (gconstpointer ptr)
+{
+	struct rspamd_dns_fail_cache_entry *elt =
+			(struct rspamd_dns_fail_cache_entry *)ptr;
+
+	/* We don't care about type when doing hashing */
+	return rspamd_cryptobox_fast_hash (elt->name, elt->namelen,
+			rspamd_hash_seed ());
+}
+
+static gboolean
+rspamd_dns_fail_equal (gconstpointer p1, gconstpointer p2)
+{
+	struct rspamd_dns_fail_cache_entry *e1 = (struct rspamd_dns_fail_cache_entry *)p1,
+			*e2 = (struct rspamd_dns_fail_cache_entry *)p2;
+
+	if (e1->type == e2->type && e1->namelen == e2->namelen) {
+		return memcmp (e1->name, e2->name, e1->namelen) == 0;
+	}
+
+	return FALSE;
+}
+
 static void
 rspamd_dns_fin_cb (gpointer arg)
 {
@@ -100,13 +130,41 @@ rspamd_dns_callback (struct rdns_reply *reply, gpointer ud)
 
 	reqdata->reply = reply;
 
+
 	if (reqdata->session) {
+		if (reply->code == RDNS_RC_SERVFAIL &&
+			reqdata->task &&
+			reqdata->task->resolver->fails_cache) {
+
+			/* Add to cache... */
+			const gchar *name = reqdata->req->requested_names[0].name;
+			gchar *target;
+			gsize namelen;
+			struct rspamd_dns_fail_cache_entry *nentry;
+
+			/* Allocate in a single entry to allow further free in a single call */
+			namelen = strlen (name);
+			nentry = g_malloc (sizeof (nentry) + namelen + 1);
+			target = ((gchar *)nentry) + sizeof (nentry);
+			rspamd_strlcpy (target, name, namelen + 1);
+			nentry->type = reqdata->req->requested_names[0].type;
+			nentry->name = target;
+			nentry->namelen = namelen;
+
+			/* Rdns request is retained there */
+			rspamd_lru_hash_insert (reqdata->task->resolver->fails_cache,
+					nentry, rdns_request_retain (reply->request),
+					reqdata->task->task_timestamp,
+					reqdata->task->resolver->fails_cache_time);
+		}
+
 		/*
 		 * Ref event to avoid double unref by
 		 * event removing
 		 */
 		rdns_request_retain (reply->request);
-		rspamd_session_remove_event (reqdata->session, rspamd_dns_fin_cb, reqdata);
+		rspamd_session_remove_event (reqdata->session,
+				rspamd_dns_fin_cb, reqdata);
 	}
 	else {
 		reqdata->cb (reply, reqdata->ud);
@@ -177,13 +235,38 @@ rspamd_dns_resolver_request (struct rspamd_dns_resolver *resolver,
 	return reqdata;
 }
 
+struct rspamd_dns_cached_delayed_cbdata {
+	struct rspamd_task *task;
+	dns_callback_type cb;
+	gpointer ud;
+	ev_timer tm;
+	struct rdns_request *req;
+};
+
+static void
+rspamd_fail_cache_cb (EV_P_ ev_timer *w, int revents)
+{
+	struct rspamd_dns_cached_delayed_cbdata *cbd =
+			(struct rspamd_dns_cached_delayed_cbdata *)w->data;
+	struct rdns_reply fake_reply;
+
+	ev_timer_stop (EV_A_ w);
+	memset (&fake_reply, 0, sizeof (fake_reply));
+	fake_reply.code = RDNS_RC_SERVFAIL;
+	fake_reply.request = cbd->req;
+	fake_reply.resolver = cbd->req->resolver;
+	fake_reply.requested_name = cbd->req->requested_names[0].name;
+	cbd->cb (&fake_reply, cbd->ud);
+	rdns_request_release (cbd->req);
+}
+
 static gboolean
 make_dns_request_task_common (struct rspamd_task *task,
-	dns_callback_type cb,
-	gpointer ud,
-	enum rdns_request_type type,
-	const char *name,
-	gboolean forced)
+							  dns_callback_type cb,
+							  gpointer ud,
+							  enum rdns_request_type type,
+							  const char *name,
+							  gboolean forced)
 {
 	struct rspamd_dns_request_ud *reqdata;
 
@@ -191,7 +274,37 @@ make_dns_request_task_common (struct rspamd_task *task,
 		return FALSE;
 	}
 
-	reqdata = rspamd_dns_resolver_request (task->resolver, task->s, task->task_pool, cb, ud,
+	if (task->resolver->fails_cache) {
+		/* Search in failures cache */
+		struct rspamd_dns_fail_cache_entry search;
+		struct rdns_request *req;
+
+		search.name = name;
+		search.namelen = strlen (name);
+		search.type = type;
+
+		if ((req = rspamd_lru_hash_lookup (task->resolver->fails_cache,
+				&search, task->task_timestamp)) != NULL) {
+			/*
+			 * We need to reply with SERVFAIL again to the API, so add a special
+			 * timer, uh-oh, and fire it
+			 */
+			struct rspamd_dns_cached_delayed_cbdata *cbd =
+					rspamd_mempool_alloc0 (task->task_pool, sizeof (*cbd));
+
+			ev_timer_init (&cbd->tm, rspamd_fail_cache_cb, 0.0, 0.0);
+			cbd->task = task;
+			cbd->cb = cb;
+			cbd->ud = ud;
+			cbd->req = rdns_request_retain (req);
+			cbd->tm.data = cbd;
+
+			return TRUE;
+		}
+	}
+
+	reqdata = rspamd_dns_resolver_request (
+			task->resolver, task->s, task->task_pool, cb, ud,
 			type, name);
 
 	if (reqdata) {
@@ -531,7 +644,8 @@ rspamd_dns_resolver_config_ucl (struct rspamd_config *cfg,
 								struct rspamd_dns_resolver *dns_resolver,
 								const ucl_object_t *dns_section)
 {
-	const ucl_object_t *fake_replies;
+	const ucl_object_t *fake_replies, *fails_cache_size, *fails_cache_time;
+	static const ev_tstamp default_fails_cache_time = 10.0;
 
 	/* Process fake replies */
 	fake_replies = ucl_object_lookup_any (dns_section, "fake_records",
@@ -544,6 +658,22 @@ rspamd_dns_resolver_config_ucl (struct rspamd_config *cfg,
 			rspamd_process_fake_reply (cfg, dns_resolver, cur_arr);
 		}
 	}
+
+	fails_cache_size = ucl_object_lookup (dns_section, "fails_cache_size");
+	if (fails_cache_size && ucl_object_type (fails_cache_size) == UCL_INT) {
+
+		dns_resolver->fails_cache_time = default_fails_cache_time;
+		fails_cache_time = ucl_object_lookup (dns_section, "fails_cache_time");
+
+		if (fails_cache_time) {
+			dns_resolver->fails_cache_time = ucl_object_todouble (fails_cache_time);
+		}
+
+		dns_resolver->fails_cache = rspamd_lru_hash_new_full (
+				ucl_object_toint (fails_cache_size),
+				g_free, (GDestroyNotify)rdns_request_release,
+				rspamd_dns_fail_hash, rspamd_dns_fail_equal);
+	}
 }
 
 struct rspamd_dns_resolver *
@@ -655,6 +785,10 @@ rspamd_dns_resolver_deinit (struct rspamd_dns_resolver *resolver)
 			rspamd_upstreams_destroy (resolver->ups);
 		}
 
+		if (resolver->fails_cache) {
+			rspamd_lru_hash_destroy (resolver->fails_cache);
+		}
+
 		g_free (resolver);
 	}
 }
diff --git a/src/libserver/dns.h b/src/libserver/dns.h
index 87cc0954b..367053ef4 100644
--- a/src/libserver/dns.h
+++ b/src/libserver/dns.h
@@ -23,6 +23,7 @@
 #include "logger.h"
 #include "rdns.h"
 #include "upstream.h"
+#include "libutil/hash.h"
 
 #ifdef  __cplusplus
 extern "C" {
@@ -34,6 +35,8 @@ struct rspamd_task;
 struct rspamd_dns_resolver {
 	struct rdns_resolver *r;
 	struct ev_loop *event_loop;
+	rspamd_lru_hash_t *fails_cache;
+	ev_tstamp fails_cache_time;
 	struct upstream_list *ups;
 	struct rspamd_config *cfg;
 	gdouble request_timeout;