[Minor] Try harder to track DNS based chaining

diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index b0bde6bce2a92915df096bd4843c41f9478b9196..358bf7dca6a06169f1e7afc731bbb9a579175b46 100644 (file)
--- a/src/libserver/dns.c
+++ b/src/libserver/dns.c
@@ -60,6 +60,10 @@ rspamd_dns_fin_cb (gpointer arg)
 {
        struct rspamd_dns_request_ud *reqdata = (struct rspamd_dns_request_ud *)arg;
 
+       if (reqdata->item) {
+               rspamd_symbols_cache_set_cur_item (reqdata->task, reqdata->item);
+       }
+
        if (reqdata->reply) {
                reqdata->cb (reqdata->reply, reqdata->ud);
        }
@@ -164,6 +168,7 @@ make_dns_request (struct rspamd_dns_resolver *resolver,
                if (pool == NULL) {
                        g_free (reqdata);
                }
+
                return NULL;
        }
 

diff --git a/src/lua/lua_dns_resolver.c b/src/lua/lua_dns_resolver.c
index d00aaa6a2aaa37610ce618cbc3c7f2cff6a4dda2..c95d7bbbac4b6cc53d9817d12d350063131013e8 100644 (file)
--- a/src/lua/lua_dns_resolver.c
+++ b/src/lua/lua_dns_resolver.c
@@ -183,6 +183,11 @@ lua_dns_resolver_callback (struct rdns_reply *reply, gpointer arg)
 
        lua_pushboolean (L, reply->authenticated);
 
+       if (cd->item) {
+               /* We also need to restore the item in case there are some chains */
+               rspamd_symbols_cache_set_cur_item (cd->task, cd->item);
+       }
+
        if (lua_pcall (L, 6, 0, err_idx) != 0) {
                tb = lua_touserdata (L, -1);
 
@@ -438,6 +443,13 @@ lua_dns_resolver_resolve_common (lua_State *L,
                        }
                }
                else {
+                       /* Fail-safety as this function can, in theory, call
+                        * lua_dns_resolver_callback without switching to the event loop
+                        */
+                       if (cbdata->item) {
+                               rspamd_symcache_item_async_inc (task, cbdata->item);
+                       }
+
                        if (forced) {
                                ret = make_dns_request_task_forced (task,
                                                lua_dns_resolver_callback,
@@ -461,9 +473,14 @@ lua_dns_resolver_resolve_common (lua_State *L,
                                }
                                /* callback was set up */
                                lua_pushboolean (L, TRUE);
-                       } else {
+                       }
+                       else {
                                lua_pushnil (L);
                        }
+
+                       if (cbdata->item) {
+                               rspamd_symcache_item_async_dec_check (task, cbdata->item);
+                       }
                }
        }
        else {