From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Tue, 19 Nov 2013 17:56:43 +0000 (+0000)
Subject: Remove duplicated resolving functions from task.
X-Git-Tag: 0.6.0~43
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=545f0f3a02365cf60f228dee3de7d3232af3207a;p=rspamd.git

Remove duplicated resolving functions from task.
---

diff --git a/src/lua/lua_dns.c b/src/lua/lua_dns.c
index fde9da422..dedadeae5 100644
--- a/src/lua/lua_dns.c
+++ b/src/lua/lua_dns.c
@@ -60,6 +60,7 @@ struct lua_dns_cbdata {
 	struct rspamd_dns_resolver	   *resolver;
 	gint							cbref;
 	const gchar                    *to_resolve;
+	const gchar                    *user_str;
 };
 
 static void
@@ -137,7 +138,14 @@ lua_dns_callback (struct rspamd_dns_reply *reply, gpointer arg)
 		lua_pushstring (cd->L, dns_strerror (reply->code));
 	}
 
-	if (lua_pcall (cd->L, 4, 0, 0) != 0) {
+	if (cd->user_str != NULL) {
+		lua_pushstring (cd->L, cd->user_str);
+	}
+	else {
+		lua_pushnil (cd->L);
+	}
+
+	if (lua_pcall (cd->L, 5, 0, 0) != 0) {
 		msg_info ("call to dns_callback failed: %s", lua_tostring (cd->L, -1));
 	}
 
@@ -210,6 +218,14 @@ lua_dns_resolver_resolve_common (lua_State *L, struct rspamd_dns_resolver *resol
 		cbdata->to_resolve = memory_pool_strdup (pool, to_resolve);
 		lua_pushvalue (L, 5);
 		cbdata->cbref = luaL_ref (L, LUA_REGISTRYINDEX);
+
+		if (lua_gettop (L) > 5) {
+			cbdata->user_str = lua_tostring (L, 6);
+		}
+		else {
+			cbdata->user_str = NULL;
+		}
+
 		if (type == DNS_REQUEST_PTR) {
 			make_dns_request (resolver, session, pool, lua_dns_callback, cbdata, type, &ina);
 		}
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 80fc048f0..3980220e2 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -51,6 +51,7 @@ LUA_FUNCTION_DEF (task, process_message);
 LUA_FUNCTION_DEF (task, set_cfg);
 LUA_FUNCTION_DEF (task, destroy);
 LUA_FUNCTION_DEF (task, get_mempool);
+LUA_FUNCTION_DEF (task, get_session);
 LUA_FUNCTION_DEF (task, get_ev_base);
 LUA_FUNCTION_DEF (task, insert_result);
 LUA_FUNCTION_DEF (task, set_pre_result);
@@ -62,9 +63,8 @@ LUA_FUNCTION_DEF (task, get_raw_headers);
 LUA_FUNCTION_DEF (task, get_raw_header);
 LUA_FUNCTION_DEF (task, get_raw_header_strong);
 LUA_FUNCTION_DEF (task, get_received_headers);
-LUA_FUNCTION_DEF (task, resolve_dns_a);
-LUA_FUNCTION_DEF (task, resolve_dns_ptr);
-LUA_FUNCTION_DEF (task, resolve_dns_txt);
+LUA_FUNCTION_DEF (task, get_resolver);
+LUA_FUNCTION_DEF (task, inc_dns_req);
 LUA_FUNCTION_DEF (task, call_rspamd_function);
 LUA_FUNCTION_DEF (task, get_recipients);
 LUA_FUNCTION_DEF (task, get_from);
@@ -100,6 +100,7 @@ static const struct luaL_reg    tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, process_message),
 	LUA_INTERFACE_DEF (task, set_cfg),
 	LUA_INTERFACE_DEF (task, get_mempool),
+	LUA_INTERFACE_DEF (task, get_session),
 	LUA_INTERFACE_DEF (task, get_ev_base),
 	LUA_INTERFACE_DEF (task, insert_result),
 	LUA_INTERFACE_DEF (task, set_pre_result),
@@ -111,9 +112,8 @@ static const struct luaL_reg    tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, get_raw_header),
 	LUA_INTERFACE_DEF (task, get_raw_header_strong),
 	LUA_INTERFACE_DEF (task, get_received_headers),
