[Feature] Selectors: Support filter_map and apply_map functions

diff --git a/lualib/lua_selectors.lua b/lualib/lua_selectors.lua
index b4496a42814522016ee234a7bfb997eb645c352a..9af5f8f92b04137e33c903a65de6df24f5a1b946 100644 (file)
--- a/lualib/lua_selectors.lua
+++ b/lualib/lua_selectors.lua
@@ -26,7 +26,10 @@ limitations under the License.
 -- Typical selector looks like this: header(User).lower.substring(1, 2):ip
 --]]
 
-local exports = {}
+local exports = {
+  maps = {} -- Defined for selectors maps, must be indexed by name
+}
+
 local logger = require 'rspamd_logger'
 local fun = require 'fun'
 local lua_util = require "lua_util"
@@ -551,6 +554,56 @@ the second argument is optional hash type (`blake2`, `sha256`, `sha1`, `sha512`,
     ['description'] = 'Regexp matching',
     ['args_schema'] = {ts.string}
   },
+  -- Returns a value if it exists in some map (or acts like a `filter` function)
+  ['filter_map'] = {
+    ['types'] = {
+      ['string'] = true
+    },
+    ['map_type'] = 'string',
+    ['process'] = function(inp, t, args)
+      local map = exports.maps[args[1]]
+
+      if not map then
+        logger.errx('invalid map name: %s', args[1])
+        return nil
+      end
+
+      local res = map:get_key(inp)
+
+      if res then
+        return inp,t
+      end
+
+      return nil
+    end,
+    ['description'] = 'Returns a value if it exists in some map (or acts like a `filter` function)',
+    ['args_schema'] = {ts.string}
+  },
+  -- Returns a value from some map corresponding to some key (or acts like a `map` function)
+  ['apply_map'] = {
+    ['types'] = {
+      ['string'] = true
+    },
+    ['map_type'] = 'string',
+    ['process'] = function(inp, t, args)
+      local map = exports.maps[args[1]]
+
+      if not map then
+        logger.errx('invalid map name: %s', args[1])
+        return nil
+      end
+
+      local res = map:get_key(inp)
+
+      if res then
+        return res,t
+      end
+
+      return nil
+    end,
+    ['description'] = 'Returns a value from some map corresponding to some key (or acts like a `map` function)',
+    ['args_schema'] = {ts.string}
+  },
   -- Drops input value and return values from function's arguments or an empty string
   ['id'] = {
     ['types'] = {