From 18c72f08061f4e0a383784e153dd5defea103529 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?= Date: Tue, 7 Jul 2026 14:14:57 +0200 Subject: [PATCH 1/2] Increase retries to be ensured at most once a day --- app/config.py | 2 + tasks/check_custom_domains.py | 16 ++-- tests/test_custom_domain_validation.py | 101 +++++++++++++++++++++++++ 3 files changed, 112 insertions(+), 7 deletions(-) diff --git a/app/config.py b/app/config.py index 9e8fdc80d..1b4434333 100644 --- a/app/config.py +++ b/app/config.py @@ -738,3 +738,5 @@ def read_hex_data(key: string, default: bytes) -> bytes: DROP_PGP_KEY_ATTACHMENTS_ON_REPLY = "DROP_PGP_KEY_ATTACHMENTS_ON_REPLY" in os.environ ENFORCE_OAUTH_CLIENT_APPROVED = "ENFORCE_OAUTH_CLIENT_APPROVED" in os.environ + +MAX_DOMAIN_CHECKS = max(int(os.environ.get("MAX_DOMAIN_CHECKS", 4)), 1) diff --git a/tasks/check_custom_domains.py b/tasks/check_custom_domains.py index 00dbc9bbf..4d86cbcd0 100644 --- a/tasks/check_custom_domains.py +++ b/tasks/check_custom_domains.py @@ -67,16 +67,18 @@ def check_single_custom_domain(custom_domain: CustomDomain): if not is_mx_equivalent(mx_domains, expected_custom_domains): user = custom_domain.user LOG.w( - "The MX record is not correctly set for %s %s %s", - custom_domain, - user, - mx_domains, + f"The MX record is not correctly set for domain {custom_domain} of user {user}. Got {mx_domains}. Retried {custom_domain.nb_failed_checks} days", ) - custom_domain.nb_failed_checks += 1 + if ( + not custom_domain.updated_at + or custom_domain.updated_at <= arrow.now().shift(days=-1) + ): + # Only update it once a day + custom_domain.nb_failed_checks += 1 - # send alert if fail for 3 consecutive days - if custom_domain.nb_failed_checks > 2: + # send alert if fail for MAX_DOMAIN_CHECKS consecutive days + if custom_domain.nb_failed_checks > config.MAX_DOMAIN_CHECKS: domain_dns_url = f"{config.URL}/dashboard/domains/{custom_domain.id}/dns" LOG.w("Alert domain MX check fails %s about %s", user, custom_domain) send_email_with_rate_control( diff --git a/tests/test_custom_domain_validation.py b/tests/test_custom_domain_validation.py index fbeb54d80..7a795d228 100644 --- a/tests/test_custom_domain_validation.py +++ b/tests/test_custom_domain_validation.py @@ -296,6 +296,74 @@ def test_custom_domain_validation_validate_dkim_records_success(): assert db_domain.dkim_verified is True +def test_custom_domain_validation_validate_dkim_records_partner_domain_success(): + """Test DKIM validation for partner domains with custom DKIM configuration.""" + dkim_domain = random_domain() + partner_dkim_domain = random_domain() + dns_client = InMemoryDNSClient() + + partner_id = get_proton_partner().id + validator = CustomDomainValidation( + dkim_domain, dns_client, partner_domains={partner_id: partner_dkim_domain} + ) + + user_domain = random_domain() + domain = create_custom_domain(user_domain) + domain.partner_id = partner_id + Session.commit() + + # Set correct partner DKIM records (preferred) + dns_client.set_cname_record( + f"dkim._domainkey.{user_domain}", f"dkim._domainkey.{partner_dkim_domain}" + ) + dns_client.set_cname_record( + f"dkim02._domainkey.{user_domain}", f"dkim02._domainkey.{partner_dkim_domain}" + ) + dns_client.set_cname_record( + f"dkim03._domainkey.{user_domain}", f"dkim03._domainkey.{partner_dkim_domain}" + ) + + res = validator.validate_dkim_records(domain) + assert len(res) == 0 + + db_domain = CustomDomain.get_by(id=domain.id) + assert db_domain.dkim_verified is True + + +def test_custom_domain_validation_validate_dkim_records_partner_domain_fallback_success(): + """Test DKIM validation for partner domains falls back to SL domain if partner DKIM not set.""" + dkim_domain = random_domain() + partner_mx_domain = random_domain() # Partner has MX config but not DKIM + dns_client = InMemoryDNSClient() + + partner_id = get_proton_partner().id + validator = CustomDomainValidation( + dkim_domain, dns_client, partner_domains={partner_id: partner_mx_domain} + ) + + user_domain = random_domain() + domain = create_custom_domain(user_domain) + domain.partner_id = partner_id + Session.commit() + + # Set correct SL DKIM records (fallback) + dns_client.set_cname_record( + f"dkim._domainkey.{user_domain}", f"dkim._domainkey.{dkim_domain}" + ) + dns_client.set_cname_record( + f"dkim02._domainkey.{user_domain}", f"dkim02._domainkey.{dkim_domain}" + ) + dns_client.set_cname_record( + f"dkim03._domainkey.{user_domain}", f"dkim03._domainkey.{dkim_domain}" + ) + + res = validator.validate_dkim_records(domain) + assert len(res) == 0 + + db_domain = CustomDomain.get_by(id=domain.id) + assert db_domain.dkim_verified is True + + # validate_ownership def test_custom_domain_validation_validate_ownership_empty_records_failure(): dns_client = InMemoryDNSClient() @@ -428,6 +496,39 @@ def test_custom_domain_validation_validate_mx_records_success(): assert db_domain.verified is True +def test_custom_domain_validation_validate_mx_records_partner_domain_success(): + """Test MX validation for partner domains with custom MX configuration.""" + dns_client = InMemoryDNSClient() + partner_id = get_proton_partner().id + partner_mx_domain = random_domain() + + validator = CustomDomainValidation( + random_domain(), dns_client, partner_domains={partner_id: partner_mx_domain} + ) + + domain = create_custom_domain(random_domain()) + domain.partner_id = partner_id + Session.commit() + + # Get expected records (should include partner MX + fallback SL MX) + expected_records = validator.get_expected_mx_records(domain) + + # Set MX records with partner MX first (preferred priority) + dns_records = {} + for priority in expected_records: + # Use the recommended record (partner MX) + dns_records[priority] = [expected_records[priority].recommended] + + dns_client.set_mx_records(domain.domain, dns_records) + res = validator.validate_mx_records(domain) + + assert res.success is True + assert len(res.errors) == 0 + + db_domain = CustomDomain.get_by(id=domain.id) + assert db_domain.verified is True + + # validate_spf_records def test_custom_domain_validation_validate_spf_records_empty_failure(): dns_client = InMemoryDNSClient() From f4c4cc50e7ee6b23809e34b58cff62e57598166a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?= Date: Tue, 7 Jul 2026 16:06:22 +0200 Subject: [PATCH 2/2] Fix test --- tests/tasks/test_check_custom_domain.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/tasks/test_check_custom_domain.py b/tests/tasks/test_check_custom_domain.py index b6b1301e5..d2eb0ad68 100644 --- a/tests/tasks/test_check_custom_domain.py +++ b/tests/tasks/test_check_custom_domain.py @@ -1,6 +1,7 @@ import arrow from unittest.mock import patch +from app.db import Session from app.models import CustomDomain from tasks.check_custom_domains import ( check_all_custom_domains, @@ -11,6 +12,7 @@ def test_check_single_custom_domain_increments_failed_checks(flask_client): user = create_new_user() + # Set updated_at to 2 days ago so the increment logic runs custom_domain = CustomDomain.create( user_id=user.id, domain=random_string(), @@ -18,6 +20,10 @@ def test_check_single_custom_domain_increments_failed_checks(flask_client): nb_failed_checks=0, commit=True, ) + # Manually set updated_at to simulate an old update + custom_domain.updated_at = arrow.now().shift(days=-2) + Session.commit() + with patch("tasks.check_custom_domains.get_mx_domains", return_value=[]), patch( "tasks.check_custom_domains.is_mx_equivalent", return_value=False ), patch("tasks.check_custom_domains.send_email_with_rate_control"): @@ -28,13 +34,18 @@ def test_check_single_custom_domain_increments_failed_checks(flask_client): def test_check_single_custom_domain_deactivates_after_threshold(flask_client): user = create_new_user() + # Start at 4 so that after one increment (4->5), 5 > MAX_DOMAIN_CHECKS (4) triggers deactivation custom_domain = CustomDomain.create( user_id=user.id, domain=random_string(), verified=True, - nb_failed_checks=2, + nb_failed_checks=4, commit=True, ) + # Set updated_at to 2 days ago so the increment logic runs + custom_domain.updated_at = arrow.now().shift(days=-2) + Session.commit() + with patch("tasks.check_custom_domains.get_mx_domains", return_value=[]), patch( "tasks.check_custom_domains.is_mx_equivalent", return_value=False ), patch("tasks.check_custom_domains.send_email_with_rate_control"):