get_object_from_key: Accept multiple allowed Confirmation types.

This allows making check_prereg_key significantly cleaner.

confirmation/models.py

@@ -4,7 +4,7 @@ __revision__ = "$Id: models.py 28 2009-10-22 15:03:02Z jarek.zgoda $"
 import datetime
 import secrets
 from base64 import b32encode
-from typing import Mapping, Optional, Union
+from typing import List, Mapping, Optional, Union
 from urllib.parse import urljoin
 
 from django.conf import settings
@@ -58,14 +58,14 @@ ConfirmationObjT = Union[MultiuseInvite, PreregistrationUser, EmailChangeStatus]
 
 
 def get_object_from_key(
-    confirmation_key: str, confirmation_type: int, activate_object: bool = True
+    confirmation_key: str, confirmation_types: List[int], activate_object: bool = True
 ) -> ConfirmationObjT:
     # Confirmation keys used to be 40 characters
     if len(confirmation_key) not in (24, 40):
         raise ConfirmationKeyException(ConfirmationKeyException.WRONG_LENGTH)
     try:
         confirmation = Confirmation.objects.get(
-            confirmation_key=confirmation_key, type=confirmation_type
+            confirmation_key=confirmation_key, type__in=confirmation_types
         )
     except Confirmation.DoesNotExist:
         raise ConfirmationKeyException(ConfirmationKeyException.DOES_NOT_EXIST)

zerver/tests/test_signup.py

@@ -855,7 +855,7 @@ class LoginTest(ZulipTestCase):
         with queries_captured() as queries, cache_tries_captured() as cache_tries:
             self.register(self.nonreg_email("test"), "test")
         # Ensure the number of queries we make is not O(streams)
-        self.assert_length(queries, 91)
+        self.assert_length(queries, 89)
 
         # We can probably avoid a couple cache hits here, but there doesn't
         # seem to be any O(N) behavior.  Some of the cache hits are related
@@ -2008,7 +2008,7 @@ so we didn't send them an invitation. We did send invitations to everyone else!"
 
         # Mainly a test of get_object_from_key, rather than of the invitation pathway
         with self.assertRaises(ConfirmationKeyException) as cm:
-            get_object_from_key(registration_key, Confirmation.INVITATION)
+            get_object_from_key(registration_key, [Confirmation.INVITATION])
         self.assertEqual(cm.exception.error_type, ConfirmationKeyException.DOES_NOT_EXIST)
 
         # Verify that using the wrong type doesn't work in the main confirm code path

zerver/views/auth.py

@@ -186,7 +186,7 @@ def maybe_send_to_registration(
     if multiuse_object_key:
         from_multiuse_invite = True
         try:
-            multiuse_obj = get_object_from_key(multiuse_object_key, Confirmation.MULTIUSE_INVITE)
+            multiuse_obj = get_object_from_key(multiuse_object_key, [Confirmation.MULTIUSE_INVITE])
         except ConfirmationKeyException:
             return render(request, "zerver/confirmation_link_expired_error.html", status=404)
 

zerver/views/realm.py

@@ -316,7 +316,7 @@ def check_subdomain_available(request: HttpRequest, subdomain: str) -> HttpRespo
 
 def realm_reactivation(request: HttpRequest, confirmation_key: str) -> HttpResponse:
     try:
-        realm = get_object_from_key(confirmation_key, Confirmation.REALM_REACTIVATION)
+        realm = get_object_from_key(confirmation_key, [Confirmation.REALM_REACTIVATION])
     except ConfirmationKeyException:
         return render(request, "zerver/realm_reactivation_link_error.html")
     do_reactivate_realm(realm)

zerver/views/registration.py

@@ -118,32 +118,22 @@ def check_prereg_key(
     Checks if the Confirmation key is valid, returning the PreregistrationUser object in case of success
     and an appropriate error page otherwise.
     """
-    try:
+    confirmation_types = [
-        confirmation: Optional[Confirmation] = Confirmation.objects.get(
-            confirmation_key=confirmation_key
-        )
-    except Confirmation.DoesNotExist:
-        confirmation = None
-
-    if confirmation is None or confirmation.type not in [
         Confirmation.USER_REGISTRATION,
         Confirmation.INVITATION,
         Confirmation.REALM_CREATION,
-    ]:
+    ]
-        return render_confirmation_key_error(
-            request, ConfirmationKeyException(ConfirmationKeyException.DOES_NOT_EXIST)
-        )
-
-    prereg_user = confirmation.content_object
-    assert prereg_user is not None
-    if prereg_user.status == confirmation_settings.STATUS_REVOKED:
-        return render(request, "zerver/confirmation_link_expired_error.html", status=404)
 
     try:
-        get_object_from_key(confirmation_key, confirmation.type, activate_object=False)
+        prereg_user = get_object_from_key(
+            confirmation_key, confirmation_types, activate_object=False
+        )
     except ConfirmationKeyException as exception:
         return render_confirmation_key_error(request, exception)
 
+    if prereg_user.status == confirmation_settings.STATUS_REVOKED:
+        return render(request, "zerver/confirmation_link_expired_error.html", status=404)
+
     return prereg_user
 
 
@@ -730,7 +720,7 @@ def accounts_home(
 def accounts_home_from_multiuse_invite(request: HttpRequest, confirmation_key: str) -> HttpResponse:
     multiuse_object = None
     try:
-        multiuse_object = get_object_from_key(confirmation_key, Confirmation.MULTIUSE_INVITE)
+        multiuse_object = get_object_from_key(confirmation_key, [Confirmation.MULTIUSE_INVITE])
         # Required for OAuth 2
     except ConfirmationKeyException as exception:
         realm = get_realm_from_request(request)

zerver/views/unsubscribe.py

@@ -18,7 +18,7 @@ def process_unsubscribe(
     unsubscribe_function: Callable[[UserProfile], None],
 ) -> HttpResponse:
     try:
-        user_profile = get_object_from_key(confirmation_key, Confirmation.UNSUBSCRIBE)
+        user_profile = get_object_from_key(confirmation_key, [Confirmation.UNSUBSCRIBE])
     except ConfirmationKeyException:
         return render(request, "zerver/unsubscribe_link_error.html")
 

zerver/views/user_settings.py

@@ -51,7 +51,7 @@ AVATAR_CHANGES_DISABLED_ERROR = gettext_lazy("Avatar changes are disabled in thi
 
 def confirm_email_change(request: HttpRequest, confirmation_key: str) -> HttpResponse:
     try:
-        email_change_object = get_object_from_key(confirmation_key, Confirmation.EMAIL_CHANGE)
+        email_change_object = get_object_from_key(confirmation_key, [Confirmation.EMAIL_CHANGE])
     except ConfirmationKeyException as exception:
         return render_confirmation_key_error(request, exception)