From 10a6115587ce230ca349f357e67bd2376338ad7c Mon Sep 17 00:00:00 2001
From: Vsevolod Stakhov <vsevolod@highsecure.ru>
Date: Fri, 14 Jul 2017 08:18:16 +0100
Subject: [PATCH] [Feature] Allow to receive signing keys from mempool vars

---
 lualib/dkim_sign_tools.lua            | 25 +++++++++++++++++++------
 src/libserver/mempool_vars_internal.h |  4 ++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/lualib/dkim_sign_tools.lua b/lualib/dkim_sign_tools.lua
index 7c49aef1e..3a34a5682 100644
--- a/lualib/dkim_sign_tools.lua
+++ b/lualib/dkim_sign_tools.lua
@@ -137,10 +137,23 @@ local function prepare_dkim_signing(N, task, settings)
     p.key = settings.domain[dkim_domain].path
   end
 
-  if not (p.key and p.selector) and not
-  (settings.try_fallback or settings.use_redis or settings.selector_map or settings.path_map) then
-    rspamd_logger.debugm(N, task, 'dkim unconfigured and fallback disabled')
-    return false,{}
+  if not p.key and p.selector then
+    local key_var = "dkim_key"
+    local selector_var = "dkim_selector"
+    if N == "arc" then
+      key_var = "arc_key"
+      selector_var = "arc_selector"
+    end
+
+    p.key = task:get_mempool():get_variable(key_var)
+    p.selector = task:get_mempool():get_variable(selector_var)
+
+    if (not p.key or not p.selector) and (not (settings.try_fallback or
+        settings.use_redis or settings.selector_map
+        or settings.path_map)) then
+      rspamd_logger.debugm(N, task, 'dkim unconfigured and fallback disabled')
+      return false,{}
+    end
   end
 
   if not p.key then
@@ -154,14 +167,14 @@ local function prepare_dkim_signing(N, task, settings)
   end
   p.domain = dkim_domain
 
-  if settings.selector_map then
+  if not p.selector and settings.selector_map then
     local data = settings.selector_map:get_key(dkim_domain)
     if data then
       p.selector = data
     end
   end
 
-  if settings.path_map then
+  if not p.key and settings.path_map then
     local data = settings.path_map:get_key(dkim_domain)
     if data then
       p.key = data
diff --git a/src/libserver/mempool_vars_internal.h b/src/libserver/mempool_vars_internal.h
index 229d56d48..903f169f9 100644
--- a/src/libserver/mempool_vars_internal.h
+++ b/src/libserver/mempool_vars_internal.h
@@ -32,6 +32,10 @@
 #define RSPAMD_MEMPOOL_DKIM_SIGNATURE "dkim-signature"
 #define RSPAMD_MEMPOOL_DMARC_CHECKS "dmarc_checks"
 #define RSPAMD_MEMPOOL_DKIM_BH_CACHE "dkim_bh_cache"
+#define RSPAMD_MEMPOOL_DKIM_SIGN_KEY "dkim_key"
+#define RSPAMD_MEMPOOL_DKIM_SIGN_SELECTOR "dkim_selector"
+#define RSPAMD_MEMPOOL_ARC_SIGN_KEY "arc_key"
+#define RSPAMD_MEMPOOL_ARC_SIGN_SELECTOR "arc_selector"
 #define RSPAMD_MEMPOOL_STAT_SIGNATURE "stat_signature"
 
 #endif
-- 
2.39.5