From b6a0e530368c1f66d9355d7cd1819e282df30d97 Mon Sep 17 00:00:00 2001
From: Jordan Robinson <hello@jordanrobinson.org>
Date: Thu, 27 Nov 2025 23:13:50 +0000
Subject: [PATCH] update domain check logic

#75
---
 main/tests/test_users.py | 21 ++++++++++++++++++++-
 main/views/general.py    | 29 ++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/main/tests/test_users.py b/main/tests/test_users.py
index 18eff47a92bcd1e7de525113e46593378e8fe615..c254bbf94c5ef34bee6d4f8ce0e98207953a2ca9 100644
--- a/main/tests/test_users.py
+++ b/main/tests/test_users.py
@@ -1,3 +1,4 @@
+from django.conf import settings
 from django.test import TestCase
 from django.urls import reverse
 
@@ -282,7 +283,7 @@ class UserDomainCheckTestCase(TestCase):
             username="alice", custom_domain="example.com"
         )
 
-    def test_domain_exists(self):
+    def test_custom_domain_exists(self):
         response = self.client.get(reverse("domain_check") + "?domain=example.com")
         self.assertEqual(response.status_code, 200)
 
@@ -290,6 +291,24 @@ class UserDomainCheckTestCase(TestCase):
         response = self.client.get(reverse("domain_check") + "?domain=randomdomain.com")
         self.assertEqual(response.status_code, 403)
 
+    def test_canonical_host(self):
+        response = self.client.get(
+            reverse("domain_check") + "?domain=" + settings.CANONICAL_HOST
+        )
+        self.assertEqual(response.status_code, 200)
+
+    def test_subdomain_with_existing_user(self):
+        response = self.client.get(
+            reverse("domain_check") + f"?domain=alice.{settings.CANONICAL_HOST}"
+        )
+        self.assertEqual(response.status_code, 200)
+
+    def test_subdomain_with_nonexistent_user(self):
+        response = self.client.get(
+            reverse("domain_check") + f"?domain=bob.{settings.CANONICAL_HOST}"
+        )
+        self.assertEqual(response.status_code, 403)
+
 class UserMarkdownLinkOnPaste(TestCase):
     def setUp(self):
         self.user = models.User.objects.create(username="alice")
diff --git a/main/views/general.py b/main/views/general.py
index 63a7792ab707b8ea6a474e870d36a580207b2578..3db4a4361fae236f267b87791d8e856b8c9c4fcd 100644
--- a/main/views/general.py
+++ b/main/views/general.py
@@ -280,15 +280,34 @@ class PostList(LoginRequiredMixin, ListView):
 
 def domain_check(request):
     """
-    This view returns 200 if domain given exists as custom domain in any
-    user account.
+    This view returns 200 if the domain given exists as:
+    * canonical host (main domain)
+    * subdomain of canonical host of existing user/blog
+    * custom domain of existing user/blog
     """
     url = request.GET.get("domain")
     if not url:
         raise PermissionDenied()
-    if not models.User.objects.filter(custom_domain=url).exists():
-        raise PermissionDenied()
-    return HttpResponse()
+    
+    # allow canonical host
+    if url == settings.CANONICAL_HOST:
+        return HttpResponse()
+
+    # allow user blog subdomains
+    host_parts = url.split(".")
+    canonical_parts = settings.CANONICAL_HOST.split(".")
+    if (
+        len(host_parts) == len(canonical_parts) + 1
+        and ".".join(host_parts[1:]) == settings.CANONICAL_HOST
+    ):
+        subdomain = host_parts[0]
+        if models.User.objects.filter(username=subdomain).exists():
+            return HttpResponse()
+
+    # allow custom domains
+    if models.User.objects.filter(custom_domain=url).exists():
+        return HttpResponse()
+    raise PermissionDenied()
 
 
 class Logout(DjLogoutView):