[Minor] DCC Upstreams support

2 files changed, 91 insertions, 33 deletions

diff --git a/conf/modules.d/dcc.conf b/conf/modules.d/dcc.conf
index ea774bdca..d7622374a 100644
--- a/conf/modules.d/dcc.conf
+++ b/conf/modules.d/dcc.conf
@@ -14,10 +14,14 @@
 # See https://rspamd.com/doc/tutorials/writing_rules.html for details
 
 dcc {
-  # host = "/var/dcc/dccifd";
-  # Port is only required if `dccifd` listens on a TCP socket
-  # port = 1234
-  timeout = 2s;
+
+  enabled = false;
+
+  # Define local socket or TCP servers in upstreams syntax
+  # When sockets and servers are definined - servers is used!
+  socket = "/var/dcc/dccifd"; # Unix socket
+  #servers = "127.0.0.1:10045" # OR TCP upstreams
+  timeout = 2s; # Timeout to wait for checks
 
   .include(try=true,priority=5) "${DBDIR}/dynamic/dcc.conf"
   .include(try=true,priority=1,duplicate=merge) "$LOCAL_CONFDIR/local.d/dcc.conf"
diff --git a/src/plugins/lua/dcc.lua b/src/plugins/lua/dcc.lua
index 131c83d36..aa3838bcb 100644
--- a/src/plugins/lua/dcc.lua
+++ b/src/plugins/lua/dcc.lua
@@ -20,8 +20,10 @@ limitations under the License.
 local N = 'dcc'
 local symbol_bulk = "DCC_BULK"
 local opts = rspamd_config:get_all_opt(N)
-local logger = require "rspamd_logger"
+local rspamd_logger = require "rspamd_logger"
+local lua_util = require "lua_util"
 local tcp = require "rspamd_tcp"
+local upstream_list = require "rspamd_upstream_list"
 local fun = require "fun"
 local lua_util = require "lua_util"
 
@@ -30,8 +32,8 @@ if confighelp then
     "Check messages for 'bulkiness' using DCC",
     [[
 dcc {
-  host = "/var/dcc/dccifd"; # Unix socket or hostname
-  port = 1234 # Port to use (needed for TCP socket)
+  socket = "/var/dcc/dccifd"; # Unix socket
+  servers = "127.0.0.1:10045" # OR TCP upstreams
   timeout = 2s; # Timeout to wait for checks
 }
 ]])
@@ -42,6 +44,22 @@ local function check_dcc (task)
   -- Connection
   local client = '0.0.0.0'
   local client_ip = task:get_from_ip()
+  local dcc_upstream = nil
+  local upstream = nil
+  local addr = nil
+  local port = nil
+  local retransmits = 2
+
+  if opts['servers'] then
+    dcc_upstream = upstream_list.create(rspamd_config, opts['servers'])
+    upstream = dcc_upstream:get_upstream_round_robin()
+    addr = upstream:get_addr()
+    port = addr:get_port()
+  else
+    lua_util.debugm(N, task, 'using socket %s', opts['socket'])
+    addr = opts['socket']
+  end
+
   if client_ip and client_ip:is_valid() then
     client = client_ip:to_string()
   end
@@ -74,27 +92,64 @@ local function check_dcc (task)
 
   -- Callback function to receive async result from DCC
   local function cb(err, data)
-    if (err) then
-      logger.warnx(task, 'DCC error: %1', err)
-      return
-    end
-    -- Parse the response
-    local _,_,result,disposition,header = tostring(data):find("(.-)\n(.-)\n(.-)\n")
-    logger.debugm(N, task, 'DCC result=%1 disposition=%2 header="%3"',
-      result, disposition, header)
-
-    if header then
-      local _,_,info = header:find("; (.-)$")
-      if (result == 'R') then
-        -- Reject
-        task:insert_result(symbol_bulk, 1.0, info)
-      elseif (result == 'T') then
-        -- Temporary failure
-        logger.warnx(task, 'DCC returned a temporary failure result')
+
+    if err then
+      if retransmits > 0 then
+        retransmits = retransmits - 1
+        -- Select a different upstream or the socket again
+        if opts['servers'] then
+          upstream = dcc_upstream:get_upstream_round_robin()
+          addr = upstream:get_addr()
+          port = addr:get_port()
+        else
+          addr = opts['socket']
+        end
+
+        lua_util.debugm(N, task, "sending query to %s:%s",tostring(addr), port)
+
+        data = {
+          "header\n",
+          client .. "\n",
+          helo .. "\n",
+          envfrom .. "\n",
+          envrcpt .. "\n",
+          "\n",
+          task:get_content()
+        }
+
+        tcp.request({
+          task = task,
+          host = tostring(addr),
+          port = port or 1,
+          timeout = opts['timeout'] or 2.0,
+          shutdown = true,
+          data = data,
+          callback = cb
+        })
+
       else
-        if result ~= 'A' and result ~= 'G' and result ~= 'S' then
-          -- Unknown result
-          logger.warnx(task, 'DCC result error: %1', result);
+        rspamd_logger.errx(task, 'failed to scan, maximum retransmits exceed')
+        upstream:fail()
+      end
+    else
+      -- Parse the response
+      local _,_,result,disposition,header = tostring(data):find("(.-)\n(.-)\n(.-)\n")
+      lua_util.debugm(N, task, 'DCC result=%1 disposition=%2 header="%3"',
+        result, disposition, header)
+
+      if header then
+        local _,_,info = header:find("; (.-)$")
+        if (result == 'R') then
+          -- Reject
+          task:insert_result(symbol_bulk, 1.0, info)
+        elseif (result == 'T') then
+          -- Temporary failure
+          rspamd_logger.warnx(task, 'DCC returned a temporary failure result')
+        else
+          if result ~= 'A' and result ~= 'G' and result ~= 'S' then
+            -- Unknown result
+            rspamd_logger.warnx(task, 'DCC result error: %1', result);
+          end
         end
       end
     end
@@ -112,13 +167,12 @@ local function check_dcc (task)
     task:get_content()
   }
 
-  logger.debugm(N, task, 'sending to dcc: client=%1 helo="%2" envfrom="%3" envrcpt="%4"',
-    client, helo, envfrom, envrcpt)
+  rspamd_logger.warnx(task, "sending to %s:%s",tostring(addr), port)
 
   tcp.request({
     task = task,
-    host = opts['host'],
-    port = opts['port'] or 1,
+    host = tostring(addr),
+    port = port or 1,
     timeout = opts['timeout'] or 2.0,
     shutdown = true,
     data = data,
@@ -127,7 +181,7 @@ local function check_dcc (task)
 end
 
 -- Configuration
-if opts and opts['host'] then
+if opts and ( opts['servers'] or opts['socket'] ) then
   rspamd_config:register_symbol({
     name = symbol_bulk,
     callback = check_dcc
@@ -141,5 +195,5 @@ if opts and opts['host'] then
   })
 else
   lua_util.disable_module(N, "config")
-  logger.infox('DCC module not configured');
+  rspamd_logger.infox('DCC module not configured');
 end