[Minor] Moved coroutine-related functionality of DNS resolver into a separated module

src/lua/CMakeLists.txt

@@ -26,6 +26,7 @@ SET(LUASRC			  ${CMAKE_CURRENT_SOURCE_DIR}/lua_common.c
 					  ${CMAKE_CURRENT_SOURCE_DIR}/lua_sqlite3.c
 					  ${CMAKE_CURRENT_SOURCE_DIR}/lua_cryptobox.c
 					  ${CMAKE_CURRENT_SOURCE_DIR}/lua_map.c
-					  ${CMAKE_CURRENT_SOURCE_DIR}/lua_thread_pool.c)
+					  ${CMAKE_CURRENT_SOURCE_DIR}/lua_thread_pool.c
+					  ${CMAKE_CURRENT_SOURCE_DIR}/lua_rspamd_dns.c)
 
 SET(RSPAMD_LUA ${LUASRC} PARENT_SCOPE)

src/lua/lua_common.c

@@ -740,6 +740,7 @@ rspamd_lua_init ()
 	luaopen_fann (L);
 	luaopen_sqlite3 (L);
 	luaopen_cryptobox (L);
+	luaopen_rspamd_dns (L);
 
 	luaL_newmetatable (L, "rspamd{ev_base}");
 	lua_pushstring (L, "class");

src/lua/lua_common.h

@@ -280,6 +280,7 @@ void luaopen_html (lua_State * L);
 void luaopen_fann (lua_State *L);
 void luaopen_sqlite3 (lua_State *L);
 void luaopen_cryptobox (lua_State *L);
+void luaopen_rspamd_dns (lua_State *L);
 
 void rspamd_lua_dostring (const gchar *line);
 

src/lua/lua_dns.c

@@ -53,6 +53,8 @@ LUA_FUNCTION_DEF (dns_resolver, resolve_mx);
 LUA_FUNCTION_DEF (dns_resolver, resolve_ns);
 LUA_FUNCTION_DEF (dns_resolver, resolve);
 
