[Feature] Allow to pass extradata from rspamd to rmilter

2 files changed, 41 insertions, 1 deletions

diff --git a/src/libserver/protocol.c b/src/libserver/protocol.c
index de1e08592..e7192dff2 100644
--- a/src/libserver/protocol.c
+++ b/src/libserver/protocol.c
@@ -984,6 +984,7 @@ rspamd_protocol_write_ucl (struct rspamd_task *task)
 	ucl_object_t *top = NULL, *obj;
 	GHashTableIter hiter;
 	GString *dkim_sig;
+	const ucl_object_t *rmilter_reply;
 	gpointer h, v;
 
 	g_hash_table_iter_init (&hiter, task->results);
@@ -1021,6 +1022,13 @@ rspamd_protocol_write_ucl (struct rspamd_task *task)
 					"dkim-signature", 0, false);
 	}
 
+	rmilter_reply = rspamd_mempool_get_variable (task->task_pool, "rmilter-reply");
+
+	if (rmilter_reply) {
+		ucl_object_insert_key (top, ucl_object_ref (rmilter_reply),
+				"rmilter", 0, false);
+	}
+
 	return top;
 }
 
diff --git a/src/lua/lua_task.c b/src/lua/lua_task.c
index d18eaa4ea..e11c3e5c1 100644
--- a/src/lua/lua_task.c
+++ b/src/lua/lua_task.c
@@ -534,11 +534,23 @@ LUA_FUNCTION_DEF (task, get_settings);
 /***
  * @method task:get_settings_id()
  * Get numeric hash of settings id if specified for this task. 0 is returned otherwise.
- * @param {any} obj any lua object that corresponds to the settings format
+ * @return {number} settings-id hash
  */
 LUA_FUNCTION_DEF (task, get_settings_id);
 
 /***
+ * @method task:set_rmilter_reply(obj)
+ * Set special reply for rmilter
+ * @param {any} obj any lua object that corresponds to the settings format
+ * @example
+task:set_rmilter_reply({
+	add_headers = {{'X-Lua', 'test'}},
+	remove_headers = {'DKIM-Signature},
+})
+ */
+LUA_FUNCTION_DEF (task, set_rmilter_reply);
+
+/***
  * @method task:process_re(params)
  * Processes the specified regexp and returns number of captures (cached or new)
  * Params is the table with the follwoing fields (mandatory fields are marked with `*`):
@@ -716,6 +728,7 @@ static const struct luaL_reg tasklib_m[] = {
 	LUA_INTERFACE_DEF (task, set_flag),
 	LUA_INTERFACE_DEF (task, get_flags),
 	LUA_INTERFACE_DEF (task, has_flag),
+	LUA_INTERFACE_DEF (task, set_rmilter_reply),
 	{"__tostring", rspamd_lua_class_tostring},
 	{NULL, NULL}
 };
@@ -2843,6 +2856,25 @@ lua_task_set_settings (lua_State *L)
 }
 
 static gint
+lua_task_set_rmilter_reply (lua_State *L)
+{
+	struct rspamd_task *task = lua_check_task (L, 1);
+	ucl_object_t *reply;
+
+	reply = ucl_object_lua_import (L, 2);
+
+	if (reply != NULL && task != NULL) {
+		rspamd_mempool_set_variable (task->task_pool, "rmilter-reply",
+				reply, (rspamd_mempool_destruct_t)ucl_object_unref);
+	}
+	else {
+		return luaL_error (L, "invalid arguments");
+	}
+
+	return 0;
+}
+
+static gint
 lua_task_get_settings (lua_State *L)
 {
 	struct rspamd_task *task = lua_check_task (L, 1);