commit 838d9dd: [Minor] Selectors: support extracting task cache variables

[Andrew Lewis]

Author: Andrew Lewis
Date: 2020-12-09 16:46:31 +0200
URL: https://github.com/rspamd/rspamd/commit/838d9dda1a443493d6348b7ce1245717d654e839 (refs/pull/3574/head)

[Minor] Selectors: support extracting task cache variables

---
 lualib/lua_selectors/extractors.lua | 19 +++++++++++++++++++
 test/lua/unit/selectors.lua         | 11 +++++++++++
 2 files changed, 30 insertions(+)

diff --git a/lualib/lua_selectors/extractors.lua b/lualib/lua_selectors/extractors.lua
index 3f6dc61ee..0c82d1499 100644
--- a/lualib/lua_selectors/extractors.lua
+++ b/lualib/lua_selectors/extractors.lua
@@ -363,6 +363,25 @@ e.g. `get_user`]],
 the second argument is optional and defines the type (string by default)]],
     ['args_schema'] = {ts.string, ts.string:is_optional()}
   },
+  -- Get value of specific key from task cache
+  ['task_cache'] = {
+    ['get_value'] = function(task, args)
+      local val = task:cache_get(args[1])
+      if not val then
+        return
+      end
+      if type(val) == 'table' then
+        if not val[1] then
+          return
+        end
+        return val, 'string_list'
+      end
+      return val, 'string'
+    end,
+    ['description'] = [[Get value of specific key from task cache. The first argument must be
+the key name]],
+    ['args_schema'] = {ts.string}
+  },
   -- Get specific HTTP request header. The first argument must be header name.
   ['request_header'] = {
     ['get_value'] = function(task, args)
diff --git a/test/lua/unit/selectors.lua b/test/lua/unit/selectors.lua
index 1f7252a77..270d06941 100644
--- a/test/lua/unit/selectors.lua
+++ b/test/lua/unit/selectors.lua
@@ -27,6 +27,8 @@ context("Selectors test", function()
     task:process_message()
     task:get_mempool():set_variable("int_var", 1)
     task:get_mempool():set_variable("str_var", "str 1")
+    task:cache_set('cachevar1', 'hello\x00world')
+    task:cache_set('cachevar2', {'hello', 'world'})
     if not res then
       assert_true(false, "failed to load message")
     end
@@ -358,6 +360,15 @@ context("Selectors test", function()
       selector = "header('Subject').regexp('.*').first",
       expect = {"Second, lower-cased header subject"}
     },
+
+    ["task cache string"] = {
+      selector = "task_cache('cachevar1')",
+      expect = {"hello\x00world"}
+    },
+    ["task cache table"] = {
+      selector = "task_cache('cachevar2')",
+      expect = {{"hello", "world"}}
+    },
   }
 
   for case_name, case in lua_util.spairs(cases) do