Dify 中的 @login_required 装饰器

本文主要介绍了 Dify 中的 @login_required 装饰器的具体实现过程，用于确保用户在调用视图函数之前已经登录和认证。源码位置：dify\api\libs\login.py。

def login_required(func):
    """
    If you decorate a view with this, it will ensure that the current user is
    logged in and authenticated before calling the actual view. (If they are
    not, it calls the :attr:`LoginManager.unauthorized` callback.) For
    example::

        @app.route('/post')
        @login_required
        def post():
            pass

    If there are only certain times you need to require that your user is
    logged in, you can do so with::

        if not current_user.is_authenticated:
            return current_app.login_manager.unauthorized()

    ...which is essentially the code that this function adds to your views.

    It can be convenient to globally turn off authentication when unit testing.
    To enable this, if the application configuration variable `LOGIN_DISABLED`
    is set to `True`, this decorator will be ignored.

    .. Note ::

        Per `W3 guidelines for CORS preflight requests
        <http://www.w3.org/TR/cors/#cross-origin-request-with-preflight-0>`_,
        HTTP ``OPTIONS`` requests are exempt from login checks.

    :param func: The view function to decorate.
    :type func: function
    """

    @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 = Account.query.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)
                                    user_logged_in.send(current_app._get_current_object(), user=_get_user())
        if request.method in EXEMPT_METHODS or dify_config.LOGIN_DISABLED:
            pass
        elif not current_user.is_authenticated:
            return current_app.login_manager.unauthorized()

        # flask 1.x compatibility
        # current_app.ensure_sync is only available in Flask >= 2.0
        if callable(getattr(current_app, "ensure_sync", None)):
            return current_app.ensure_sync(func)(*args, **kwargs)
        return func(*args, **kwargs)

    return decorated_view

1.文档字符串部分

"""
If you decorate a view with this, it will ensure that the current user is
logged in and authenticated before calling the actual view. (If they are
not, it calls the :attr:`LoginManager.unauthorized` callback.) For
example::

    @app.route('/post')
    @login_required
    def post():
        pass

• 这段文档说明了装饰器的功能：在调用视图函数前确保用户已登录并认证，如果未认证，则调用 LoginManager.unauthorized 回调。
• 提供了一个示例，展示如何将该装饰器应用于 Flask 路由。

2.功能说明

If there are only certain times you need to require that your user is
logged in, you can do so with::

    if not current_user.is_authenticated:
        return current_app.login_manager.unauthorized()

• 说明可以在特定情况下手动检查用户认证状态的方式。

...which is essentially the code that this function adds to your views.

• 说明这个装饰器的功能相当于将相同的认证检查代码添加到视图中。

3.单元测试说明

It can be convenient to globally turn off authentication when unit testing.
To enable this, if the application configuration variable `LOGIN_DISABLED`
is set to `True`, this decorator will be ignored.

• 提及在单元测试期间可以通过设置 LOGIN_DISABLED 为 True 来禁用认证检查，便于测试。

4.CORS 说明

.. Note ::
    Per `W3 guidelines for CORS preflight requests
    <http://www.w3.org/TR/cors/#cross-origin-request-with-preflight-0>`_,
    HTTP ``OPTIONS`` requests are exempt from login checks.

• 指出根据 W3 的 CORS 指南，HTTP OPTIONS 请求不需要进行登录检查。

5.装饰器主体

:param func: The view function to decorate.
:type func: function

• 指定参数 func 是被装饰的视图函数。

@wraps(func)
def decorated_view(*args, **kwargs):

• 使用 functools.wraps 来保持原始视图函数的元数据（如文档字符串和名称）。
• 定义内部函数 decorated_view，接受任意位置和关键字参数。

6.获取 Authorization 头部

auth_header = request.headers.get("Authorization")

• 从请求头中获取 Authorization 字段。

7.检查是否启用 ADMIN_API_KEY

if dify_config.ADMIN_API_KEY_ENABLE:

• 如果配置中启用了 ADMIN_API_KEY，则继续进行后续的身份验证。

8.验证 Authorization 格式

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.")

• 如果 auth_header 存在，检查其格式是否正确。要求格式为 Bearer <api-key>。
• 如果格式不正确，则抛出 Unauthorized 异常。

9.验证 ADMIN_API_KEY

admin_api_key = dify_config.ADMIN_API_KEY
if admin_api_key:
    if admin_api_key == auth_token:

• 获取配置中的 ADMIN_API_KEY，并与从请求中提取的 auth_token 进行比较。

10.获取工作区 ID

workspace_id = request.headers.get("X-WORKSPACE-ID")
if workspace_id:

• 从请求头中获取工作区 ID（X-WORKSPACE-ID）。

11.查询数据库

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()
)

• 查询数据库，检查该工作区 ID 是否存在，并验证当前用户是否为该租户的所有者。

12.登录管理员

if tenant_account_join:
    tenant, ta = tenant_account_join
    account = Account.query.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)
        user_logged_in.send(current_app._get_current_object(), user=_get_user())

• 如果找到了租户账户关系，获取账户信息并更新请求上下文，以登录该用户，并发送用户登录的信号。

13.处理请求方法

if request.method in EXEMPT_METHODS or dify_config.LOGIN_DISABLED:
    pass

• 检查请求方法是否在被豁免的方法列表中，或是否禁用登录。如果是，则跳过认证检查。

14.检查用户是否认证

elif not current_user.is_authenticated:
    return current_app.login_manager.unauthorized()

• 如果用户未认证，则返回未授权响应。

15.Flask 兼容性

if callable(getattr(current_app, "ensure_sync", None)):
    return current_app.ensure_sync(func)(*args, **kwargs)

• 检查 Flask 版本是否支持 ensure_sync，如果支持，则异步调用装饰的视图函数。

16.返回原始视图函数

return func(*args, **kwargs)

• 如果用户已认证，直接调用并返回原始的视图函数。

总结：这个装饰器用于确保用户在访问特定视图时已登录，并且可以处理使用 API 密钥的认证，支持灵活的配置以方便测试和 CORS 处理。