From 2e43b3e2be1d7d0cf677b965e4484742a581b292 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Tue, 4 Apr 2017 11:58:31 +0100 Subject: [PATCH] [Feature] Add task:headers_foreach function --- src/lua/lua_common.h | 16 ++-- src/lua/lua_mimepart.c | 2 +- src/lua/lua_task.c | 193 ++++++++++++++++++++++++++++++++--------- 3 files changed, 161 insertions(+), 50 deletions(-) diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h index ecf46d32e..141d614b8 100644 --- a/src/lua/lua_common.h +++ b/src/lua/lua_common.h @@ -87,6 +87,7 @@ struct rspamd_lua_regexp { struct rspamd_map; struct lua_map_callback_data; struct radix_tree_compressed; +struct rspamd_mime_header; enum rspamd_lua_map_type { RSPAMD_LUA_MAP_RADIX = 0, @@ -190,15 +191,18 @@ struct rspamd_lua_ip * lua_check_ip (lua_State * L, gint pos); struct rspamd_lua_text * lua_check_text (lua_State * L, gint pos); + +gint rspamd_lua_push_header (lua_State *L, + struct rspamd_mime_header *h, + gboolean full, + gboolean raw); /** * Push specific header to lua */ -gint rspamd_lua_push_header (lua_State * L, - GPtrArray *hdrs, - const gchar *name, - gboolean strong, - gboolean full, - gboolean raw); +gint rspamd_lua_push_header_array (lua_State *L, + GPtrArray *hdrs, + gboolean full, + gboolean raw); /** * Check for task at the specified position diff --git a/src/lua/lua_mimepart.c b/src/lua/lua_mimepart.c index 21640aab5..dba53a2d7 100644 --- a/src/lua/lua_mimepart.c +++ b/src/lua/lua_mimepart.c @@ -741,7 +741,7 @@ lua_mimepart_get_header_common (lua_State *L, gboolean full, gboolean raw) ar = rspamd_message_get_header_from_hash (part->raw_headers, NULL, name, FALSE); - return rspamd_lua_push_header (L, ar, name, FALSE, full, raw); + return rspamd_lua_push_header_array (L, ar, full, raw); } lua_pushnil (L); diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c index 53e4881d8..41bf8d792 100644 --- a/src/lua/lua_task.c +++ b/src/lua/lua_task.c @@ -773,6 +773,20 @@ LUA_FUNCTION_DEF (task, store_in_file); */ LUA_FUNCTION_DEF (task, get_protocol_reply); +/*** + * @method task:headers_foreach(callback, [params]) + * This method calls `callback` for each header that satisfies some condition. + * By default, all headers are iterated unless `callback` returns `true`. Nil or + * false means continue of iterations. + * Params could be as following: + * + * - `full`: header value is full table of all attributes @see task:get_header_full for details + * - `regexp`: return headers that satisfies the specified regexp + * @param {function} callback function from header name and header value + * @param {table} params optional parameters + */ +LUA_FUNCTION_DEF (task, headers_foreach); + static const struct luaL_reg tasklib_f[] = { {NULL, NULL} }; @@ -859,6 +873,7 @@ static const struct luaL_reg tasklib_m[] = { LUA_INTERFACE_DEF (task, get_digest), LUA_INTERFACE_DEF (task, store_in_file), LUA_INTERFACE_DEF (task, get_protocol_reply), + LUA_INTERFACE_DEF (task, headers_foreach), {"__tostring", rspamd_lua_class_tostring}, {NULL, NULL} }; @@ -1556,17 +1571,64 @@ lua_task_set_request_header (lua_State *L) } gint -rspamd_lua_push_header (lua_State * L, +rspamd_lua_push_header (lua_State *L, struct rspamd_mime_header *rh, + gboolean full, gboolean raw) +{ + const gchar *val; + + if (full) { + /* Create new associated table for a header */ + lua_createtable (L, 0, 7); + rspamd_lua_table_set (L, "name", rh->name); + + if (rh->value) { + rspamd_lua_table_set (L, "value", rh->value); + } + + if (rh->decoded) { + rspamd_lua_table_set (L, "decoded", rh->decoded); + } + + lua_pushstring (L, "tab_separated"); + lua_pushboolean (L, rh->tab_separated); + lua_settable (L, -3); + lua_pushstring (L, "empty_separator"); + lua_pushboolean (L, rh->empty_separator); + lua_settable (L, -3); + rspamd_lua_table_set (L, "separator", rh->separator); + lua_pushstring (L, "order"); + lua_pushnumber (L, rh->order); + lua_settable (L, -3); + } + else { + if (!raw) { + val = rh->decoded; + } + else { + val = rh->value; + } + + if (val) { + lua_pushstring (L, val); + } + else { + lua_pushnil (L); + } + } + + return 1; +} + +gint +rspamd_lua_push_header_array (lua_State * L, GPtrArray *ar, - const gchar *name, - gboolean strong, gboolean full, gboolean raw) { struct rspamd_mime_header *rh; guint i; - const gchar *val; + if (ar == NULL || ar->len == 0) { lua_pushnil (L); @@ -1579,46 +1641,11 @@ rspamd_lua_push_header (lua_State * L, PTR_ARRAY_FOREACH (ar, i, rh) { if (full) { - /* Create new associated table for a header */ - lua_createtable (L, 0, 7); - rspamd_lua_table_set (L, "name", rh->name); - - if (rh->value) { - rspamd_lua_table_set (L, "value", rh->value); - } - - if (rh->decoded) { - rspamd_lua_table_set (L, "decoded", rh->decoded); - } - - lua_pushstring (L, "tab_separated"); - lua_pushboolean (L, rh->tab_separated); - lua_settable (L, -3); - lua_pushstring (L, "empty_separator"); - lua_pushboolean (L, rh->empty_separator); - lua_settable (L, -3); - rspamd_lua_table_set (L, "separator", rh->separator); - lua_pushstring (L, "order"); - lua_pushnumber (L, rh->order); - lua_settable (L, -3); + rspamd_lua_push_header (L, rh, full, raw); lua_rawseti (L, -2, i + 1); } else { - if (!raw) { - val = rh->decoded; - } - else { - val = rh->value; - } - - if (val) { - lua_pushstring (L, val); - } - else { - lua_pushnil (L); - } - - return 1; + return rspamd_lua_push_header (L, rh, full, raw); } } @@ -1642,8 +1669,7 @@ lua_task_get_header_common (lua_State *L, gboolean full, gboolean raw) ar = rspamd_message_get_header_array (task, name, strong); - return rspamd_lua_push_header (L, ar, name, - strong, full, raw); + return rspamd_lua_push_header_array (L, ar, full, raw); } else { return luaL_error (L, "invalid arguments"); @@ -3972,6 +3998,87 @@ lua_task_get_protocol_reply (lua_State *L) return 1; } +static gint +lua_task_headers_foreach (lua_State *L) +{ + struct rspamd_task *task = lua_check_task (L, 1); + gboolean full = FALSE, raw = FALSE; + struct rspamd_lua_regexp *re = NULL; + GList *cur; + struct rspamd_mime_header *hdr; + + if (task && lua_isfunction (L, 2)) { + if (lua_istable (L, 3)) { + lua_pushstring (L, "full"); + lua_gettable (L, 3); + + if (lua_isboolean (L, -1)) { + full = lua_toboolean (L, -1); + } + + lua_pop (L, 1); + + lua_pushstring (L, "raw"); + lua_gettable (L, 3); + + if (lua_isboolean (L, -1)) { + raw = lua_toboolean (L, -1); + } + + lua_pop (L, 1); + + lua_pushstring (L, "regexp"); + lua_gettable (L, 3); + + if (lua_isuserdata (L, -1)) { + re = *(struct rspamd_lua_regexp **) + rspamd_lua_check_udata (L, -1, "rspamd{regexp}"); + } + + lua_pop (L, 1); + } + + if (task->headers_order) { + cur = task->headers_order->head; + + while (cur) { + hdr = cur->data; + + if (re && re->re) { + if (!rspamd_regexp_match (re->re, hdr->name, + strlen (hdr->name),FALSE)) { + cur = g_list_next (cur); + continue; + } + } + lua_pushvalue (L, 2); + lua_pushstring (L, hdr->name); + rspamd_lua_push_header (L, hdr, full, raw); + + if (lua_pcall (L, 3, 1, 0) != 0) { + msg_err_task ("call to header_foreach failed: %s", + lua_tostring (L, -1)); + lua_pop (L, 1); + break; + } + else { + if (lua_isboolean (L, -1)) { + if (lua_toboolean (L, -1)) { + lua_pop (L, 1); + break; + } + } + } + + lua_pop (L, 1); + cur = g_list_next (cur); + } + } + } + + return 0; +} + /* Image functions */ static gint lua_image_get_width (lua_State *L) -- 2.39.5