-	LUA_INTERFACE_DEF (task, resolve_dns_a),
-	LUA_INTERFACE_DEF (task, resolve_dns_ptr),
-	LUA_INTERFACE_DEF (task, resolve_dns_txt),
+	LUA_INTERFACE_DEF (task, get_resolver),
+	LUA_INTERFACE_DEF (task, inc_dns_req),
 	LUA_INTERFACE_DEF (task, call_rspamd_function),
 	LUA_INTERFACE_DEF (task, get_recipients),
 	LUA_INTERFACE_DEF (task, get_from),
@@ -363,6 +363,23 @@ lua_task_get_mempool (lua_State * L)
 	return 1;
 }
 
+static int
+lua_task_get_session (lua_State * L)
+{
+	struct rspamd_async_session   **psession;
+	struct worker_task             *task = lua_check_task (L);
+
+	if (task != NULL) {
+		psession = lua_newuserdata (L, sizeof (void *));
+		lua_setclass (L, "rspamd{session}", -1);
+		*psession = task->s;
+	}
+	else {
+		lua_pushnil (L);
+	}
+	return 1;
+}
+
 static int
 lua_task_get_ev_base (lua_State * L)
 {
@@ -666,282 +683,33 @@ lua_task_get_received_headers (lua_State * L)
 	return 1;
 }
 
