feat(api/oauth): switch-to-stateful-authentication

api/app.py

@@ -1,8 +1,9 @@
 import os
 
 from configs.app_configs import DifyConfigs
+from extensions.ext_redis import redis_client
 
-if not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true':
+if not os.environ.get("DEBUG") or os.environ.get("DEBUG", "false").lower() != 'true':
     from gevent import monkey
 
     monkey.patch_all()
@@ -169,6 +170,9 @@ def load_user_from_request(request_from_flask_login):
 
         decoded = PassportService().verify(auth_token)
         user_id = decoded.get('user_id')
+        # check if account id in redis.
+        if not redis_client.get(f"account:{user_id}"):
+            return None
 
         return AccountService.load_user(user_id)
     else:

api/controllers/console/auth/login.py

@@ -5,6 +5,7 @@ from flask_restful import Resource, reqparse
 import services
 from controllers.console import api
 from controllers.console.setup import setup_required
+from extensions.ext_redis import redis_client
 from libs.helper import email
 from libs.password import valid_password
 from services.account_service import AccountService, TenantService
@@ -17,50 +18,57 @@ class LoginApi(Resource):
     def post(self):
         """Authenticate user and login."""
         parser = reqparse.RequestParser()
-        parser.add_argument('email', type=email, required=True, location='json')
-        parser.add_argument('password', type=valid_password, required=True, location='json')
-        parser.add_argument('remember_me', type=bool, required=False, default=False, location='json')
+        parser.add_argument("email", type=email, required=True, location="json")
+        parser.add_argument("password", type=valid_password, required=True, location="json")
+        parser.add_argument("remember_me", type=bool, required=False, default=False, location="json")
         args = parser.parse_args()
 
         # todo: Verify the recaptcha
 
         try:
-            account = AccountService.authenticate(args['email'], args['password'])
+            account = AccountService.authenticate(args["email"], args["password"])
         except services.errors.account.AccountLoginError as e:
-            return {'code': 'unauthorized', 'message': str(e)}, 401
+            return {"code": "unauthorized", "message": str(e)}, 401
 
         # SELF_HOSTED only have one workspace
         tenants = TenantService.get_join_tenants(account)
         if len(tenants) == 0:
-            return {'result': 'fail', 'data': 'workspace not found, please contact system admin to invite you to join in a workspace'}
+            return {
+                "result": "fail",
+                "data": "workspace not found, please contact system admin to invite you to join in a workspace",
+            }
 
         AccountService.update_last_login(account, request)
 
         # todo: return the user info
         token = AccountService.get_account_jwt_token(account)
 
-        return {'result': 'success', 'data': token}
+        # put account id into redis with 30 days tts.
+        redis_client.set(f"account:{account.id}", token, ex=30 * 24 * 60 * 60)
+
+        return {"result": "success", "data": token}
 
 
 class LogoutApi(Resource):
-
     @setup_required
     def get(self):
+        # remove account id and token from redis
+        redis_client.delete(f"account:{flask_login.current_user.id}")
         flask_login.logout_user()
-        return {'result': 'success'}
+        return {"result": "success"}
 
 
 class ResetPasswordApi(Resource):
     @setup_required
     def get(self):
         parser = reqparse.RequestParser()
-        parser.add_argument('email', type=email, required=True, location='json')
+        parser.add_argument("email", type=email, required=True, location="json")
         args = parser.parse_args()
 
         # import mailchimp_transactional as MailchimpTransactional
         # from mailchimp_transactional.api_client import ApiClientError
 
-        account = {'email': args['email']}
+        account = {"email": args["email"]}
         # account = AccountService.get_by_email(args['email'])
         # if account is None:
         #     raise ValueError('Email not found')
@@ -68,21 +76,21 @@ class ResetPasswordApi(Resource):
         # AccountService.update_password(account, new_password)
 
         # todo: Send email
-        MAILCHIMP_API_KEY = current_app.config['MAILCHIMP_TRANSACTIONAL_API_KEY']
+        MAILCHIMP_API_KEY = current_app.config["MAILCHIMP_TRANSACTIONAL_API_KEY"]
         # mailchimp = MailchimpTransactional(MAILCHIMP_API_KEY)
 
         message = {
-            'from_email': 'noreply@example.com',
-            'to': [{'email': account.email}],
-            'subject': 'Reset your Dify password',
-            'html': """
+            "from_email": "noreply@example.com",
+            "to": [{"email": account['email']}],
+            "subject": "Reset your Dify password",
+            "html": """
                 <p>Dear User,</p>
                 <p>The Dify team has generated a new password for you, details as follows:</p> 
                 <p><strong>{new_password}</strong></p>
                 <p>Please change your password to log in as soon as possible.</p>
                 <p>Regards,</p>
                 <p>The Dify Team</p> 
-            """
+            """,
         }
 
         # response = mailchimp.messages.send({
@@ -98,8 +106,8 @@ class ResetPasswordApi(Resource):
         #     # handle error
         #     pass
 
-        return {'result': 'success'}
+        return {"result": "success"}
 
 
-api.add_resource(LoginApi, '/login')
-api.add_resource(LogoutApi, '/logout')
+api.add_resource(LoginApi, "/login")
+api.add_resource(LogoutApi, "/logout")

api/controllers/console/auth/oauth.py

@@ -8,6 +8,7 @@ from flask_restful import Resource
 
 from constants.languages import languages
 from extensions.ext_database import db
+from extensions.ext_redis import redis_client
 from libs.oauth import GitHubOAuth, GoogleOAuth, OAuthUserInfo
 from models.account import Account, AccountStatus
 from services.account_service import AccountService, RegisterService, TenantService
@@ -82,6 +83,9 @@ class OAuthCallback(Resource):
 
         token = AccountService.get_account_jwt_token(account)
 
+        # put account id into redis with 30 days tts.
+        redis_client.set(f"account:{account.id}", token, ex=30 * 24 * 60 * 60)
+
         return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?console_token={token}')
 
 

api/services/__init__.py

@@ -1 +1,3 @@
-import services.errors
+from . import errors
+
+__all__ = ['errors']

api/services/errors/__init__.py

@@ -1,6 +1,29 @@
-__all__ = [
-    'base', 'conversation', 'message', 'index', 'app_model_config', 'account', 'document', 'dataset',
-    'app', 'completion', 'audio', 'file'
-]
+from . import (
+    account,
+    app,
+    app_model_config,
+    audio,
+    base,
+    completion,
+    conversation,
+    dataset,
+    document,
+    file,
+    index,
+    message,
+)
 
-from . import *
+__all__ = [
+    "base",
+    "conversation",
+    "message",
+    "index",
+    "app_model_config",
+    "account",
+    "document",
+    "dataset",
+    "app",
+    "completion",
+    "audio",
+    "file",
+]