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

2 files changed, 34 insertions, 48 deletions

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()