api/extensions/ext_login.py
Rama main · réplica de lectura del repositorio original.
import json
import logging
from collections.abc import Callable
from typing import assert_never, cast, override
import flask_login
from flask import Flask, Request, Response, current_app, request
from flask_login import user_loaded_from_request, user_logged_in
from sqlalchemy import select
from sqlalchemy.orm import Session
from werkzeug.exceptions import NotFound, Unauthorized
from configs import dify_config
from constants import HEADER_NAME_APP_CODE
from core.db.session_factory import session_factory
from core.logging.context import set_identity_context
from dify_app import DifyApp
from libs.passport import PassportService
from libs.token import (
extract_access_token,
extract_console_cookie_token,
extract_webapp_passport,
is_admin_api_key_request,
)
from models import Account, Tenant, TenantAccountJoin
from models.enums import EndUserType
from models.model import AppMCPServer, EndUser
logger = logging.getLogger(__name__)
type LoginUser = Account | EndUser
type AccountLoader = Callable[[str], Account | None]
ACCOUNT_ONLY_ADMISSION_ATTR = "__account_only_admission__"
def bind_account_loader(app: Flask, loader: AccountLoader) -> None:
app.extensions["account_identity_loader"] = loader
def bind_account_only_loader(app: Flask, loader: AccountLoader) -> None:
"""Bind identity lookup for explicitly account-scoped Console resources."""
app.extensions["account_only_identity_loader"] = loader
def _console_account_loader() -> AccountLoader:
"""Select the account loader from the registered view's admission metadata.
Requires a Flask request context. Only explicitly marked resource handlers
use account-only lookup; missing views or metadata retain workspace admission.
Inherited handlers follow the resource class's method resolution order.
"""
view = current_app.view_functions.get(request.endpoint or "")
resource = view.__dict__.get("view_class") if view is not None else None
if not isinstance(resource, type):
return load_account
method = request.method.lower()
handler = next((base.__dict__[method] for base in resource.__mro__ if method in base.__dict__), None)
if handler is not None and handler.__dict__.get(ACCOUNT_ONLY_ADMISSION_ATTR, False):
loader = current_app.extensions["account_only_identity_loader"]
if not callable(loader):
raise TypeError("account_only_identity_loader must be callable")
return loader
return load_account
def load_account(account_id: str) -> Account | None:
return cast(AccountLoader, current_app.extensions["account_identity_loader"])(account_id)
class DifyLoginManager(flask_login.LoginManager):
"""Project-specific Flask-Login manager with a stable unauthorized contract.
Dify registers `unauthorized_handler` below to always return a JSON `Response`.
Overriding this method lets callers rely on that narrower return type instead of
Flask-Login's broader callback contract.
"""
@override
def unauthorized(self) -> Response:
"""Return the registered unauthorized handler result as a Flask `Response`."""
return cast(Response, super().unauthorized())
def load_user_from_request_context(self) -> None:
"""Populate Flask-Login's request-local user cache for the current request."""
self._load_user()
login_manager = DifyLoginManager()
# Flask-Login configuration
@login_manager.request_loader
def load_user_from_request(request_from_flask_login: Request) -> LoginUser | None:
"""Load user based on the request."""
with session_factory.create_session() as session:
return _load_user_from_request(request_from_flask_login, session)
def _load_user_from_request(request_from_flask_login: Request, session: Session) -> LoginUser | None:
"""Load user based on the request using an explicit database session."""
del request_from_flask_login
# Skip authentication for documentation endpoints
if dify_config.SWAGGER_UI_ENABLED and request.path.endswith((dify_config.SWAGGER_UI_PATH, "/swagger.json")):
return None
auth_token = extract_access_token(request)
# Check for admin API key authentication first
if is_admin_api_key_request(request):
workspace_id = request.headers.get("X-WORKSPACE-ID")
if workspace_id:
tenant_account_join = session.execute(
select(Tenant, TenantAccountJoin).where(
Tenant.id == workspace_id,
TenantAccountJoin.tenant_id == Tenant.id,
TenantAccountJoin.role == "owner",
)
).one_or_none()
if tenant_account_join:
tenant, ta = tenant_account_join
account = session.scalar(select(Account).where(Account.id == ta.account_id))
if account:
account.set_current_tenant_with_session(tenant, session=session)
return account
if request.blueprint in {"console", "inner_api"}:
if not auth_token:
raise Unauthorized("Invalid Authorization token.")
decoded = PassportService().verify(auth_token)
user_id = decoded.get("user_id")
source = decoded.get("token_source")
if source:
raise Unauthorized("Invalid Authorization token.")
if not user_id:
raise Unauthorized("Invalid Authorization token.")
logged_in_account = _console_account_loader()(user_id)
return logged_in_account
elif request.blueprint == "openapi":
# Account-branch device-flow approval routes (approve / deny /
# approval-context) sit under @login_required and authenticate via
# the console session cookie. Cookie-only on purpose — bearer
# tokens (dfoa_/dfoe_) live on the Authorization header and are
# validated by AppPipeline, not flask-login.
cookie_token = extract_console_cookie_token(request)
if not cookie_token:
return None
try:
decoded = PassportService().verify(cookie_token)
except Exception:
return None
user_id = decoded.get("user_id")
source = decoded.get("token_source")
if source or not user_id:
return None
return load_account(user_id)
elif request.blueprint == "web":
app_code = request.headers.get(HEADER_NAME_APP_CODE)
webapp_token = extract_webapp_passport(app_code, request) if app_code else None
if webapp_token:
decoded = PassportService().verify(webapp_token)
end_user_id = decoded.get("end_user_id")
if not end_user_id:
raise Unauthorized("Invalid Authorization token.")
end_user = session.scalar(select(EndUser).where(EndUser.id == end_user_id))
if not end_user:
raise NotFound("End user not found.")
return end_user
else:
if not auth_token:
raise Unauthorized("Invalid Authorization token.")
decoded = PassportService().verify(auth_token)
end_user_id = decoded.get("end_user_id")
if end_user_id:
end_user = session.scalar(select(EndUser).where(EndUser.id == end_user_id))
if not end_user:
raise NotFound("End user not found.")
return end_user
else:
raise Unauthorized("Invalid Authorization token for web API.")
elif request.blueprint == "mcp":
server_code = request.view_args.get("server_code") if request.view_args else None
if not server_code:
raise Unauthorized("Invalid Authorization token.")
app_mcp_server = session.scalar(select(AppMCPServer).where(AppMCPServer.server_code == server_code).limit(1))
if not app_mcp_server:
raise NotFound("App MCP server not found.")
end_user = session.scalar(
select(EndUser).where(EndUser.session_id == app_mcp_server.id, EndUser.type == EndUserType.MCP).limit(1)
)
if not end_user:
raise NotFound("End user not found.")
return end_user
return None
@user_logged_in.connect
@user_loaded_from_request.connect
def on_user_logged_in(_sender: object, user: LoginUser) -> None:
"""Snapshot authenticated identity into the side-effect-free logging context.
Note: AccountIdentityGateway.load_user will populate user.current_tenant_id
through the load_user method, which calls account.set_tenant_id_with_session().
"""
set_identity_context()
try:
match user:
case Account():
set_identity_context(tenant_id=user.current_tenant_id, user_id=user.id, user_type="account")
case EndUser():
set_identity_context(tenant_id=user.tenant_id, user_id=user.id, user_type=user.type or "end_user")
case _ as unreachable:
assert_never(unreachable)
except Exception:
# Logging enrichment must never make authentication fail.
logger.exception("Failed to set logging identity context")
return
@login_manager.unauthorized_handler
def unauthorized_handler() -> Response:
"""Handle unauthorized requests."""
# Keep this as a concrete `Response`; `DifyLoginManager.unauthorized()` narrows
# Flask-Login's callback contract based on this override.
return Response(
json.dumps({"code": "unauthorized", "message": "Unauthorized.", "status": 401}),
status=401,
content_type="application/json",
)
def init_app(app: DifyApp) -> None:
login_manager.init_app(app)