From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Thu, 11 Apr 2019 14:06:01 +0000 (+0100)
Subject: [Minor] Allow to set a resolver for task
X-Git-Tag: 1.9.2~34
X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=5080c1565010de3a198507bc35dadd0cc3aad210;p=rspamd.git

[Minor] Allow to set a resolver for task
---

diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index 2aaf670a4..3157c565e 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -314,6 +314,7 @@ struct memory_pool_s * rspamd_lua_check_mempool (lua_State * L, gint pos);
 struct rspamd_config * lua_check_config (lua_State * L, gint pos);
 struct rspamd_async_session* lua_check_session (lua_State * L, gint pos);
 struct event_base* lua_check_ev_base (lua_State * L, gint pos);
+struct rspamd_dns_resolver * lua_check_dns_resolver (lua_State * L, gint pos);
 
 /**
  * Extract an arguments from lua table according to format string. Supported arguments are:
diff --git a/src/lua/lua_dns_resolver.c b/src/lua/lua_dns_resolver.c
index e06db3a44..4234504ee 100644
--- a/src/lua/lua_dns_resolver.c
+++ b/src/lua/lua_dns_resolver.c
@@ -43,8 +43,6 @@ end
 
 static const gchar *M = "rspamd lua dns resolver";
 
-struct rspamd_dns_resolver * lua_check_dns_resolver (lua_State * L);
-
 /* Lua bindings */
 LUA_FUNCTION_DEF (dns_resolver, init);
 LUA_FUNCTION_DEF (dns_resolver, resolve_a);
@@ -73,10 +71,10 @@ static const struct luaL_reg dns_resolverlib_m[] = {
 };
 
 struct rspamd_dns_resolver *
-lua_check_dns_resolver (lua_State * L)
+lua_check_dns_resolver (lua_State * L, gint pos)
 {
-	void *ud = rspamd_lua_check_udata (L, 1, "rspamd{resolver}");
-	luaL_argcheck (L, ud != NULL, 1, "'resolver' expected");
+	void *ud = rspamd_lua_check_udata (L, pos, "rspamd{resolver}");
+	luaL_argcheck (L, ud != NULL, pos, "'resolver' expected");
 	return ud ? *((struct rspamd_dns_resolver **)ud) : NULL;
 }
 
@@ -532,7 +530,7 @@ err:
 static int
 lua_dns_resolver_resolve_a (lua_State *L)
 {
-	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L);
+	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L, 1);
 
 	if (dns_resolver) {
 		return lua_dns_resolver_resolve_common (L,
@@ -562,7 +560,7 @@ lua_dns_resolver_resolve_a (lua_State *L)
 static int
 lua_dns_resolver_resolve_ptr (lua_State *L)
 {
-	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L);
+	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L, 1);
 
 	if (dns_resolver) {
 		return lua_dns_resolver_resolve_common (L,
@@ -592,7 +590,7 @@ lua_dns_resolver_resolve_ptr (lua_State *L)
 static int
 lua_dns_resolver_resolve_txt (lua_State *L)
 {
-	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L);
+	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L, 1);
 
 	if (dns_resolver) {
 		return lua_dns_resolver_resolve_common (L,
@@ -622,7 +620,7 @@ lua_dns_resolver_resolve_txt (lua_State *L)
 static int
 lua_dns_resolver_resolve_mx (lua_State *L)
 {
-	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L);
+	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L, 1);
 
 	if (dns_resolver) {
 		return lua_dns_resolver_resolve_common (L,
@@ -652,7 +650,7 @@ lua_dns_resolver_resolve_mx (lua_State *L)
 static int
 lua_dns_resolver_resolve_ns (lua_State *L)
 {
-	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L);
+	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L, 1);
 
 	if (dns_resolver) {
 		return lua_dns_resolver_resolve_common (L,
@@ -671,7 +669,7 @@ lua_dns_resolver_resolve_ns (lua_State *L)
 static int
 lua_dns_resolver_resolve (lua_State *L)
 {
-	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L);
+	struct rspamd_dns_resolver *dns_resolver = lua_check_dns_resolver (L, 1);
 	int type;
 
 	type = lua_dns_get_type (L, 2);
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 8e739f725..7bc4438cd 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -394,9 +394,14 @@ end
  */
 LUA_FUNCTION_DEF (task, get_resolver);
 /***
- * @method task:inc_dns_req()
- * Increment number of DNS requests for the task. Is used just for logging purposes.
+ * @method task:set_resolver(resolver)
+ * Sets rspamd_resolver for a specified task.
  */
+LUA_FUNCTION_DEF (task, set_resolver);
+/***
+* @method task:inc_dns_req()
+* Increment number of DNS requests for the task. Is used just for logging purposes.
+*/
 LUA_FUNCTION_DEF (task, inc_dns_req);
 /***
  * @method task:get_dns_req()
@@ -1058,6 +1063,7 @@ static const struct luaL_reg tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, get_queue_id),
 	LUA_INTERFACE_DEF (task, get_uid),
 	LUA_INTERFACE_DEF (task, get_resolver),
+	LUA_INTERFACE_DEF (task, set_resolver),
 	LUA_INTERFACE_DEF (task, inc_dns_req),
 	LUA_INTERFACE_DEF (task, get_dns_req),
 	LUA_INTERFACE_DEF (task, has_recipients),
@@ -2664,6 +2670,23 @@ lua_task_get_resolver (lua_State *L)
 	return 1;
 }
 
+static gint
+lua_task_set_resolver (lua_State *L)
+{
+	LUA_TRACE_POINT;
+	struct rspamd_task *task = lua_check_task (L, 1);
+	struct rspamd_dns_resolver *resolver = lua_check_dns_resolver (L, 2);
+
+	if (task != NULL && resolver != NULL) {
+		task->resolver = resolver;
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 0;
+}
+
 static gint
 lua_task_inc_dns_req (lua_State *L)
 {
diff --git a/src/rspamadm/rspamadm.c b/src/rspamadm/rspamadm.c
index 650ae5834..f3760ce3e 100644
--- a/src/rspamadm/rspamadm.c
+++ b/src/rspamadm/rspamadm.c
@@ -348,7 +348,7 @@ rspamadm_add_lua_globals (struct rspamd_dns_resolver *resolver)
 	presolver = lua_newuserdata (L, sizeof (struct rspamd_dns_resolver *));
 	rspamd_lua_setclass (L, "rspamd{resolver}", -1);
 	*presolver = resolver;
-	lua_setglobal (L, "rspamadm_resolver");
+	lua_setglobal (L, "rspamadm_dns_resolver");
 }
 
 gint