[Minor] Try harder to track DNS based chaining

2 files changed, 23 insertions, 1 deletions

diff --git a/src/libserver/dns.c b/src/libserver/dns.c
index b0bde6bce..358bf7dca 100644
--- 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 d00aaa6a2..c95d7bbba 100644
--- 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 {