From 7b4f8690e093b2fc9982a8e91a7e39ef16efe331 Mon Sep 17 00:00:00 2001
From: Andrew Lewis <nerf@judo.za.org>
Date: Wed, 8 Feb 2017 12:24:04 +0200
Subject: [PATCH] [Feature] Multimap: Received header processing

---
 src/plugins/lua/multimap.lua | 91 ++++++++++++++++++++++++++++++++++++
 1 file changed, 91 insertions(+)

diff --git a/src/plugins/lua/multimap.lua b/src/plugins/lua/multimap.lua
index 3df9e0a0c..dd474d0c8 100644
--- a/src/plugins/lua/multimap.lua
+++ b/src/plugins/lua/multimap.lua
@@ -22,6 +22,7 @@ local cdb = require "rspamd_cdb"
 local util = require "rspamd_util"
 local regexp = require "rspamd_regexp"
 local rspamd_expression = require "rspamd_expression"
+local rspamd_ip = require "rspamd_ip"
 local redis_params
 local fun = require "fun"
 local N = 'multimap'
@@ -68,6 +69,9 @@ local value_types = {
   country = {
     get_value = function(val) return val end,
   },
+  received = {
+    get_value = function(val) return val end,
+  },
   mempool = {
     get_value = function(val) return val end,
   },
@@ -531,6 +535,45 @@ local function multimap_callback(task, rule)
     match_rule(r, fn)
   end
 
+  local function match_received_header(r, pos, total, h)
+    local filter = r['filter'] or 'real_ip'
+    local v = h[filter]
+    if v then
+      local min_pos = tonumber(r['min_pos'])
+      local max_pos = tonumber(r['max_pos'])
+      if min_pos then
+        if min_pos < 0 then
+          if (pos - (min_pos*-1)) < pos then
+            return
+          end
+        else
+          if pos < min_pos then
+            return
+          end
+        end
+      end
+      if max_pos then
+        if max_pos < 0 then
+          if (pos - (max_pos*-1)) > pos then
+            return
+          end
+        else
+          if pos > max_pos then
+            return
+          end
+        end
+      end
+      if filter == 'real_ip' or filter == 'from_ip' then
+        v = rspamd_ip.from_string(v)
+        if v and v:is_valid() then
+          match_rule(r, v)
+        end
+      else
+        match_rule(r, v)
+      end
+    end
+  end
+
   local function match_content(r)
     local data
 
@@ -660,6 +703,15 @@ local function multimap_callback(task, rule)
         match_rule(rule, var)
       end
     end,
+    received = function()
+      local hdrs = task:get_received_headers()
+      if hdrs and hdrs[1] then
+        local num_hdrs = #hdrs
+        for pos, h in ipairs(hdrs) do
+          match_received_header(rule, pos, num_hdrs, h)
+        end
+      end
+    end,
   }
   process_rule_funcs.ip = process_rule_funcs.dnsbl
   local f = process_rule_funcs[rt]
@@ -758,6 +810,45 @@ local function add_multimap_rule(key, newrule)
           rspamd_logger.warnx(rspamd_config, 'Cannot add rule: map doesn\'t exists: %1',
             newrule['map'])
         end
+      elseif newrule['type'] == 'received' then
+        local filter = newrule['filter'] or 'real_ip'
+        if filter == 'real_ip' or filter == 'from_ip' then
+          newrule['radix'] = rspamd_config:add_map ({
+            url = newrule['map'],
+            description = newrule['description'],
+            type = 'radix'
+          })
+          if newrule['radix'] then
+            ret = true
+          end
+        else
+          if newrule['regexp'] then
+            newrule['hash'] = rspamd_config:add_map ({
+              url = newrule['map'],
+              description = newrule['description'],
+              type = 'regexp'
+            })
+          else
+            newrule['hash'] = rspamd_config:add_map ({
+              url = newrule['map'],
+              description = newrule['description'],
+              type = 'hash'
+            })
+          end
+          if newrule['hash'] then
+            ret = true
+            if type(newrule['map']) == 'string' then
+              urls[newrule['map']] = {
+                type = newrule['type'],
+                map = newrule['hash'],
+                regexp = newrule['regexp']
+              }
+            end
+          else
+            rspamd_logger.warnx(rspamd_config, 'Cannot add rule: map doesn\'t exists: %1',
+              newrule['map'])
+          end
+        end
       elseif newrule['type'] == 'header'
         or newrule['type'] == 'rcpt'
         or newrule['type'] == 'from'
-- 
2.39.5