-struct lua_dns_callback_data {
-	lua_State                      *L;
-	struct worker_task             *task;
-	union {
-		const gchar                *cbname;
-		gint 						ref;
-	} callback;
-	gboolean						cb_is_ref;
-	const gchar                    *to_resolve;
-	gint                            cbtype;
-	union {
-		gpointer                    string;
-		gboolean                    boolean;
-		gdouble                     number;
-	}                               cbdata;
-};
-
-static void
-lua_dns_callback (struct rspamd_dns_reply *reply, gpointer arg)
+static gint
+lua_task_get_resolver (lua_State *L)
 {
-	struct lua_dns_callback_data   *cd = arg;
-	gint                            i = 0;
-	struct worker_task            **ptask;
-	union rspamd_reply_element     *elt;
-	GList                          *cur;
-
-	if (cd->cb_is_ref) {
-		lua_rawgeti (cd->L, LUA_REGISTRYINDEX, cd->callback.ref);
-	}
-	else {
-		lua_getglobal (cd->L, cd->callback.cbname);
-	}
-	ptask = lua_newuserdata (cd->L, sizeof (struct worker_task *));
-	lua_setclass (cd->L, "rspamd{task}", -1);
-
-	*ptask = cd->task;
-	lua_pushstring (cd->L, cd->to_resolve);
-
-	if (reply->code == DNS_RC_NOERROR) {
-		if (reply->type == DNS_REQUEST_A) {
-
-			lua_newtable (cd->L);
-			cur = reply->elements;
-			while (cur) {
-				elt = cur->data;
-				lua_ip_push (cd->L, AF_INET, &elt->a.addr);
-				lua_rawseti (cd->L, -2, ++i);
-				cur = g_list_next (cur);
-			}
-			lua_pushnil (cd->L);
-		}
-		if (reply->type == DNS_REQUEST_AAA) {
-
-			lua_newtable (cd->L);
-			cur = reply->elements;
-			while (cur) {
-				elt = cur->data;
-				lua_ip_push (cd->L, AF_INET6, &elt->aaa.addr);
-				lua_rawseti (cd->L, -2, ++i);
-				cur = g_list_next (cur);
-			}
-			lua_pushnil (cd->L);
-		}
-		else if (reply->type == DNS_REQUEST_PTR) {
-			lua_newtable (cd->L);
-			cur = reply->elements;
-			while (cur) {
-				elt = cur->data;
-				lua_pushstring (cd->L, elt->ptr.name);
-				lua_rawseti (cd->L, -2, ++i);
-				cur = g_list_next (cur);
-			}
-			lua_pushnil (cd->L);
-
-		}
-		else if (reply->type == DNS_REQUEST_TXT) {
-			lua_newtable (cd->L);
-			cur = reply->elements;
-			while (cur) {
-				elt = cur->data;
-				lua_pushstring (cd->L, elt->txt.data);
-				lua_rawseti (cd->L, -2, ++i);
-				cur = g_list_next (cur);
-			}
-			lua_pushnil (cd->L);
+	struct worker_task             *task = lua_check_task (L);
+	struct rspamd_dns_resolver    **presolver;
 
-		}
-		else {
-			lua_pushnil (cd->L);
-			lua_pushstring (cd->L, "Unknown reply type");
-		}
+	if (task != NULL && task->resolver != NULL) {
+		presolver = lua_newuserdata (L, sizeof (void *));
+		lua_setclass (L, "rspamd{resolver}", -1);
+		*presolver = task->resolver;
 	}
 	else {
-		lua_pushnil (cd->L);
-		lua_pushstring (cd->L, dns_strerror (reply->code));
-	}
-
-	switch (cd->cbtype) {
-	case LUA_TBOOLEAN:
-		lua_pushboolean (cd->L, cd->cbdata.boolean);
-		break;
-	case LUA_TNUMBER:
-		lua_pushnumber (cd->L, cd->cbdata.number);
-		break;
-	case LUA_TSTRING:
-		lua_pushstring (cd->L, cd->cbdata.string);
-		break;
-	default:
-		lua_pushnil (cd->L);
-		break;
-	}
-
-	if (lua_pcall (cd->L, 5, 0, 0) != 0) {
-		msg_info ("call to %s failed: %s", cd->cb_is_ref ? "local function" :
-				cd->callback.cbname, lua_tostring (cd->L, -1));
-	}
-
-	/* Unref function */
-	if (cd->cb_is_ref) {
-		luaL_unref (cd->L, LUA_REGISTRYINDEX, cd->callback.ref);
+		lua_pushnil (L);
 	}
-}
-
-static gint
-lua_task_resolve_dns_a (lua_State * L)
-{
-	struct worker_task             *task = lua_check_task (L);
-	struct lua_dns_callback_data   *cd;
-
-	if (task) {
-		cd = memory_pool_alloc (task->task_pool, sizeof (struct lua_dns_callback_data));
-		cd->task = task;
-		cd->L = L;
-		cd->to_resolve = memory_pool_strdup (task->task_pool, luaL_checkstring (L, 2));
-
-		/* Check what type we have */
-		if (lua_type (L, 3) == LUA_TSTRING) {
-			cd->cb_is_ref = FALSE;
-			cd->callback.cbname = memory_pool_strdup (task->task_pool, luaL_checkstring (L, 3));
-		}
-		else {
-			lua_pushvalue (L, 3);
-			cd->cb_is_ref = TRUE;
-			cd->callback.ref = luaL_ref (L, LUA_REGISTRYINDEX);
-		}
-
-		cd->cbtype = lua_type (L, 4);
-		if (cd->cbtype != LUA_TNONE && cd->cbtype != LUA_TNIL) {
-			switch (cd->cbtype) {
-			case LUA_TBOOLEAN:
-				cd->cbdata.boolean = lua_toboolean (L, 4);
-				break;
-			case LUA_TNUMBER:
-				cd->cbdata.number = lua_tonumber (L, 4);
-				break;
-			case LUA_TSTRING:
-				cd->cbdata.string = memory_pool_strdup (task->task_pool, lua_tostring (L, 4));
-				break;
-			default:
-				msg_warn ("cannot handle type %s as callback data, try using closures", lua_typename (L, cd->cbtype));
-				cd->cbtype = LUA_TNONE;
-				break;
-			}
-		}
 
-		if (!cd->to_resolve) {
-			msg_info ("invalid parameters passed to function");
-			return 0;
-		}
-		if (make_dns_request (task->resolver, task->s, task->task_pool, lua_dns_callback, (void *)cd, DNS_REQUEST_A, cd->to_resolve)) {
-			task->dns_requests ++;
-		}
-	}
-	return 0;
+	return 1;
 }
 
 static gint
-lua_task_resolve_dns_txt (lua_State * L)
+lua_task_inc_dns_req (lua_State *L)
 {
 	struct worker_task             *task = lua_check_task (L);
-	struct lua_dns_callback_data   *cd;
 
-	if (task) {
-		cd = memory_pool_alloc (task->task_pool, sizeof (struct lua_dns_callback_data));
-		cd->task = task;
-		cd->L = L;
-		cd->to_resolve = memory_pool_strdup (task->task_pool, luaL_checkstring (L, 2));
-		/* Check what type we have */
-		if (lua_type (L, 3) == LUA_TSTRING) {
-			cd->cb_is_ref = FALSE;
-			cd->callback.cbname = memory_pool_strdup (task->task_pool, luaL_checkstring (L, 3));
-		}
-		else {
-			lua_pushvalue (L, 3);
-			cd->cb_is_ref = TRUE;
-			cd->callback.ref = luaL_ref (L, LUA_REGISTRYINDEX);
-		}
-		cd->cbtype = lua_type (L, 4);
-		if (cd->cbtype != LUA_TNONE && cd->cbtype != LUA_TNIL) {
-			switch (cd->cbtype) {
-			case LUA_TBOOLEAN:
-				cd->cbdata.boolean = lua_toboolean (L, 4);
-				break;
-			case LUA_TNUMBER:
-				cd->cbdata.number = lua_tonumber (L, 4);
-				break;
-			case LUA_TSTRING:
-				cd->cbdata.string = memory_pool_strdup (task->task_pool, lua_tostring (L, 4));
-				break;
-			default:
-				msg_warn ("cannot handle type %s as callback data", lua_typename (L, cd->cbtype));
-				cd->cbtype = LUA_TNONE;
-				break;
-			}
-		}
-		if (!cd->to_resolve) {
-			msg_info ("invalid parameters passed to function");
-			return 0;
-		}
-		if (make_dns_request (task->resolver, task->s, task->task_pool, lua_dns_callback, (void *)cd, DNS_REQUEST_TXT, cd->to_resolve)) {
-			task->dns_requests ++;
-		}
+	if (task != NULL) {
+		task->dns_requests ++;
 	}
-	return 0;
-}
 
-static gint
-lua_task_resolve_dns_ptr (lua_State * L)
-{
-	struct worker_task             *task = lua_check_task (L);
-	struct lua_dns_callback_data   *cd;
-	struct in_addr                 *ina;
-
-	if (task) {
-		cd = memory_pool_alloc (task->task_pool, sizeof (struct lua_dns_callback_data));
-		cd->task = task;
-		cd->L = L;
-		cd->to_resolve = memory_pool_strdup (task->task_pool, luaL_checkstring (L, 2));
-		/* Check what type we have */
-		if (lua_type (L, 3) == LUA_TSTRING) {
-			cd->cb_is_ref = FALSE;
-			cd->callback.cbname = memory_pool_strdup (task->task_pool, luaL_checkstring (L, 3));
-		}
-		else {
-			lua_pushvalue (L, 3);
-			cd->cb_is_ref = TRUE;
-			cd->callback.ref = luaL_ref (L, LUA_REGISTRYINDEX);
-		}
-		cd->cbtype = lua_type (L, 4);
-		if (cd->cbtype != LUA_TNONE && cd->cbtype != LUA_TNIL) {
-			switch (cd->cbtype) {
-			case LUA_TBOOLEAN:
-				cd->cbdata.boolean = lua_toboolean (L, 4);
-				break;
-			case LUA_TNUMBER:
-				cd->cbdata.number = lua_tonumber (L, 4);
-				break;
-			case LUA_TSTRING:
-				cd->cbdata.string = memory_pool_strdup (task->task_pool, lua_tostring (L, 4));
-				break;
-			default:
-				msg_warn ("cannot handle type %s as callback data", lua_typename (L, cd->cbtype));
-				cd->cbtype = LUA_TNONE;
-				break;
-			}
-		}
-		ina = memory_pool_alloc (task->task_pool, sizeof (struct in_addr));
-		if (!cd->to_resolve || !inet_aton (cd->to_resolve, ina)) {
-			msg_info ("invalid parameters passed to function");
-			return 0;
-		}
-		if (make_dns_request (task->resolver, task->s, task->task_pool,
-				lua_dns_callback, (void *)cd, DNS_REQUEST_PTR, ina)) {
-			task->dns_requests ++;
-		}
-	}
 	return 0;
 }
 
diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index faed356c4..60fc1fcd8 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -38,23 +38,21 @@ local function split(str, delim, maxNb)
 	return result
 end
 
-local function is_rbl(str,tail)
-	return tail == '' or string.sub(str,-string.len(tail))==tail
-end
 
-local function multimap_rbl_cb(task, to_resolve, results, err)
-	if results then
-		-- Get corresponding rule by rbl name
-		for _,rule in pairs(rules) do
-			if is_rbl(to_resolve, rule['map']) then
-				task:insert_result(rule['symbol'], 1, rule['map'])
-				return
+local function check_multimap(task)
+	local function multimap_rbl_cb(resolver, to_resolve, results, err, rbl)
+		task:inc_dns_req()
+		if results then
+			-- Get corresponding rule by rbl name
+			for _,rule in pairs(rules) do
+				if rule == rbl then
+					task:insert_result(rule['symbol'], 1, rule['map'])
+					return
+				end
 			end
 		end
 	end
-end
 
-local function check_multimap(task)
 	for _,rule in pairs(rules) do
 		if rule['type'] == 'ip' then
 			if rule['cdb'] then
@@ -102,7 +100,8 @@ local function check_multimap(task)
 		elseif rule['type'] == 'dnsbl' then
 			local ip = task:get_from_ip()
 			if ip then
-				task:resolve_dns_a(ip_to_rbl(ip, rule['map']), multimap_rbl_cb)
+				task:get_resolver():resolve_a(task:get_session(), task:get_mempool(),
+					ip_to_rbl(ip, rule['map']), multimap_rbl_cb, rule['map'])
 			end
 		elseif rule['type'] == 'rcpt' then
 			-- First try to get rcpt field
diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua
index 2bc2cac42..4e6a7e353 100644
--- a/src/plugins/lua/ratelimit.lua
+++ b/src/plugins/lua/ratelimit.lua
@@ -118,11 +118,11 @@ end
 --- Make rate key
 local function make_rate_key(from, to, ip)
 	if from and ip then
-		return string.format('%s:%s:%s', from, to, ip)
+		return string.format('%s:%s:%s', from, to, ip:to_string())
 	elseif from then
 		return string.format('%s:%s', from, to)
 	elseif ip then
-		return string.format('%s:%s', to, ip)
+		return string.format('%s:%s', to, ip:to_string())
 	elseif to then
 		return to
 	else
diff --git a/src/plugins/lua/rbl.lua b/src/plugins/lua/rbl.lua
index 096c7f3c9..74c731e05 100644
--- a/src/plugins/lua/rbl.lua
+++ b/src/plugins/lua/rbl.lua
@@ -27,27 +27,30 @@ local function ip_to_rbl(ip, rbl)
 	return str
 end
 
-local function rbl_dns_cb(task, to_resolve, results, err, sym)
-	if results then
-		task:insert_result(sym, 1)
+local function rbl_cb (task)
+	local function rbl_dns_cb(resolver, to_resolve, results, err, sym)
+		if results then
+			task:insert_result(sym, 1)
+		end
+		task:inc_dns_req()
 	end
-end
 
-local function rbl_cb (task)
 	local rip = task:get_from_ip()
 	if(rip ~= nil) then
 		for _,rbl in pairs(rbls) do
-			task:resolve_dns_a(ip_to_rbl(rip, rbl['rbl']), rbl_dns_cb, rbl['symbol'])
+			task:get_resolver():resolve_a(task:get_session(), task:get_mempool(), 
+				ip_to_rbl(rip, rbl['rbl']), rbl_dns_cb, rbl['symbol'])
 		end
 	end
 	local recvh = task:get_received_headers()
 	for _,rh in ipairs(recvh) do
 		if rh['real_ip'] then
 			for _,rbl in pairs(rbls) do
-				task:resolve_dns_a(ip_to_rbl(rip, rbl['rbl']), rbl_dns_cb, rbl['symbol'])
+				task:get_resolver():resolve_a(task:get_session(), task:get_mempool(), 
+					ip_to_rbl(rip, rbl['rbl']), rbl_dns_cb, rbl['symbol'])
 			end
-       		end
     	end
+	end
 end
 
 -- Registration