]> source.dussan.org Git - rspamd.git/commitdiff
Remove duplicated resolving functions from task.
authorVsevolod Stakhov <vsevolod@highsecure.ru>
Tue, 19 Nov 2013 17:56:43 +0000 (17:56 +0000)
committerVsevolod Stakhov <vsevolod@highsecure.ru>
Tue, 19 Nov 2013 17:56:43 +0000 (17:56 +0000)
src/lua/lua_dns.c
src/lua/lua_task.c
src/plugins/lua/multimap.lua
src/plugins/lua/ratelimit.lua
src/plugins/lua/rbl.lua

index fde9da42276297398406705b2ab3d2c05f770fe0..dedadeae5a84c7cfed1db0514f098de25d55b0a8 100644 (file)
@@ -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);
                }
index 80fc048f0ba78764a41bdb6c027d587e0c289ce6..3980220e2b439a2dc4ae40d721fabd934fdb1352 100644 (file)
@@ -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;
 }
 
index faed356c450f7ce496ee251b6d340888cd50cd42..60fc1fcd87275c3cf6dec00fefe4e7d06fa76b1f 100644 (file)
@@ -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
index 2bc2cac42aab72819a594582d618df055c3ade83..4e6a7e3533ff1c2cabfffd64c26fcfa24e771f19 100644 (file)
@@ -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
index 096c7f3c9775f1c1355cde9764249ec8f620210d..74c731e05518a5e16bb45b19bfd1ec0ce9e6f8cf 100644 (file)
@@ -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