feat: update raise error

This commit is contained in:
Joe 2024-09-26 15:49:37 +08:00
parent 9e4ee2beb1
commit b249f2b9f9
3 changed files with 17 additions and 13 deletions

View File

@ -49,3 +49,9 @@ class EmailPasswordLoginLimitError(BaseHTTPException):
"Too many incorrect password attempts. Please verify your identity with the email code to complete login." "Too many incorrect password attempts. Please verify your identity with the email code to complete login."
) )
code = 429 code = 429
class EmailCodeLoginRateLimitExceededError(BaseHTTPException):
error_code = "email_code_login_rate_limit_exceeded"
description = "Too many login emails have been sent. Please try again in 5 minutes."
code = 429

View File

@ -43,14 +43,10 @@ class ForgotPasswordSendEmailApi(Resource):
) )
else: else:
raise NotAllowedRegister() raise NotAllowedRegister()
elif account: else:
try: token = AccountService.send_reset_password_email(
token = AccountService.send_reset_password_email( account=account, email=args["email"], language=args["language"] or "en-US"
account=account, email=args["email"], language=args["language"] or "en-US" )
)
except RateLimitExceededError:
logging.warning(f"Rate limit exceeded for email: {args['email']}")
raise PasswordResetRateLimitExceededError()
return {"result": "success", "data": token} return {"result": "success", "data": token}

View File

@ -12,7 +12,6 @@ from werkzeug.exceptions import Unauthorized
from configs import dify_config from configs import dify_config
from constants.languages import language_timezone_mapping, languages from constants.languages import language_timezone_mapping, languages
from controllers.console.auth.error import EmailPasswordLoginLimitError
from events.tenant_event import tenant_was_created from events.tenant_event import tenant_was_created
from extensions.ext_redis import redis_client from extensions.ext_redis import redis_client
from libs.helper import RateLimiter, TokenManager from libs.helper import RateLimiter, TokenManager
@ -34,7 +33,6 @@ from services.errors.account import (
LinkAccountIntegrateError, LinkAccountIntegrateError,
MemberNotInTenantError, MemberNotInTenantError,
NoPermissionError, NoPermissionError,
RateLimitExceededError,
RoleAlreadyAssignedError, RoleAlreadyAssignedError,
TenantNotFoundError, TenantNotFoundError,
) )
@ -45,7 +43,7 @@ from tasks.mail_reset_password_task import send_reset_password_mail_task
class AccountService: class AccountService:
reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=5, time_window=60 * 60) reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=1, time_window=60 * 1)
email_code_login_rate_limiter = RateLimiter( email_code_login_rate_limiter = RateLimiter(
prefix="email_code_login_rate_limit", max_attempts=1, time_window=60 * 1 prefix="email_code_login_rate_limit", max_attempts=1, time_window=60 * 1
) )
@ -263,7 +261,9 @@ class AccountService:
account_language = account.interface_language if account else language account_language = account.interface_language if account else language
if cls.reset_password_rate_limiter.is_rate_limited(account_email): if cls.reset_password_rate_limiter.is_rate_limited(account_email):
raise RateLimitExceededError(f"Rate limit exceeded for email: {account_email}. Please try again later.") from controllers.console.auth.error import PasswordResetRateLimitExceededError
raise PasswordResetRateLimitExceededError()
code = "".join([str(random.randint(0, 9)) for _ in range(6)]) code = "".join([str(random.randint(0, 9)) for _ in range(6)])
token = TokenManager.generate_token( token = TokenManager.generate_token(
@ -290,7 +290,9 @@ class AccountService:
cls, account: Optional[Account] = None, email: Optional[str] = None, language: Optional[str] = "en-US" cls, account: Optional[Account] = None, email: Optional[str] = None, language: Optional[str] = "en-US"
): ):
if cls.email_code_login_rate_limiter.is_rate_limited(email): if cls.email_code_login_rate_limiter.is_rate_limited(email):
raise EmailPasswordLoginLimitError() from controllers.console.auth.error import EmailCodeLoginRateLimitExceededError
raise EmailCodeLoginRateLimitExceededError()
code = "".join([str(random.randint(0, 9)) for _ in range(6)]) code = "".join([str(random.randint(0, 9)) for _ in range(6)])
token = TokenManager.generate_token( token = TokenManager.generate_token(