Only allow to set default aliases for ownership AND mx verified domains

app/api/views/setting.py

@@ -7,12 +7,11 @@
 from app.models import (
     User,
     AliasGeneratorEnum,
-    SLDomain,
-    CustomDomain,
     SenderFormatEnum,
     AliasSuffixEnum,
 )
 from app.proton.proton_unlink import perform_proton_account_unlink
+from app.user_settings import CannotSetAlias, set_default_alias_domain
 
 
 def setting_to_dict(user: User):
@@ -85,25 +84,11 @@ def update_setting():
 
     if "random_alias_default_domain" in data:
         default_domain = data["random_alias_default_domain"]
-        sl_domain: SLDomain = SLDomain.get_by(domain=default_domain)
-        if sl_domain:
-            if sl_domain.premium_only and not user.is_premium():
-                return jsonify(error="You cannot use this domain"), 400
-
-            user.default_alias_public_domain_id = sl_domain.id
-            user.default_alias_custom_domain_id = None
-        else:
-            custom_domain = CustomDomain.get_by(domain=default_domain)
-            if not custom_domain:
-                return jsonify(error="invalid domain"), 400
-
-            # sanity check
-            if custom_domain.user_id != user.id or not custom_domain.verified:
-                LOG.w("%s cannot use domain %s", user, default_domain)
-                return jsonify(error="invalid domain"), 400
-            else:
-                user.default_alias_custom_domain_id = custom_domain.id
-                user.default_alias_public_domain_id = None
+        try:
+            set_default_alias_domain(user, default_domain)
+        except CannotSetAlias as e:
+            LOG.w("%s cannot use domain %s", user, default_domain)
+            return jsonify(error=e.msg), 400
 
     Session.commit()
     return jsonify(setting_to_dict(user))

app/custom_domain_validation.py

@@ -8,6 +8,7 @@
     DNSClient,
     get_network_dns_client,
 )
+from app.log import LOG
 from app.models import CustomDomain
 from app.user_audit_log_utils import emit_user_audit_log, UserAuditLogAction
 from app.utils import random_string
