From 66eb93fe53088078eebebf4ebd85f88c0a0ea8f2 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Wed, 17 Mar 2021 10:59:13 +0100
Subject: [PATCH] fix sanitize header

---
 app/email_utils.py | 13 +++++++++++++
 app/utils.py       |  6 ------
 email_handler.py   | 11 ++++++-----
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/app/email_utils.py b/app/email_utils.py
index ba2d889b..40a418f6 100644
--- a/app/email_utils.py
+++ b/app/email_utils.py
@@ -421,6 +421,19 @@ def delete_header(msg: Message, header: str):
             del msg._headers[i]
 
 
+def sanitize_header(msg: Message, header: str):
+    """remove trailing space and remove linebreak from a header"""
+    for i in reversed(range(len(msg._headers))):
+        header_name = msg._headers[i][0].lower()
+        if header_name == header.lower():
+            # msg._headers[i] is a tuple like ('From', 'hey@google.com')
+            if msg._headers[i][1]:
+                msg._headers[i] = (
+                    msg._headers[i][0],
+                    msg._headers[i][1].strip().replace("\n", " "),
+                )
+
+
 def delete_all_headers_except(msg: Message, headers: [str]):
     headers = [h.lower() for h in headers]
 
diff --git a/app/utils.py b/app/utils.py
index 07a6582b..dde4e0b9 100644
--- a/app/utils.py
+++ b/app/utils.py
@@ -66,9 +66,3 @@ def sanitize_email(email_address: str) -> str:
     if email_address:
         return email_address.lower().strip().replace(" ", "").replace("\n", " ")
     return email_address
-
-
-def sanitize_header(header: Optional[str]) -> Optional[str]:
-    if header:
-        return header.strip().replace("\n", " ")
-    return header
diff --git a/email_handler.py b/email_handler.py
index 754dd02a..3115ba3a 100644
--- a/email_handler.py
+++ b/email_handler.py
@@ -103,6 +103,7 @@ from app.email_utils import (
     parse_id_from_bounce,
     spf_pass,
     sl_sendmail,
+    sanitize_header,
 )
 from app.extensions import db
 from app.greylisting import greylisting_needed
@@ -118,7 +119,7 @@ from app.models import (
     TransactionalEmail,
 )
 from app.pgp_utils import PGPException, sign_data_with_pgpy, sign_data
-from app.utils import sanitize_email, sanitize_header
+from app.utils import sanitize_email
 from init_app import load_pgp_public_keys
 from server import create_app, create_light_app
 
@@ -1515,10 +1516,10 @@ def handle(envelope: Envelope) -> str:
     msg = email.message_from_bytes(envelope.original_content)
 
     # sanitize email headers
-    msg["from"] = sanitize_header(msg["from"])
-    msg["to"] = sanitize_header(msg["to"])
-    msg["cc"] = sanitize_header(msg["cc"])
-    msg["reply-to"] = sanitize_header(msg["reply-to"])
+    sanitize_header(msg, "from")
+    sanitize_header(msg, "to")
+    sanitize_header(msg, "cc")
+    sanitize_header(msg, "reply-to")
 
     LOG.d(
         "==>> Handle mail_from:%s, rcpt_tos:%s, header_from:%s, header_to:%s, "