perf: 优化忘记密码

apps/authentication/api/password.py

@@ -1,3 +1,5 @@
+import time
+
 from django.core.cache import cache
 from django.http import HttpResponseRedirect
 from django.shortcuts import reverse
@@ -7,7 +9,7 @@ from rest_framework.generics import CreateAPIView
 from rest_framework.permissions import AllowAny
 from rest_framework.response import Response
 
-from authentication.errors import PasswordInvalid
+from authentication.errors import PasswordInvalid, IntervalTooShort
 from authentication.mixins import AuthMixin
 from authentication.mixins import authenticate
 from authentication.serializers import (
@@ -38,18 +40,18 @@ class UserResetPasswordSendCodeApi(CreateAPIView):
             return None, err_msg
         return user, None
 
-    def create(self, request, *args, **kwargs):
-        token = request.GET.get('token')
-        userinfo = cache.get(token)
-        if not userinfo:
-            return HttpResponseRedirect(reverse('authentication:forgot-previewing'))
+    @staticmethod
+    def safe_send_code(token, code, target, form_type, content):
+        token_sent_key = '{}_send_at'.format(token)
+        token_send_at = cache.get(token_sent_key, 0)
+        if token_send_at:
+            raise IntervalTooShort(60)
+        SendAndVerifyCodeUtil(target, code, backend=form_type, **content).gen_and_send_async()
+        cache.set(token_sent_key, int(time.time()), 60)
 
-        serializer = self.get_serializer(data=request.data)
-        serializer.is_valid(raise_exception=True)
-        username = userinfo.get('username')
+    def prepare_code_data(self, user_info, serializer):
+        username = user_info.get('username')
         form_type = serializer.validated_data['form_type']
-        code = random_string(6, lower=False, upper=False)
-        other_args = {}
 
         target = serializer.validated_data[form_type]
         if form_type == 'sms':
@@ -59,15 +61,30 @@ class UserResetPasswordSendCodeApi(CreateAPIView):
             query_key = form_type
         user, err = self.is_valid_user(username=username, **{query_key: target})
         if not user:
-            return Response({'error': err}, status=400)
+            raise ValueError(err)
 
+        code = random_string(6, lower=False, upper=False)
         subject = '%s: %s' % (get_login_title(), _('Forgot password'))
         context = {
             'user': user, 'title': subject, 'code': code,
         }
         message = render_to_string('authentication/_msg_reset_password_code.html', context)
-        other_args['subject'], other_args['message'] = subject, message
-        SendAndVerifyCodeUtil(target, code, backend=form_type, **other_args).gen_and_send_async()
+        content = {'subject': subject, 'message': message}
+        return code, target, form_type, content
+
+    def create(self, request, *args, **kwargs):
+        token = request.GET.get('token')
+        user_info = cache.get(token)
+        if not user_info:
+            return HttpResponseRedirect(reverse('authentication:forgot-previewing'))
+
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        try:
+            code, target, form_type, content = self.prepare_code_data(user_info, serializer)
+        except ValueError as e:
+            return Response({'error': str(e)}, status=400)
+        self.safe_send_code(token, code, target, form_type, content)
         return Response({'data': 'ok'}, status=200)
 
 

apps/authentication/errors/mfa.py

@@ -36,3 +36,11 @@ class FeiShuNotBound(JMSException):
 class PasswordInvalid(JMSException):
     default_code = 'passwd_invalid'
     default_detail = _('Your password is invalid')
+
+
+class IntervalTooShort(JMSException):
+    default_code = 'interval_too_short'
+    default_detail = _('Please wait for %s seconds before retry')
+
+    def __init__(self, interval, *args, **kwargs):
+        super().__init__(detail=self.default_detail % interval, *args, **kwargs)

apps/common/exceptions.py

@@ -1,8 +1,8 @@
 # -*- coding: utf-8 -*-
 #
 from django.utils.translation import gettext_lazy as _
-from rest_framework.exceptions import APIException
 from rest_framework import status
+from rest_framework.exceptions import APIException
 
 
 class JMSException(APIException):

apps/common/utils/verify_code.py

@@ -26,7 +26,6 @@ class SendAndVerifyCodeUtil(object):
         self.target = target
         self.backend = backend
         self.key = key or self.KEY_TMPL.format(target)
-        self.verify_key = self.key + '_verify'
         self.timeout = settings.VERIFY_CODE_TTL if timeout is None else timeout
         self.other_args = kwargs
 
@@ -48,11 +47,6 @@ class SendAndVerifyCodeUtil(object):
             raise
 
     def verify(self, code):
-        times = cache.get(self.verify_key, 0)
-        if times >= 3:
-            self.__clear()
-            raise CodeExpired
-        cache.set(self.verify_key, times + 1, timeout=self.timeout)
         right = cache.get(self.key)
         if not right:
             raise CodeExpired
@@ -65,7 +59,6 @@ class SendAndVerifyCodeUtil(object):
 
     def __clear(self):
         cache.delete(self.key)
-        cache.delete(self.verify_key)
 
     def __ttl(self):
         return cache.ttl(self.key)

apps/users/models/user.py

@@ -606,7 +606,8 @@ class TokenMixin:
 
     def generate_reset_token(self):
         token = random_string(50)
-        self.set_cache(token)
+        key = self.CACHE_KEY_USER_RESET_PASSWORD_PREFIX.format(token)
+        cache.set(key, {'id': self.id, 'email': self.email}, 3600)
         return token
 
     @classmethod
@@ -626,10 +627,6 @@ class TokenMixin:
             logger.error(e, exc_info=True)
             return None
 
-    def set_cache(self, token):
-        key = self.CACHE_KEY_USER_RESET_PASSWORD_PREFIX.format(token)
-        cache.set(key, {'id': self.id, 'email': self.email}, 3600)
-
     @classmethod
     def expired_reset_password_token(cls, token):
         key = cls.CACHE_KEY_USER_RESET_PASSWORD_PREFIX.format(token)

apps/users/views/profile/reset.py

@@ -2,6 +2,8 @@
 
 from __future__ import unicode_literals
 
+import time
+
 from django.conf import settings
 from django.core.cache import cache
 from django.shortcuts import redirect, reverse
@@ -9,6 +11,7 @@ from django.urls import reverse_lazy
 from django.utils.translation import gettext as _
 from django.views.generic import FormView, RedirectView
 
+from authentication.errors import IntervalTooShort
 from common.utils import FlashMessageUtil, get_object_or_none, random_string
 from common.utils.verify_code import SendAndVerifyCodeUtil
 from users.notifications import ResetPasswordSuccessMsg
@@ -37,6 +40,20 @@ class UserForgotPasswordPreviewingView(FormView):
     def get_redirect_url(token):
         return reverse('authentication:forgot-password') + '?token=%s' % token
 
+    @staticmethod
+    def generate_previewing_token(user):
+        sent_ttl = 60
+        token_sent_at_key = '{}_send_at'.format(user.username)
+        token_sent_at = cache.get(token_sent_at_key, 0)
+
+        if token_sent_at:
+            raise IntervalTooShort(sent_ttl)
+        token = random_string(36)
+        user_info = {'username': user.username, 'phone': user.phone, 'email': user.email}
+        cache.set(token, user_info, 5 * 60)
+        cache.set(token_sent_at_key, time.time(), sent_ttl)
+        return token
+
     def form_valid(self, form):
         username = form.cleaned_data['username']
         user = get_object_or_none(User, username=username)
@@ -49,9 +66,11 @@ class UserForgotPasswordPreviewingView(FormView):
             form.add_error('username', error)
             return super().form_invalid(form)
 
-        token = random_string(36)
-        user_map = {'username': user.username, 'phone': user.phone, 'email': user.email}
-        cache.set(token, user_map, 5 * 60)
+        try:
+            token = self.generate_previewing_token(user)
+        except IntervalTooShort as e:
+            form.add_error('username', e)
+            return super().form_invalid(form)
         return redirect(self.get_redirect_url(token))
 
 
@@ -103,13 +122,25 @@ class UserForgotPasswordView(FormView):
         reset_password_url = reverse('authentication:reset-password')
         return reset_password_url + query_params
 
+    @staticmethod
+    def safe_verify_code(token, target, form_type, code):
+        token_verified_key = '{}_verified'.format(token)
+        token_verified_times = cache.get(token_verified_key, 0)
+
+        if token_verified_times >= 3:
+            cache.delete(token)
+            raise ValueError('Verification code has been used more than 3 times, please re-verify')
+        cache.set(token_verified_key, token_verified_times + 1, 5 * 60)
+        sender_util = SendAndVerifyCodeUtil(target, backend=form_type)
+        return sender_util.verify(code)
+
     def form_valid(self, form):
         token = self.request.GET.get('token')
-        userinfo = cache.get(token)
-        if not userinfo:
+        user_info = cache.get(token)
+        if not user_info:
             return redirect(self.get_redirect_url(return_previewing=True))
 
-        username = userinfo.get('username')
+        username = user_info.get('username')
         form_type = form.cleaned_data['form_type']
         target = form.cleaned_data[form_type]
         code = form.cleaned_data['code']
@@ -120,8 +151,9 @@ class UserForgotPasswordView(FormView):
             target = target.lstrip('+')
 
         try:
-            sender_util = SendAndVerifyCodeUtil(target, backend=form_type)
-            sender_util.verify(code)
+            self.safe_verify_code(token, target, form_type, code)
+        except ValueError as e:
+            return redirect(self.get_redirect_url(return_previewing=True))
         except Exception as e:
             form.add_error('code', str(e))
             return super().form_invalid(form)