Merge pull request #77 from simple-login/email-header

Fix bugs due to an email header can appear multiple times

.github/workflows/main.yml

@@ -36,7 +36,7 @@ jobs:
     - name: Publish to Docker Registry
       uses: elgohr/Publish-Docker-Github-Action@master
       with:
-        name: simplelogin/app
+        name: simplelogin/app-ci
         username: ${{ secrets.DOCKER_USERNAME }}
         password: ${{ secrets.DOCKER_PASSWORD }}
 

app/email_utils.py

@@ -268,9 +268,7 @@ def get_email_domain_part(email):
 
 
 def add_dkim_signature(msg: Message, email_domain: str):
-    if msg["DKIM-Signature"]:
+    delete_header(msg, "DKIM-Signature")
-        LOG.d("Remove DKIM-Signature %s", msg["DKIM-Signature"])
-        del msg["DKIM-Signature"]
 
     # Specify headers in "byte" form
     # Generate message signature
@@ -285,23 +283,24 @@ def add_dkim_signature(msg: Message, email_domain: str):
 
     # remove linebreaks from sig
     sig = sig.replace("\n", " ").replace("\r", "")
-
+    msg["DKIM-Signature"] = sig[len("DKIM-Signature: ") :]
-    msg.add_header("DKIM-Signature", sig[len("DKIM-Signature: ") :])
 
 
 def add_or_replace_header(msg: Message, header: str, value: str):
-    try:
+    """
-        msg.add_header(header, value)
+    Remove all occurrences of `header` and add `header` with `value`.
-    except ValueError:
+    """
-        # the header exists already
+    delete_header(msg, header)
-        msg.replace_header(header, value)
+    msg[header] = value
 
 
 def delete_header(msg: Message, header: str):
     """a header can appear several times in message."""
-    for h in msg._headers:
+    # inspired from https://stackoverflow.com/a/47903323/1428034
-        if h[0].lower() == header.lower():
+    for i in reversed(range(len(msg._headers))):
-            msg._headers.remove(h)
+        header_name = msg._headers[i][0].lower()
+        if header_name == header.lower():
+            del msg._headers[i]
 
 
 def email_belongs_to_alias_domains(email: str) -> bool:

tests/test_email_utils.py

@@ -1,3 +1,5 @@
+from email.message import EmailMessage
+
 from app.email_utils import (
     get_email_name,
     get_email_part,
@@ -5,6 +7,8 @@ from app.email_utils import (
     get_email_domain_part,
     email_belongs_to_alias_domains,
     can_be_used_as_personal_email,
+    delete_header,
+    add_or_replace_header,
 )
 from app.extensions import db
 from app.models import User, CustomDomain
@@ -55,3 +59,26 @@ def test_can_be_used_as_personal_email(flask_client):
     CustomDomain.create(user_id=user.id, domain="ab.cd", verified=True)
     db.session.commit()
     assert not can_be_used_as_personal_email("hey@ab.cd")
+
+
+def test_delete_header():
+    msg = EmailMessage()
+    assert msg._headers == []
+
+    msg["H"] = "abcd"
+    msg["H"] = "xyzt"
+
+    assert msg._headers == [("H", "abcd"), ("H", "xyzt")]
+
+    delete_header(msg, "H")
+    assert msg._headers == []
+
+
+def test_add_or_replace_header():
+    msg = EmailMessage()
+    msg["H"] = "abcd"
+    msg["H"] = "xyzt"
+    assert msg._headers == [("H", "abcd"), ("H", "xyzt")]
+
+    add_or_replace_header(msg, "H", "new")
+    assert msg._headers == [("H", "new")]