[Minor] Fix taskless resolver invocation

1 files changed, 91 insertions, 49 deletions

diff --git a/src/lua/lua_dns_resolver.c b/src/lua/lua_dns_resolver.c
index dc475f2ae..7c224243e 100644
--- a/src/lua/lua_dns_resolver.c
+++ b/src/lua/lua_dns_resolver.c
@@ -81,10 +81,11 @@ lua_check_dns_resolver (lua_State * L)
 
 struct lua_dns_cbdata {
 	struct rspamd_task *task;
+	rspamd_mempool_t *pool;
 	struct rspamd_dns_resolver *resolver;
 	gint cbref;
-	const gchar *to_resolve;
-	const gchar *user_str;
+	gchar *to_resolve;
+	gchar *user_str;
 	struct rspamd_async_watcher *w;
 	struct rspamd_async_session *s;
 };
@@ -141,10 +142,17 @@ lua_dns_resolver_callback (struct rdns_reply *reply, gpointer arg)
 	struct rspamd_dns_resolver **presolver;
 	lua_State *L;
 	struct lua_callback_state cbs;
+	rspamd_mempool_t *pool;
+	gint err_idx;
+	GString *tb = NULL;
 
+	pool = cd->pool;
 	lua_thread_pool_prepare_callback (cd->resolver->cfg->lua_thread_pool, &cbs);
 	L = cbs.L;
 
+	lua_pushcfunction (L, &rspamd_lua_traceback);
+	err_idx = lua_gettop (L);
+
 	lua_rawgeti (L, LUA_REGISTRYINDEX, cd->cbref);
 
 	presolver = lua_newuserdata (L, sizeof (gpointer));
@@ -176,19 +184,26 @@ lua_dns_resolver_callback (struct rdns_reply *reply, gpointer arg)
 
 	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);
+	if (lua_pcall (L, 6, 0, err_idx) != 0) {
+		msg_info_pool_check ("call to dns callback failed: %s", tb->str);
+		g_string_free (tb, TRUE);
 	}
 
+	lua_settop (L, err_idx - 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);
 	}
+
+	if (!cd->pool) {
+		g_free (cd->to_resolve);
+		g_free (cd->user_str);
+		g_free (cd);
+	}
 }
 
 void
@@ -347,45 +362,67 @@ lua_dns_resolver_resolve_common (lua_State *L,
 		pool = task->task_pool;
 		session = task->s;
 	}
-	else if (!session || !pool) {
-		return luaL_error (L, "invalid arguments: either 'task' or 'session'/'mempool' should be set");
-	}
 
-	if (pool != NULL && to_resolve != NULL) {
-		cbdata = rspamd_mempool_alloc0 (pool, sizeof (struct lua_dns_cbdata));
-		cbdata->resolver = resolver;
-		cbdata->cbref = cbref;
-		cbdata->user_str = rspamd_mempool_strdup (pool, user_str);
+	if (to_resolve != NULL) {
+		if (pool != NULL) {
+			cbdata = rspamd_mempool_alloc0 (pool, sizeof (struct lua_dns_cbdata));
+			cbdata->user_str = rspamd_mempool_strdup (pool, user_str);
 
-		if (type != RDNS_REQUEST_PTR) {
-			cbdata->to_resolve = rspamd_mempool_strdup (pool, to_resolve);
+			if (type != RDNS_REQUEST_PTR) {
+				cbdata->to_resolve = rspamd_mempool_strdup (pool, to_resolve);
+			}
+			else {
+				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);
+					goto err;
+				}
+
+				cbdata->to_resolve = rspamd_mempool_strdup (pool, ptr_str);
+				to_resolve = cbdata->to_resolve;
+				free (ptr_str);
+			}
 		}
 		else {
-			char *ptr_str;
+			cbdata = g_malloc0 (sizeof (struct lua_dns_cbdata));
+			cbdata->user_str = user_str ? g_strdup (user_str) : NULL;
 
-			ptr_str = rdns_generate_ptr_from_str (to_resolve);
+			if (type != RDNS_REQUEST_PTR) {
+				cbdata->to_resolve = g_strdup (to_resolve);
+			}
+			else {
+				char *ptr_str;
 
-			if (ptr_str == NULL) {
-				msg_err_task_check ("wrong resolve string to PTR request: %s",
-						to_resolve);
-				lua_pushnil (L);
+				ptr_str = rdns_generate_ptr_from_str (to_resolve);
 
-				return 1;
-			}
+				if (ptr_str == NULL) {
+					msg_err_task_check ("wrong resolve string to PTR request: %s",
+							to_resolve);
+					goto err;
+				}
 
-			cbdata->to_resolve = rspamd_mempool_strdup (pool, ptr_str);
-			to_resolve = cbdata->to_resolve;
-			free (ptr_str);
+				cbdata->to_resolve = g_strdup (ptr_str);
+				free (ptr_str);
+			}
 		}
 
+		cbdata->resolver = resolver;
+		cbdata->cbref = cbref;
+		cbdata->task = task;
+		cbdata->pool = pool;
+
 		if (task == NULL) {
-			if ( make_dns_request (resolver,
-								   session,
-								   pool,
-								   lua_dns_resolver_callback,
-								   cbdata,
-								   type,
-								   to_resolve)) {
+			if (make_dns_request (resolver,
+					session,
+					pool,
+					lua_dns_resolver_callback,
+					cbdata,
+					type,
+					to_resolve)) {
 
 				lua_pushboolean (L, TRUE);
 
@@ -396,25 +433,22 @@ lua_dns_resolver_resolve_common (lua_State *L,
 				}
 			}
 			else {
-				lua_pushnil (L);
+				goto err;
 			}
 		}
 		else {
-			cbdata->task = task;
-
 			if (forced) {
 				ret = make_dns_request_task_forced (task,
-													lua_dns_resolver_callback,
-													cbdata,
-													type,
-													to_resolve);
-			}
-			else {
+						lua_dns_resolver_callback,
+						cbdata,
+						type,
+						to_resolve);
+			} else {
 				ret = make_dns_request_task (task,
-											 lua_dns_resolver_callback,
-											 cbdata,
-											 type,
-											 to_resolve);
+						lua_dns_resolver_callback,
+						cbdata,
+						type,
+						to_resolve);
 			}
 
 			if (ret) {
@@ -423,8 +457,7 @@ lua_dns_resolver_resolve_common (lua_State *L,
 				rspamd_session_watcher_push (session);
 				/* callback was set up */
 				lua_pushboolean (L, TRUE);
-			}
-			else {
+			} else {
 				lua_pushnil (L);
 			}
 		}
@@ -434,7 +467,16 @@ lua_dns_resolver_resolve_common (lua_State *L,
 	}
 
 	return 1;
+err:
+	if (!pool) {
+		/* Free resources */
+		g_free (cbdata->to_resolve);
+		g_free (cbdata->user_str);
+	}
 
+	lua_pushnil (L);
+
+	return 1;
 }
 
 /***