From 7c5daac8b6e70793a83be494b96753e0ee9f89a9 Mon Sep 17 00:00:00 2001 From: Mateusz Mandera Date: Thu, 2 Dec 2021 16:34:05 +0100 Subject: [PATCH] get_object_from_key: Accept multiple allowed Confirmation types. This allows making check_prereg_key significantly cleaner. --- confirmation/models.py | 6 +++--- zerver/tests/test_signup.py | 4 ++-- zerver/views/auth.py | 2 +- zerver/views/realm.py | 2 +- zerver/views/registration.py | 28 +++++++++------------------- zerver/views/unsubscribe.py | 2 +- zerver/views/user_settings.py | 2 +- 7 files changed, 18 insertions(+), 28 deletions(-) diff --git a/confirmation/models.py b/confirmation/models.py index 1640201488..f778bbf0c5 100644 --- a/confirmation/models.py +++ b/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) diff --git a/zerver/tests/test_signup.py b/zerver/tests/test_signup.py index 6e9d220fc9..d019420948 100644 --- a/zerver/tests/test_signup.py +++ b/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 diff --git a/zerver/views/auth.py b/zerver/views/auth.py index 853237d045..aed7d0d8a2 100644 --- a/zerver/views/auth.py +++ b/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) diff --git a/zerver/views/realm.py b/zerver/views/realm.py index 073d69b47f..8dbff794b7 100644 --- a/zerver/views/realm.py +++ b/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) diff --git a/zerver/views/registration.py b/zerver/views/registration.py index 9c8505907d..ea5edfd6c4 100644 --- a/zerver/views/registration.py +++ b/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: Optional[Confirmation] = Confirmation.objects.get( - confirmation_key=confirmation_key - ) - except Confirmation.DoesNotExist: - confirmation = None - - if confirmation is None or confirmation.type not in [ + confirmation_types = [ 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) diff --git a/zerver/views/unsubscribe.py b/zerver/views/unsubscribe.py index ecede8baf1..80da41a522 100644 --- a/zerver/views/unsubscribe.py +++ b/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") diff --git a/zerver/views/user_settings.py b/zerver/views/user_settings.py index f5feac76dc..db74f15583 100644 --- a/zerver/views/user_settings.py +++ b/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)