[Minor] Add function to inject urls into a task

3 files changed, 35 insertions, 1 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
@@ -227,6 +227,12 @@ LUA_FUNCTION_DEF (task, get_urls);
  */
 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
  * @return {text} the data contained in the 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),
@@ -2268,6 +2275,31 @@ lua_task_has_urls (lua_State * L)
 }
 
 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)
 {
 	LUA_TRACE_POINT;
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}");