Merge pull request #58 from simple-login/avoid-loop

Avoid loop

app/dashboard/views/custom_domain.py

@@ -5,6 +5,7 @@ from wtforms import StringField, validators
 
 from app.config import EMAIL_SERVERS_WITH_PRIORITY
 from app.dashboard.base import dashboard_bp
+from app.email_utils import get_email_domain_part
 from app.extensions import db
 from app.models import CustomDomain
 
@@ -30,9 +31,15 @@ def custom_domain():
                 return redirect(url_for("dashboard.custom_domain"))
 
             if new_custom_domain_form.validate():
-                new_domain = new_custom_domain_form.domain.data.strip()
+                new_domain = new_custom_domain_form.domain.data.lower().strip()
                 if CustomDomain.get_by(domain=new_domain):
                     flash(f"{new_domain} already added", "warning")
+                elif get_email_domain_part(current_user.email) == new_domain:
+                    flash(
+                        "You cannot add a domain that you are currently using for your personal email. "
+                        "Please change your personal email to your real email",
+                        "error",
+                    )
                 else:
                     new_custom_domain = CustomDomain.create(
                         domain=new_domain, user_id=current_user.id

tests/dashboard/test_custom_domain.py

@@ -0,0 +1,38 @@
+from flask import url_for
+
+from app.extensions import db
+from tests.utils import login
+
+
+def test_add_domain_success(flask_client):
+    user = login(flask_client)
+    user.lifetime = True
+    db.session.commit()
+
+    r = flask_client.post(
+        url_for("dashboard.custom_domain"),
+        data={"form-name": "create", "domain": "ab.cd"},
+        follow_redirects=True,
+    )
+
+    assert r.status_code == 200
+    assert b"New domain ab.cd is created" in r.data
+
+
+def test_add_domain_same_as_user_email(flask_client):
+    """cannot add domain if user personal email uses this domain"""
+    user = login(flask_client)
+    user.lifetime = True
+    db.session.commit()
+
+    r = flask_client.post(
+        url_for("dashboard.custom_domain"),
+        data={"form-name": "create", "domain": "b.c"},  # user email is a@b.c
+        follow_redirects=True,
+    )
+
+    assert r.status_code == 200
+    assert (
+        b"You cannot add a domain that you are currently using for your personal email"
+        in r.data
+    )