AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-06-03 18:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: Consolidate Flask-Login Authentication Logic (#20235)

Browse Source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2025-05-26 18:22:01 +08:00 committed by GitHub
parent 6f982eb7e4
commit f21e6e03a3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 23 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
api/extensions/ext_login.py
View File

@ -6,10 +6,11 @@ from flask_login import user_loaded_from_request, user_logged_in
from werkzeug.exceptions import NotFound, Unauthorized
import contexts
from configs import dify_config
from dify_app import DifyApp
from extensions.ext_database import db
from libs.passport import PassportService
from models.account import Account
from models.account import Account, Tenant, TenantAccountJoin
from models.model import EndUser
from services.account_service import AccountService
@ -32,6 +33,26 @@ def load_user_from_request(request_from_flask_login):
else:
auth_token = request.args.get("_token")
# Check for admin API key authentication first
if dify_config.ADMIN_API_KEY_ENABLE and auth_header:
admin_api_key = dify_config.ADMIN_API_KEY
if admin_api_key and admin_api_key == auth_token:
workspace_id = request.headers.get("X-WORKSPACE-ID")
if workspace_id:
tenant_account_join = (
db.session.query(Tenant, TenantAccountJoin)
.filter(Tenant.id == workspace_id)
.filter(TenantAccountJoin.tenant_id == Tenant.id)
.filter(TenantAccountJoin.role == "owner")
.one_or_none()
)
if tenant_account_join:
tenant, ta = tenant_account_join
account = db.session.query(Account).filter_by(id=ta.account_id).first()
if account:
account.current_tenant = tenant
return account
if request.blueprint in {"console", "inner_api"}:
if not auth_token:
raise Unauthorized("Invalid Authorization token.")

35
api/libs/login.py
View File

@ -2,14 +2,11 @@ from functools import wraps
from typing import Any
from flask import current_app, g, has_request_context, request
from flask_login import user_logged_in # type: ignore
from flask_login.config import EXEMPT_METHODS # type: ignore
from werkzeug.exceptions import Unauthorized
from werkzeug.local import LocalProxy
from configs import dify_config
from extensions.ext_database import db
from models.account import Account, Tenant, TenantAccountJoin
from models.account import Account
from models.model import EndUser
#: A proxy for the current user. If no user is logged in, this will be an
@ -53,36 +50,6 @@ def login_required(func):
@wraps(func)
def decorated_view(*args, **kwargs):
auth_header = request.headers.get("Authorization")
if dify_config.ADMIN_API_KEY_ENABLE:
if auth_header:
if " " not in auth_header:
raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
auth_scheme, auth_token = auth_header.split(None, 1)
auth_scheme = auth_scheme.lower()
if auth_scheme != "bearer":
raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
admin_api_key = dify_config.ADMIN_API_KEY
if admin_api_key:
if admin_api_key == auth_token:
workspace_id = request.headers.get("X-WORKSPACE-ID")
if workspace_id:
tenant_account_join = (
db.session.query(Tenant, TenantAccountJoin)
.filter(Tenant.id == workspace_id)
.filter(TenantAccountJoin.tenant_id == Tenant.id)
.filter(TenantAccountJoin.role == "owner")
.one_or_none()
)
if tenant_account_join:
tenant, ta = tenant_account_join
account = db.session.query(Account).filter_by(id=ta.account_id).first()
# Login admin
if account:
account.current_tenant = tenant
current_app.login_manager._update_request_context_with_user(account) # type: ignore
user_logged_in.send(current_app._get_current_object(), user=_get_user()) # type: ignore
if request.method in EXEMPT_METHODS or dify_config.LOGIN_DISABLED:
pass
elif not current_user.is_authenticated: