[Fix] milter_headers: Header fields may be inserted at wrong position.

3 files changed, 51 insertions, 50 deletions

diff --git a/lualib/lua_mime.lua b/lualib/lua_mime.lua
index 50c6ccd43..1087b7fd5 100644
--- a/lualib/lua_mime.lua
+++ b/lualib/lua_mime.lua
@@ -19,6 +19,7 @@ limitations under the License.
 -- This module contains helper functions to modify mime parts
 --]]
 
+local logger = require "rspamd_logger"
 local rspamd_util = require "rspamd_util"
 local rspamd_text = require "rspamd_text"
 local ucl = require "ucl"
@@ -519,21 +520,35 @@ exports.modify_headers = function(task, hdr_alterations)
   local add = hdr_alterations.add or {}
   local remove = hdr_alterations.remove or {}
 
+  local add_headers = {} -- For Milter reply
   local hdr_flattened = {} -- For C API
 
   local function flatten_add_header(hname, hdr)
+    if not add_headers[hname] then
+      add_headers[hname] = {}
+    end
     if not hdr_flattened[hname] then
       hdr_flattened[hname] = {add = {}}
     end
     local add_tbl = hdr_flattened[hname].add
     if hdr.value then
+      table.insert(add_headers[hname], {
+        order = (tonumber(hdr.order) or -1),
+        value = hdr.value,
+      })
       table.insert(add_tbl, {tonumber(hdr.order) or -1, hdr.value})
     elseif type(hdr) == 'table' then
       for _,v in ipairs(hdr) do
-        table.insert(add_tbl, {-1, v})
+        flatten_add_header(hname, v)
       end
     elseif type(hdr) == 'string' then
+      table.insert(add_headers[hname], {
+        order = -1,
+        value = hdr,
+      })
       table.insert(add_tbl, {-1, hdr})
+    else
+      logger.errx(task, 'invalid modification of header: %s', hdr)
     end
   end
   if hdr_alterations.order then
@@ -566,7 +581,7 @@ exports.modify_headers = function(task, hdr_alterations)
   end
 
   task:set_milter_reply({
-    add_headers = hdr_alterations.add,
+    add_headers = add_headers,
     remove_headers = hdr_alterations.remove
   })
 
diff --git a/src/libserver/milter.c b/src/libserver/milter.c
index f8ed7be32..8c78a6c7a 100644
--- a/src/libserver/milter.c
+++ b/src/libserver/milter.c
@@ -1652,17 +1652,21 @@ rspamd_milter_remove_header_safe (struct rspamd_milter_session *session,
 		hname = g_string_new (key);
 		hvalue = g_string_new ("");
 
-		if (nhdr >= 1) {
-			rspamd_milter_send_action (session,
-					RSPAMD_MILTER_CHGHEADER,
-					nhdr, hname, hvalue);
+		if (nhdr > 0) {
+			if (ar->len >= nhdr) {
+				rspamd_milter_send_action (session,
+						RSPAMD_MILTER_CHGHEADER,
+						nhdr, hname, hvalue);
+				priv->cur_hdr --;
+			}
 		}
-		else if (nhdr == 0 && ar->len > 0) {
+		else if (nhdr == 0) {
 			/* We need to clear all headers */
 			for (i = ar->len; i > 0; i --) {
 				rspamd_milter_send_action (session,
 						RSPAMD_MILTER_CHGHEADER,
 						i, hname, hvalue);
+				priv->cur_hdr --;
 			}
 		}
 		else {
@@ -1671,11 +1675,17 @@ rspamd_milter_remove_header_safe (struct rspamd_milter_session *session,
 				rspamd_milter_send_action (session,
 						RSPAMD_MILTER_CHGHEADER,
 						ar->len + nhdr + 1, hname, hvalue);
+				priv->cur_hdr --;
 			}
 		}
 
 		g_string_free (hname, TRUE);
 		g_string_free (hvalue, TRUE);
+
+		if (priv->cur_hdr < 0) {
+			msg_err_milter("negative header count after removing %s", key);
+			priv->cur_hdr = 0;
+		}
 	}
 }
 
@@ -1715,10 +1725,19 @@ rspamd_milter_extract_single_header (struct rspamd_milter_session *session,
 			else {
 				/* Calculate negative offset */
 
-				if (-idx <= priv->cur_hdr) {
+				if (idx == -1) {
+					rspamd_milter_send_action (session,
+							RSPAMD_MILTER_ADDHEADER,
+							hname, hvalue);
+				}
+				else if (-idx <= priv->cur_hdr) {
+					/*
+					 * Note: We should account MTA's own "Received:" field
+					 * which wasn't passed by Milter's header command.
+					 */
 					rspamd_milter_send_action (session,
 							RSPAMD_MILTER_INSHEADER,
-							priv->cur_hdr + idx + 1,
+							priv->cur_hdr + idx + 2,
 							hname, hvalue);
 				}
 				else {
@@ -1735,6 +1754,8 @@ rspamd_milter_extract_single_header (struct rspamd_milter_session *session,
 					hname, hvalue);
 		}
 
+		priv->cur_hdr ++;
+
 		g_string_free (hname, TRUE);
 		g_string_free (hvalue, TRUE);
 	}
diff --git a/src/plugins/lua/milter_headers.lua b/src/plugins/lua/milter_headers.lua
index 750b9c453..d83115fc5 100644
--- a/src/plugins/lua/milter_headers.lua
+++ b/src/plugins/lua/milter_headers.lua
@@ -195,48 +195,13 @@ local function milter_headers(task)
 
   local function add_header(name, value, stop_chars, order)
     local hname = settings.routines[name].header
-    if order then
-      if not add[hname] then
-        add[hname] = {
-          order = order,
-          value = lua_util.fold_header(task, hname, value, stop_chars)
-        }
-      else
-        if not add[hname][1] then
-          -- Convert to a table
-          add[hname] = {
-            [1] = add[hname]
-          }
-        end
-
-        table.insert(add[hname], {
-          order = order,
-          value = lua_util.fold_header(task, hname, value, stop_chars)
-        })
-      end
-    else
-      if not add[hname] then
-        add[hname] = lua_util.fold_header(task, hname, value, stop_chars)
-      else
-        if not add[hname][1] then
-          -- Convert to a table
-          add[hname] = {
-            [1] = add[hname]
-          }
-        end
-
-        if settings.default_headers_order then
-          table.insert(add[hname], {
-            order = settings.default_headers_order,
-            value = lua_util.fold_header(task, hname, value, stop_chars)
-          })
-        else
-          table.insert(add[hname],
-              lua_util.fold_header(task, hname, value, stop_chars))
-        end
-
-      end
+    if not add[hname] then
+      add[hname] = {}
     end
+    table.insert(add[hname], {
+      order = (order or settings.default_headers_order or -1),
+      value = lua_util.fold_header(task, hname, value, stop_chars)
+    })
   end
 
   routines['x-spamd-result'] = function()