@@ -230,6 +231,14 @@ def validate_mx_records(
                 for mx_domain in mx_domains[prio]:
                     errors.append(f"{prio} {mx_domain}")
             return DomainValidationResult(success=False, errors=errors)
+        elif not custom_domain.ownership_verified:
+            # Only verify MX if the ownership is verified
+            LOG.i(
+                f"Not marking custom domain {custom_domain.id} ({custom_domain.domain}) as MX verified: ownership is not verified"
+            )
+            return DomainValidationResult(
+                success=False, errors=["The domain ownership must be verified first"]
+            )
         else:
             custom_domain.verified = True
             emit_user_audit_log(

app/dashboard/views/domain_detail.py

@@ -73,6 +73,16 @@ def domain_detail_dns(custom_domain_id):
                 ownership_errors = ownership_validation_result.errors
 
         elif request.form.get("form-name") == "check-mx":
+            if not custom_domain.ownership_verified:
+                flash("You need to verify the domain ownership first", "error")
+                return redirect(
+                    url_for(
+                        "dashboard.domain_detail_dns",
+                        custom_domain_id=custom_domain.id,
+                        _anchor="ownership-form",
+                    )
+                )
+
             mx_validation_result = domain_validator.validate_mx_records(custom_domain)
             if mx_validation_result.success:
                 flash(

app/models.py

@@ -1097,6 +1097,9 @@ def available_domains_for_random_alias(
             res.append((True, domain.domain))
 
         for custom_domain in self.verified_custom_domains():
+            # Request domains to also be MX verified
+            if not custom_domain.verified:
+                continue
             res.append((False, custom_domain.domain))
 
         return res

app/user_settings.py

@@ -27,30 +27,38 @@ def set_default_alias_domain(user: User, domain_name: Optional[str]):
         Session.flush()
         return
 
-    sl_domain: SLDomain = SLDomain.get_by(domain=domain_name)
-    if sl_domain:
-        if sl_domain.hidden:
-            LOG.i(f"User {user} has tried to set up a hidden domain as default domain")
-            raise CannotSetAlias("Domain does not exist")
-        if sl_domain.premium_only and not user.is_premium():
-            LOG.i(f"User {user} has tried to set up a premium domain as default domain")
-            raise CannotSetAlias("You cannot use this domain")
+    # only accept a domain that we would offer to the user in the first place. Relying on
+    # the same list that the settings page and the API expose keeps this check from
+    # drifting away from what the user is actually allowed to pick
+    is_sl_domain = None
+    for is_public, available_domain in user.available_domains_for_random_alias():
+        if available_domain == domain_name:
+            is_sl_domain = is_public
+            break
+
+    # is_sl_domain will be true/false if the domain is allowed for the user
+    if is_sl_domain is None:
+        LOG.i(
+            f"User {user} has tried to set up {domain_name} as default domain but it is not available to them"
+        )
+        raise CannotSetAlias("Domain does not exist or it hasn't been verified")
+
+    if is_sl_domain:
+        sl_domain: SLDomain = SLDomain.get_by(domain=domain_name)
         LOG.i(f"User {user} has set public {sl_domain} as default domain")
         user.default_alias_public_domain_id = sl_domain.id
         user.default_alias_custom_domain_id = None
         Session.flush()
         return
-    custom_domain = CustomDomain.get_by(domain=domain_name)
+
+    custom_domain: Optional[CustomDomain] = CustomDomain.get_by(
+        domain=domain_name, user_id=user.id
+    )
     if not custom_domain:
         LOG.i(
             f"User {user} has tried to set up an non existing domain as default domain"
         )
         raise CannotSetAlias("Domain does not exist or it hasn't been verified")
-    if custom_domain.user_id != user.id or not custom_domain.verified:
-        LOG.i(
-            f"User {user} has tried to set domain {custom_domain} as default domain that does not belong to the user or that is not verified"
-        )
-        raise CannotSetAlias("Domain does not exist or it hasn't been verified")
     LOG.i(f"User {user} has set custom {custom_domain} as default domain")
     user.default_alias_public_domain_id = None
     user.default_alias_custom_domain_id = custom_domain.id

tests/api/test_setting.py

@@ -45,7 +45,11 @@ def test_update_settings_alias_generator(flask_client):
 def test_update_settings_random_alias_default_domain(flask_client):
     user = login(flask_client)
     custom_domain = CustomDomain.create(
-        domain=random_domain(), verified=True, user_id=user.id, flush=True
+        domain=random_domain(),
+        verified=True,
+        ownership_verified=True,
+        user_id=user.id,
+        flush=True,
     )
     assert user.default_random_alias_domain() == "sl.lan"
 
@@ -67,6 +71,23 @@ def test_update_settings_random_alias_default_domain(flask_client):
     assert user.default_random_alias_domain() == custom_domain.domain
 
 
+def test_update_settings_random_alias_default_domain_without_ownership(flask_client):
+    user = login(flask_client)
+    custom_domain = CustomDomain.create(
+        domain=random_domain(),
+        verified=True,
+        ownership_verified=False,
+        user_id=user.id,
+        flush=True,
+    )
+
+    r = flask_client.patch(
+        "/api/setting", json={"random_alias_default_domain": custom_domain.domain}
+    )
+    assert r.status_code == 400
+    assert user.default_alias_custom_domain_id is None
+
+
 def test_update_settings_sender_format(flask_client):
     user = login(flask_client)
     assert user.sender_format == SenderFormatEnum.AT.value

tests/test_custom_domain_validation.py

@@ -22,8 +22,13 @@ def setup_module():
     Session.commit()
 
 
-def create_custom_domain(domain: str) -> CustomDomain:
-    return CustomDomain.create(user_id=user.id, domain=domain, commit=True)
+def create_custom_domain(domain: str, ownership_verified: bool = False) -> CustomDomain:
+    return CustomDomain.create(
+        user_id=user.id,
+        domain=domain,
+        ownership_verified=ownership_verified,
+        commit=True,
+    )
 
 
 def test_custom_domain_validation_get_dkim_records():
@@ -479,7 +484,7 @@ def test_custom_domain_validation_validate_mx_records_success():
     dns_client = InMemoryDNSClient()
     validator = CustomDomainValidation(random_domain(), dns_client)
 
-    domain = create_custom_domain(random_domain())
+    domain = create_custom_domain(random_domain(), ownership_verified=True)
 
     mx_records_by_prio = validator.get_expected_mx_records(domain)
     dns_records = {
@@ -496,6 +501,26 @@ def test_custom_domain_validation_validate_mx_records_success():
     assert db_domain.verified is True
 
 
+def test_custom_domain_validation_validate_mx_records_without_ownership_failure():
+    dns_client = InMemoryDNSClient()
+    validator = CustomDomainValidation(random_domain(), dns_client)
+
+    domain = create_custom_domain(random_domain(), ownership_verified=False)
+
+    mx_records_by_prio = validator.get_expected_mx_records(domain)
+    dns_records = {
+        priority: mx_records_by_prio[priority].allowed
+        for priority in mx_records_by_prio
+    }
+    dns_client.set_mx_records(domain.domain, dns_records)
+    res = validator.validate_mx_records(domain)
+
+    assert res.success is False
+
+    db_domain = CustomDomain.get_by(id=domain.id)
+    assert db_domain.verified is False
+
+
 def test_custom_domain_validation_validate_mx_records_partner_domain_success():
     """Test MX validation for partner domains with custom MX configuration."""
     dns_client = InMemoryDNSClient()
@@ -506,7 +531,7 @@ def test_custom_domain_validation_validate_mx_records_partner_domain_success():
         random_domain(), dns_client, partner_domains={partner_id: partner_mx_domain}
     )
 
-    domain = create_custom_domain(random_domain())
+    domain = create_custom_domain(random_domain(), ownership_verified=True)
     domain.partner_id = partner_id
     Session.commit()
 

tests/user_settings/test_set_default_alias_domain.py

@@ -20,6 +20,7 @@ def setup_module():
         catch_all=True,
         domain=random_token() + ".com",
         verified=True,
+        ownership_verified=True,
         flush=True,
     ).domain
     sl_domain_name = SLDomain.create(
@@ -95,6 +96,7 @@ def test_set_other_user_custom_domain():
         catch_all=True,
         domain=random_token() + ".com",
         verified=True,
+        ownership_verified=True,
     ).domain
     Session.flush()
     with pytest.raises(user_settings.CannotSetAlias):
@@ -122,6 +124,21 @@ def test_set_custom_domain():
     assert user.default_alias_custom_domain_id == domain.id
 
 
+def test_set_custom_domain_without_ownership_verified():
+    user = User.get(user_id)
+    user.lifetime = True
+    domain = CustomDomain.get_by(domain=custom_domain_name)
+    domain.verified = True
+    domain.ownership_verified = False
+    Session.flush()
+    try:
+        with pytest.raises(user_settings.CannotSetAlias):
+            user_settings.set_default_alias_domain(user, custom_domain_name)
+    finally:
+        domain.ownership_verified = True
+        Session.flush()
+
+
 def test_set_invalid_custom_domain():
     user = User.get(user_id)
     with pytest.raises(user_settings.CannotSetAlias):