[Minor] DNS: Add servfail cache

diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index 08454c6563b132e539a2de9959a9588ee3dcb811..710f96b3adbed0a0a075e794250b9e5ce7f936d2 100644 (file)
--- 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 87cc0954b8fd551b4df28e8f58dfd332480947d1..367053ef44e8af47902c2fd311eaf82aaf626193 100644 (file)
--- 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;