From 7b6227cfaa3b5b4273e6841c8d5811514fff28e2 Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Mon, 7 Jan 2019 16:08:37 +0000
Subject: [PATCH] [Minor] Core: Allow to pass flags to squeezed rules

---
 lualib/lua_squeeze_rules.lua | 19 ++++++++++++++++---
 src/lua/lua_config.c         | 21 ++++++++++++++++++++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/lualib/lua_squeeze_rules.lua b/lualib/lua_squeeze_rules.lua
index d4298957b..1df8a4ec1 100644
--- a/lualib/lua_squeeze_rules.lua
+++ b/lualib/lua_squeeze_rules.lua
@@ -31,8 +31,20 @@ local squeezed_groups = {}
 local function gen_lua_squeeze_function(order)
   return function(task)
     local symbols_disabled = task:cache_get('squeezed_disable')
+    local mime_task = task:has_flag('mime')
     for _,data in ipairs(squeezed_rules[order]) do
-      if not symbols_disabled or not symbols_disabled[data[2]] then
+      local disable = false
+      if symbols_disabled and symbols_disabled[data[2]] then
+        disable = true
+      end
+
+      if data[3] and data[3].flags.mime then
+        if not mime_task then
+          disable = true
+        end
+      end
+
+      if not disable then
         local function real_call()
           return {data[1](task)}
         end
@@ -99,13 +111,14 @@ local function gen_lua_squeeze_function(order)
   end
 end
 
-exports.squeeze_rule = function(s, func)
+exports.squeeze_rule = function(s, func, flags)
   if s then
     if not squeezed_symbols[s] then
       squeezed_symbols[s] = {
         cb = func,
         order = 0,
         sym = s,
+        flags = flags or {}
       }
       lua_util.debugm(SN, rspamd_config, 'squeezed rule: %s', s)
     else
@@ -115,7 +128,7 @@ exports.squeeze_rule = function(s, func)
     -- Unconditionally add function to the squeezed rules
     local id = tostring(#squeezed_rules)
     lua_util.debugm(SN, rspamd_config, 'squeezed unnamed rule: %s', id)
-    table.insert(squeezed_rules[1], {func, 'unnamed: ' .. id})
+    table.insert(squeezed_rules[1], {func, 'unnamed: ' .. id, squeezed_symbols[s]})
   end
 
   if not squeeze_function_ids[1] then
diff --git a/src/lua/lua_config.c b/src/lua/lua_config.c
index cc0ba7aed..31d4af79b 100644
--- a/src/lua/lua_config.c
+++ b/src/lua/lua_config.c
@@ -1278,8 +1278,27 @@ rspamd_register_symbol_fromlua (lua_State *L,
 			/* Push function reference */
 			lua_rawgeti (L, LUA_REGISTRYINDEX, ref);
 
+			/* Flags */
+			lua_createtable (L, 0, 0);
+
+			if (type & SYMBOL_TYPE_MIME_ONLY) {
+				lua_pushstring (L, "mime");
+				lua_pushboolean (L, true);
+				lua_settable (L, -3);
+			}
+			if (type & SYMBOL_TYPE_FINE) {
+				lua_pushstring (L, "fine");
+				lua_pushboolean (L, true);
+				lua_settable (L, -3);
+			}
+			if (type & SYMBOL_TYPE_NOSTAT) {
+				lua_pushstring (L, "nostat");
+				lua_pushboolean (L, true);
+				lua_settable (L, -3);
+			}
+
 			/* Now call for squeeze function */
-			if (lua_pcall (L, 2, 1, err_idx) != 0) {
+			if (lua_pcall (L, 3, 1, err_idx) != 0) {
 				GString *tb = lua_touserdata (L, -1);
 				msg_err_config ("call to squeeze_rule failed: %v", tb);
 
-- 
2.39.5