commit f688e88: [Minor] Add caching for attachments digests

[Vsevolod Stakhov]

Author: Vsevolod Stakhov
Date: 2020-07-14 16:48:23 +0100
URL: https://github.com/rspamd/rspamd/commit/f688e88f182a0bceb99ce203c2d9933855cbedc3 (HEAD -> master)

[Minor] Add caching for attachments digests

---
 lualib/lua_selectors/common.lua     | 43 +++++++++++++++++++++++++++++++++----
 lualib/lua_selectors/extractors.lua | 10 ++-------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/lualib/lua_selectors/common.lua b/lualib/lua_selectors/common.lua
index cb2ffc37f..336c16695 100644
--- a/lualib/lua_selectors/common.lua
+++ b/lualib/lua_selectors/common.lua
@@ -27,12 +27,10 @@ end
 
 exports.digest_schema = digest_schema
 
-local function create_digest(data, args)
-
-  local encoding = args[1] or 'hex'
+local function create_raw_digest(data, args)
   local ht = args[2] or 'blake2'
 
-  local h, s
+  local h
 
   if ht == 'blake2' then
     -- Hack to be compatible with various 'get_digest' methods
@@ -41,6 +39,13 @@ local function create_digest(data, args)
     h = cr_hash.create_specific(ht):update(data)
   end
 
+  return h
+end
+
+local function encode_digest(h, args)
+  local encoding = args[1] or 'hex'
+
+  local s
   if encoding == 'hex' then
     s = h:hex()
   elseif encoding == 'base32' then
@@ -56,6 +61,36 @@ local function create_digest(data, args)
   return s
 end
 
+local function create_digest(data, args)
+  local h = create_raw_digest(data, args)
+  return encode_digest(h, args)
+end
+
+
+local function get_cached_or_raw_digest(task, idx, mime_part, args)
+  if #args == 0 then
+    -- Optimise as we already have this hash in the API
+    return mime_part:get_digest()
+  end
+
+  local ht = args[2] or 'blake2'
+  local cache_key = 'mp_digest_' .. ht .. tostring(idx)
+
+  local cached = task:cache_get(cache_key)
+
+  if cached then
+    return encode_digest(cached, args)
+  end
+
+  local h = create_raw_digest(mime_part:get_content('raw_parsed'), args)
+  task:cache_set(cache_key, h)
+
+  return encode_digest(h, args)
+end
+
 exports.create_digest = create_digest
+exports.create_raw_digest = create_raw_digest
+exports.get_cached_or_raw_digest = get_cached_or_raw_digest
+exports.encode_digest = encode_digest
 
 return exports
\ No newline at end of file
diff --git a/lualib/lua_selectors/extractors.lua b/lualib/lua_selectors/extractors.lua
index 532c56592..d88eaa216 100644
--- a/lualib/lua_selectors/extractors.lua
+++ b/lualib/lua_selectors/extractors.lua
@@ -134,15 +134,9 @@ uses any type by default)]],
     ['get_value'] = function(task, args)
       local parts = task:get_parts() or E
       local digests = {}
-      for _,p in ipairs(parts) do
+      for i,p in ipairs(parts) do
         if p:get_filename() then
-          if #args == 0 then
-            -- Optimise as we already have this hash in the API
-            table.insert(digests, p:get_digest())
-          else
-            table.insert(digests, common.create_digest(p:get_content('raw_parsed'), args))
-          end
-
+          table.insert(digests, common.get_cached_or_raw_digest(task, i, p, args))
         end
       end