From da7957647314a46a3cf04d47845ffa44a8d5f97c Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Fri, 23 Aug 2019 13:27:27 +0100
Subject: [PATCH] [Minor] Add function to inject urls into a task

---
 src/lua/lua_common.h         |  2 ++
 src/lua/lua_task.c           | 32 ++++++++++++++++++++++++++++++++
 src/lua/lua_url.c            |  2 +-
 src/plugins/lua/multimap.lua |  5 +++--
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h
index bef163c3c..40bbea772 100644
--- a/src/lua/lua_common.h
+++ b/src/lua/lua_common.h
@@ -370,6 +370,8 @@ struct ev_loop *lua_check_ev_base (lua_State *L, gint pos);
 
 struct rspamd_dns_resolver *lua_check_dns_resolver (lua_State *L, gint pos);
 
+struct rspamd_lua_url *lua_check_url (lua_State * L, gint pos);
+
 enum rspamd_lua_parse_arguments_flags {
 	RSPAMD_LUA_PARSE_ARGUMENTS_DEFAULT = 0,
 	RSPAMD_LUA_PARSE_ARGUMENTS_IGNORE_MISSING,
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index 9069d3f85..1664e415f 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -226,6 +226,12 @@ LUA_FUNCTION_DEF (task, get_urls);
  * @return {boolean} true if a task has urls (urls or emails if `need_emails` is true)
  */
 LUA_FUNCTION_DEF (task, has_urls);
+/***
+ * @method task:inject_url(url)
+ * Inserts an url into a task (useful for redirected urls)
+ * @param {lua_url} url url to inject
+ */
+LUA_FUNCTION_DEF (task, inject_url);
 /***
  * @method task:get_content()
  * Get raw content for the specified task
@@ -1092,6 +1098,7 @@ static const struct luaL_reg tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, append_message),
 	LUA_INTERFACE_DEF (task, has_urls),
 	LUA_INTERFACE_DEF (task, get_urls),
+	LUA_INTERFACE_DEF (task, inject_url),
 	LUA_INTERFACE_DEF (task, get_content),
 	LUA_INTERFACE_DEF (task, get_filename),
 	LUA_INTERFACE_DEF (task, get_rawbody),
@@ -2267,6 +2274,31 @@ lua_task_has_urls (lua_State * L)
 	return 1;
 }
 
+static gint
+lua_task_inject_url (lua_State * L)
+{
+	LUA_TRACE_POINT;
+	struct rspamd_task *task = lua_check_task (L, 1);
+	struct rspamd_lua_url *url = lua_check_url (L, 2);
+
+	if (task && task->message && url && url->url) {
+		struct rspamd_url *existing;
+
+		if ((existing = g_hash_table_lookup (MESSAGE_FIELD (task, urls),
+				url->url)) == NULL) {
+			g_hash_table_insert (MESSAGE_FIELD (task, urls), url->url, url->url);
+		}
+		else {
+			existing->count ++;
+		}
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 0;
+}
+
 static gint
 lua_task_get_content (lua_State * L)
 {
diff --git a/src/lua/lua_url.c b/src/lua/lua_url.c
index 5dec8ddaa..b30e560c9 100644
--- a/src/lua/lua_url.c
+++ b/src/lua/lua_url.c
@@ -102,7 +102,7 @@ static const struct luaL_reg urllib_f[] = {
 	{NULL, NULL}
 };
 
-static struct rspamd_lua_url *
+struct rspamd_lua_url *
 lua_check_url (lua_State * L, gint pos)
 {
 	void *ud = rspamd_lua_check_udata (L, pos, "rspamd{url}");
diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index da8ab01e3..ae97afeb8 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -430,12 +430,13 @@ local function multimap_callback(task, rule)
   local pre_filter = rule['prefilter']
 
   local function match_element(r, value, callback)
-    lua_util.debugm(N, task, 'check value %s for multimap %s', value,
-        rule.symbol)
     if not value then
       return false
     end
 
+    lua_util.debugm(N, task, 'check value %s for multimap %s', value,
+        rule.symbol)
+
     local ret = false
 
     if r['cdb'] then
-- 
2.39.5