+void lua_push_dns_reply (lua_State *L, const struct rdns_reply *reply);
+
 static const struct luaL_reg dns_resolverlib_f[] = {
 	LUA_INTERFACE_DEF (dns_resolver, init),
 	{NULL, NULL}
@@ -78,7 +80,6 @@ lua_check_dns_resolver (lua_State * L)
 }
 
 struct lua_dns_cbdata {
-	struct thread_entry *thread;
 	struct rspamd_task *task;
 	struct rspamd_dns_resolver *resolver;
 	gint cbref;
@@ -137,31 +138,66 @@ static void
 lua_dns_callback (struct rdns_reply *reply, gpointer arg)
 {
 	struct lua_dns_cbdata *cd = arg;
-	gint i = 0, naddrs = 0;
 	struct rspamd_dns_resolver **presolver;
-	struct rdns_reply_entry *elt;
-	rspamd_inet_addr_t *addr;
 	lua_State *L;
 	struct lua_callback_state cbs;
 
-	if (cd->cbref != -1) {
-		lua_thread_pool_prepare_callback (cd->resolver->cfg->lua_thread_pool, &cbs);
-		L = cbs.L;
+	lua_thread_pool_prepare_callback (cd->resolver->cfg->lua_thread_pool, &cbs);
+	L = cbs.L;
 
-		lua_rawgeti (L, LUA_REGISTRYINDEX, cd->cbref);
+	lua_rawgeti (L, LUA_REGISTRYINDEX, cd->cbref);
 
-		presolver = lua_newuserdata (L, sizeof (gpointer));
-		rspamd_lua_setclass (L, "rspamd{resolver}", -1);
+	presolver = lua_newuserdata (L, sizeof (gpointer));
+	rspamd_lua_setclass (L, "rspamd{resolver}", -1);
 
-		*presolver = cd->resolver;
-		lua_pushstring (L, cd->to_resolve);
-	} else {
-		L = cd->thread->lua_state;
-	}
+	*presolver = cd->resolver;
+	lua_pushstring (L, cd->to_resolve);
+
+	lua_push_dns_reply (L, reply);
 
 	/*
-	 * XXX: rework to handle different request types
+	 * 1 - resolver
+	 * 2 - to_resolve
+	 * 3 - entries | nil
+	 * 4 - error | nil
+	 * 5 - user_str
+	 * 6 - reply->authenticated
 	 */
+	if (reply->code != RDNS_RC_NOERROR) {
+		lua_pushnil (L);
+		lua_pushstring (L, rdns_strerror (reply->code));
+	}
+	if (cd->user_str != NULL) {
+		lua_pushstring (L, cd->user_str);
+	}
+	else {
+		lua_pushnil (L);
+	}
+
+	lua_pushboolean (L, reply->authenticated);
+
+	if (lua_pcall (L, 6, 0, 0) != 0) {
+		msg_info ("call to dns callback failed: %s", lua_tostring (L, -1));
+		lua_pop (L, 1);
+	}
+
+	/* Unref function */
+	luaL_unref (L, LUA_REGISTRYINDEX, cd->cbref);
+
+	lua_thread_pool_restore_callback (&cbs);
+
+	if (cd->s) {
+		rspamd_session_watcher_pop (cd->s, cd->w);
+	}
+}
+
+void
+lua_push_dns_reply (lua_State *L, const struct rdns_reply *reply)
+{
+	gint i = 0, naddrs = 0;
+	struct rdns_reply_entry *elt;
+	rspamd_inet_addr_t *addr;
+
 	if (reply->code == RDNS_RC_NOERROR) {
 		LL_FOREACH (reply->entries, elt) {
 			naddrs ++;
@@ -234,70 +270,6 @@ lua_dns_callback (struct rdns_reply *reply, gpointer arg)
 		}
 		lua_pushnil (L);
 	}
-
-	if (cd->cbref != -1) {
-		/*
-		 * 1 - resolver
-		 * 2 - to_resolve
-		 * 3 - entries | nil
-		 * 4 - error | nil
-		 * 5 - user_str
-		 * 6 - reply->authenticated
-		 */
-		if (reply->code != RDNS_RC_NOERROR) {
-			lua_pushnil (L);
-			lua_pushstring (L, rdns_strerror (reply->code));
-		}
-		if (cd->user_str != NULL) {
-			lua_pushstring (L, cd->user_str);
-		}
-		else {
-			lua_pushnil (L);
-		}
-
-		lua_pushboolean (L, reply->authenticated);
-
-		if (lua_pcall (L, 6, 0, 0) != 0) {
-			msg_info ("call to dns callback failed: %s", lua_tostring (L, -1));
-			lua_pop (L, 1);
-		}
-
-		/* Unref function */
-		luaL_unref (L, LUA_REGISTRYINDEX, cd->cbref);
-
-		lua_thread_pool_restore_callback (&cbs);
-	} else {
-		/*
-		 * 1 - true | false in the case of error
-		 * 2. string - error message or table {
-		 * [0] -> entry 1
-		 * [1] -> entry 2
-		 * ...
-		 * is_authenticated = true|false
-		 * }
-		 */
-		if (reply->code != RDNS_RC_NOERROR) {
-			lua_pushboolean (L, false);
-			lua_pushstring (L, rdns_strerror (reply->code));
-		}
-		else {
-			lua_pushboolean (L, reply->authenticated);
-			lua_setfield (L, -3, "authenticated");
-
-			/* result 1 - not and error */
-			lua_pushboolean (L, true);
-			/* push table into stack, result 2 - results itself */
-			lua_pushvalue (L, -3);
-		}
-
-		g_assert (L == cd->thread->lua_state);
-
-		lua_resume_thread (cd->task, cd->thread, 2);
-	}
-
-	if (cd->s) {
-		rspamd_session_watcher_pop (cd->s, cd->w);
-	}
 }
 
 /***
@@ -357,7 +329,7 @@ lua_dns_resolver_resolve_common (lua_State *L,
 
 	/* Check arguments */
 	if (!rspamd_lua_parse_table_arguments (L, first, &err,
-			"session=U{session};mempool=U{mempool};*name=S;callback=F;"
+			"session=U{session};mempool=U{mempool};*name=S;*callback=F;"
 			"option=S;task=U{task};forced=B",
 			&session, &pool, &to_resolve, &cbref, &user_str, &task, &forced)) {
 
@@ -428,7 +400,6 @@ lua_dns_resolver_resolve_common (lua_State *L,
 			}
 		}
 		else {
-			cbdata->thread = lua_thread_pool_get_running_entry (task->cfg->lua_thread_pool);
 			cbdata->task = task;
 
 			if (forced) {
@@ -450,13 +421,8 @@ lua_dns_resolver_resolve_common (lua_State *L,
 				cbdata->s = session;
 				cbdata->w = rspamd_session_get_watcher (session);
 				rspamd_session_watcher_push (session);
-				if (cbdata->cbref != -1) {
-					/* callback was set up */
-					lua_pushboolean (L, TRUE);
-				} else {
-					/* this is coroutine-based call */
-					return lua_yield_thread (cbdata->thread, 0);
-				}
+				/* callback was set up */
+				lua_pushboolean (L, TRUE);
 			}
 			else {
 				lua_pushnil (L);

src/lua/lua_dns.h

@@ -0,0 +1,16 @@
+#ifndef RSPAMD_LUA_DNS_H
+#define RSPAMD_LUA_DNS_H
+
+typedef struct lua_State lua_State;
+struct rdns_reply;
+
+/**
+ * Pushes dns reply onto Lua stack
+ *
+ * @param L
+ * @param reply
+ */
+void
+lua_push_dns_reply (lua_State *L, const struct rdns_reply *reply);
+
+#endif //RSPAMD_LUA_DNS_H

src/lua/lua_rspamd_dns.c

@@ -0,0 +1,170 @@
+/*-
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "lua_common.h"
+#include "lua_dns.h"
+#include "lua_thread_pool.h"
+
+LUA_FUNCTION_DEF (dns, request);
+
+static const struct luaL_reg dns_f[] = {
+		LUA_INTERFACE_DEF (dns, request),
+		{"__tostring", rspamd_lua_class_tostring},
+		{NULL, NULL}
+};
+
+void
+lua_rspamd_dns_callback (struct rdns_reply *reply, void *arg);
+
+struct lua_rspamd_dns_cbdata {
+	struct thread_entry *thread;
+	struct rspamd_task *task;
+	struct rspamd_dns_resolver *resolver;
+	struct rspamd_async_watcher *w;
+	struct rspamd_async_session *s;
+};
+
+static gint
+lua_dns_request (lua_State *L)
+{
+	GError *err = NULL;
+	struct rspamd_async_session *session = NULL;
+	struct lua_rspamd_dns_cbdata *cbdata = NULL;
+	const gchar *to_resolve = NULL;
+	const gchar *type_str = NULL;
+	struct rspamd_task *task = NULL;
+	rspamd_mempool_t *pool = NULL;
+	gint ret = 0;
+	gboolean forced = FALSE;
+
+	/* Check arguments */
+	if (!rspamd_lua_parse_table_arguments (L, 1, &err,
+										   "*name=S;*task=U{task};*type=S;forced=B",
+										   &to_resolve, &task, &type_str, &forced)) {
+
+		if (err) {
+			ret = luaL_error (L, "invalid arguments: %s", err->message);
+			g_error_free (err);
+
+			return ret;
+		}
+
+		return luaL_error (L, "invalid arguments");
+	}
+
+	if (task) {
+		session = task->s;
+		pool = task->task_pool;
+	}
+	else {
+		return luaL_error (L, "invalid arguments: either task or session/config should be set");
+	}
+
+	enum rdns_request_type type = rdns_type_fromstr (type_str);
+
+	if (type == RDNS_REQUEST_INVALID) {
+		return luaL_error (L, "invalid arguments: this record type is not supported");
+	}
+
+	cbdata = rspamd_mempool_alloc0 (pool, sizeof (*cbdata));
+
+	cbdata->task = task;
+
+	if (type == RDNS_REQUEST_PTR) {
+		char *ptr_str;
+
+		ptr_str = rdns_generate_ptr_from_str (to_resolve);
+
+		if (ptr_str == NULL) {
+			msg_err_task_check ("wrong resolve string to PTR request: %s",
+								to_resolve);
+			lua_pushnil (L);
+
+			return 1;
+		}
+
+		to_resolve = rspamd_mempool_strdup (pool, ptr_str);
+		free (ptr_str);
+	}
+
+	if (forced) {
+		ret = make_dns_request_task_forced (task,
+											lua_rspamd_dns_callback,
+											cbdata,
+											type,
+											to_resolve);
+	}
+	else {
+		ret = make_dns_request_task (task,
+									 lua_rspamd_dns_callback,
+									 cbdata,
+									 type,
+									 to_resolve);
+	}
+
+	if (ret) {
+		cbdata->thread = lua_thread_pool_get_running_entry (task->cfg->lua_thread_pool);
+		cbdata->s = session;
+		cbdata->w = rspamd_session_get_watcher (session);
+		rspamd_session_watcher_push (session);
+		return lua_yield_thread (cbdata->thread, 0);
+	}
+	else {
+		lua_pushnil (L);
+		return 1;
+	}
+}
+
+void
+lua_rspamd_dns_callback (struct rdns_reply *reply, void *arg)
+{
+	struct lua_rspamd_dns_cbdata *cbdata = arg;
+	lua_State *L = cbdata->thread->lua_state;
+
+	if (reply->code != RDNS_RC_NOERROR) {
+		lua_pushboolean (L, false);
+		lua_pushstring (L, rdns_strerror (reply->code));
+	}
+	else {
+		lua_push_dns_reply (L, reply);
+
+		lua_pushboolean (L, reply->authenticated);
+		lua_setfield (L, -3, "authenticated");
+
+		/* result 1 - not and error */
+		lua_pushboolean (L, true);
+		/* push table into stack, result 2 - results itself */
+		lua_pushvalue (L, -3);
+	}
+
+	lua_resume_thread (cbdata->task, cbdata->thread, 2);
+
+	if (cbdata->s) {
+		rspamd_session_watcher_pop (cbdata->s, cbdata->w);
+	}
+}
+
+static gint
+lua_load_rspamd_dns (lua_State * L)
+{
+	lua_newtable (L);
+	luaL_register (L, NULL, dns_f);
+
+	return 1;
+}
+
+void
+luaopen_rspamd_dns (lua_State * L)
+{
+	rspamd_lua_add_preload (L, "rspamd_dns", lua_load_rspamd_dns);
+}

src/plugins/lua/multimap.lua

@@ -28,6 +28,7 @@ local regexp = require "rspamd_regexp"
 local rspamd_expression = require "rspamd_expression"
 local rspamd_ip = require "rspamd_ip"
 local lua_util = require "lua_util"
+local rspamd_dns = require "rspamd_dns"
 local redis_params
 local fun = require "fun"
 local N = 'multimap'
@@ -703,12 +704,13 @@ local function multimap_callback(task, rule)
         else
           local to_resolve = ip_to_rbl(ip, rule['map'])
 
-          local is_ok, results = task:get_resolver():resolve_a({
+          local is_ok, results = rspamd_dns.request({
+            type = "a",
             task = task,
             name = to_resolve,
           })
 
-          lua_util.debugm(N, rspamd_config, 'resolve_a() finished: results=%1, is_ok=%2, to_resolve=%3', results, is_ok, to_resolve)
+          lua_util.debugm(N, rspamd_config, 'resolve() finished: results=%1, is_ok=%2, to_resolve=%3', results, is_ok, to_resolve)
 
           if not is_ok and (results ~= 'requested record is not found' and results ~= 'no records with this name') then
             rspamd_logger.errx(task, 'error looking up %s: %s', to_resolve, results)

src/plugins/lua/reputation.lua

@@ -25,6 +25,7 @@ local N = 'reputation'
 
 local rspamd_logger = require "rspamd_logger"
 local rspamd_util = require "rspamd_util"
+local rspamd_dns = require "rspamd_dns"
 local lua_util = require "lua_util"
 local lua_maps = require "lua_maps"
 local hash = require 'rspamd_cryptobox_hash'
@@ -716,11 +717,12 @@ end
 --]]
 
 local function reputation_dns_get_token(task, rule, token, continuation_cb)
-  local r = task:get_resolver()
+  -- local r = task:get_resolver()
   local key = gen_token_key(token, rule)
   local dns_name = key .. '.' .. rule.backend.config.list
 
-  local is_ok, results = r:resolve_a({
+  local is_ok, results = rspamd_dns.request({
+    type = 'a',
     task = task,
     name = dns_name,
     forced = true,