From 78b1aabb67662941b6ee99cacda8d86246840d45 Mon Sep 17 00:00:00 2001
From: Joe <1264204425@qq.com>
Date: Tue, 3 Sep 2024 15:53:19 +0800
Subject: [PATCH] feat: add NotAllowedCreateWorkspace and NotAllowedRegister

---
 api/controllers/console/auth/forgot_password.py |  3 ++-
 api/controllers/console/error.py                | 12 ++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/api/controllers/console/auth/forgot_password.py b/api/controllers/console/auth/forgot_password.py
index e3826b6ef6..2b12a77a2b 100644
--- a/api/controllers/console/auth/forgot_password.py
+++ b/api/controllers/console/auth/forgot_password.py
@@ -15,6 +15,7 @@ from controllers.console.auth.error import (
     PasswordMismatchError,
     PasswordResetRateLimitExceededError,
 )
+from controllers.console.error import NotAllowedRegister
 from controllers.console.setup import setup_required
 from extensions.ext_database import db
 from libs.helper import email, get_remote_ip
@@ -37,7 +38,7 @@ class ForgotPasswordSendEmailApi(Resource):
             if dify_config.ALLOW_REGISTER:
                 token = AccountService.send_reset_password_email(email=args["email"])
             else:
-                raise InvalidEmailError()
+                raise NotAllowedRegister()
         elif account:
             try:
                 token = AccountService.send_reset_password_email(account=account, email=args["email"])
diff --git a/api/controllers/console/error.py b/api/controllers/console/error.py
index 1c70ea6c59..3cc7bf9603 100644
--- a/api/controllers/console/error.py
+++ b/api/controllers/console/error.py
@@ -40,3 +40,15 @@ class AlreadyActivateError(BaseHTTPException):
     error_code = "already_activate"
     description = "Auth Token is invalid or account already activated, please check again."
     code = 403
+
+
+class NotAllowedCreateWorkspace(BaseHTTPException):
+    error_code = "unauthorized"
+    description = "Workspace not found, please contact system admin to invite you to join in a workspace."
+    code = 400
+
+
+class NotAllowedRegister(BaseHTTPException):
+    error_code = "unauthorized"
+    description = "Account not found."
+    code = 404