[Minor] Allow to set a resolver for task

diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index 2aaf670a4fb1a83dda3a824f981e7ffdccb1ab67..3157c565e1519a1a5116cc4488638b868de70502 100644 (file)
--- 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 e06db3a44bfe24c0efbb1f769cf2eaf722a0d2da..4234504ee3f16b4c7e879b43e9fc33432b6fc89a 100644 (file)
--- 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 8e739f725443d5dab715f82dd2c40f1afdaecb85..7bc4438cdeca5ff84cdfc335075fbf47fa8b4259 100644 (file)
--- 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 650ae5834675b0abb5e9a4753978c30e5e7a9bf3..f3760ce3e46c1b707d48d342d2e8af7eee54f20d 100644 (file)
--- 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