From 9fa16c0eadb1b11edc41c7cb72ab4cb3493b18f6 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Wed, 6 Jun 2018 16:48:35 +0100
Subject: [PATCH] [Feature] Allow to set task flags via settings

---
 src/plugins/lua/settings.lua | 50 ++++++++++++++++++------------------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/src/plugins/lua/settings.lua b/src/plugins/lua/settings.lua
index e1a4cef74..8e4009162 100644
--- a/src/plugins/lua/settings.lua
+++ b/src/plugins/lua/settings.lua
@@ -37,6 +37,26 @@ local rspamd_regexp = require "rspamd_regexp"
 local ucl = require "ucl"
 local fun = require "fun"
 
+local function apply_settings(task, to_apply)
+  task:set_settings(to_apply)
+  task:cache_set('settings', to_apply)
+  lua_squeeze.handle_settings(task, to_apply)
+
+  if to_apply['add_headers'] or to_apply['remove_headers'] then
+    local rep = {
+      add_headers = to_apply['add_headers'] or {},
+      remove_headers = to_apply['remove_headers'] or {},
+    }
+    task:set_rmilter_reply(rep)
+  end
+
+  if to_apply.flags and type(to_apply.flags) == 'table' then
+    for _,fl in ipairs(to_apply.flags) do
+      task:set_flag(fl)
+    end
+  end
+end
+
 -- Checks for overridden settings within query params and returns 'true' if
 -- settings are overridden
 local function check_query_settings(task)
@@ -47,9 +67,7 @@ local function check_query_settings(task)
     local res,err = parser:parse_string(tostring(query_set))
     if res then
       local settings_obj = parser:get_object()
-      task:set_settings(settings_obj)
-      task:cache_set('settings', settings_obj)
-      lua_squeeze.handle_settings(task, settings_obj)
+      apply_settings(task, settings_obj)
 
       return true
     else
@@ -76,10 +94,7 @@ local function check_query_settings(task)
         nset['default']['actions']['add header'] = ss
       end
 
-      task:set_settings(nset)
-      task:cache_set('settings', nset)
-      lua_squeeze.handle_settings(task, nset)
-
+      apply_settings(task, nset)
       return true
     end
   end
@@ -90,19 +105,8 @@ local function check_query_settings(task)
     local id_str = tostring(settings_id)
     local elt = settings_ids[id_str]
     if elt and elt['apply'] then
-      task:set_settings(elt['apply'])
-      lua_squeeze.handle_settings(task, elt['apply'])
-      task:cache_set('settings', elt['apply'])
-
-      if elt.apply['add_headers'] or elt.apply['remove_headers'] then
-        local rep = {
-          add_headers = elt.apply['add_headers'] or {},
-          remove_headers = elt.apply['remove_headers'] or {},
-        }
-        task:set_rmilter_reply(rep)
-      end
+      apply_settings(task, elt['apply'])
       rspamd_logger.infox(task, "applying settings id %s", id_str)
-
       return true
     end
   end
@@ -340,9 +344,7 @@ local function check_settings(task)
           rspamd_logger.infox(task, "<%1> apply settings according to rule %2",
             task:get_message_id(), s.name)
           if rule['apply'] then
-            task:set_settings(rule['apply'])
-            lua_squeeze.handle_settings(task, rule['apply'])
-            task:cache_set('settings', rule['apply'])
+            apply_settings(task, rule['apply'])
             applied = true
           end
           if rule['symbols'] then
@@ -653,9 +655,7 @@ local function gen_redis_callback(handler, id)
               local obj = parser:get_object()
               rspamd_logger.infox(task, "<%1> apply settings according to redis rule %2",
                 task:get_message_id(), id)
-              task:set_settings(obj)
-              lua_squeeze.handle_settings(task, obj)
-              task:cache_set('settings', obj)
+              apply_settings(task, obj)
               break
             end
           end
-- 
2.39.5