diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index caee32d3174..4b91a2320a3 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -15,7 +15,12 @@ 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 +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 @@ -66,24 +71,22 @@ def _load_user_from_request(request_from_flask_login: Request, session: Session) auth_token = extract_access_token(request) # Check for admin API key authentication first - if dify_config.ADMIN_API_KEY_ENABLE and auth_token: - 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 = 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 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: diff --git a/api/libs/token.py b/api/libs/token.py index ce868dc3d8a..1b28b60a538 100644 --- a/api/libs/token.py +++ b/api/libs/token.py @@ -1,3 +1,4 @@ +import hmac import logging import re from datetime import UTC, datetime, timedelta @@ -84,6 +85,19 @@ def extract_access_token(request: Request) -> str | None: return extract_console_cookie_token(request) or _try_extract_from_header(request) +def is_admin_api_key_request(request: Request) -> bool: + """Return whether the request carries the configured admin API key as a bearer token. + + Admin API key authentication is header-only so an unrelated console session cookie + cannot shadow the bearer token used by server-to-server clients. + """ + admin_api_key = dify_config.ADMIN_API_KEY + bearer_token = _try_extract_from_header(request) + if not dify_config.ADMIN_API_KEY_ENABLE or not admin_api_key or not bearer_token: + return False + return hmac.compare_digest(bearer_token, admin_api_key) + + def extract_webapp_access_token(request: Request) -> str | None: return request.cookies.get(_real_cookie_name(COOKIE_NAME_WEBAPP_ACCESS_TOKEN)) or _try_extract_from_header(request) @@ -183,10 +197,8 @@ def build_force_logout_cookie_headers() -> list[str]: def check_csrf_token(request: Request, user_id: str): # some apis are sent by beacon, so we need to bypass csrf token check # since these APIs are post, they are already protected by SameSite: Lax, so csrf is not required. - if dify_config.ADMIN_API_KEY_ENABLE: - auth_token = extract_access_token(request) - if auth_token and auth_token == dify_config.ADMIN_API_KEY: - return + if is_admin_api_key_request(request): + return def _unauthorized(): raise Unauthorized("CSRF token is missing or invalid.") diff --git a/api/tests/unit_tests/extensions/test_ext_login.py b/api/tests/unit_tests/extensions/test_ext_login.py index 2207dac1e6a..feb4930292f 100644 --- a/api/tests/unit_tests/extensions/test_ext_login.py +++ b/api/tests/unit_tests/extensions/test_ext_login.py @@ -3,8 +3,9 @@ from typing import cast from unittest import mock import pytest -from flask import Response +from flask import Flask, Response, request +from constants import COOKIE_NAME_ACCESS_TOKEN from core.logging.context import clear_request_context, get_identity_context from extensions import ext_login from extensions.ext_login import unauthorized_handler @@ -74,3 +75,31 @@ def test_on_user_logged_in_logs_unsupported_user_type(caplog: pytest.LogCaptureF assert get_identity_context() == ("", "", "") assert "Failed to set logging identity context" in caplog.text + + +def test_admin_api_key_header_takes_precedence_over_console_cookie(monkeypatch: pytest.MonkeyPatch) -> None: + app = Flask(__name__) + session = mock.Mock(spec=ext_login.Session) + tenant = mock.Mock(spec=ext_login.Tenant) + tenant_account_join = mock.Mock(spec=ext_login.TenantAccountJoin) + account = mock.Mock(spec=ext_login.Account) + session.execute.return_value.one_or_none.return_value = (tenant, tenant_account_join) + session.scalar.return_value = account + monkeypatch.setattr(ext_login.dify_config, "ADMIN_API_KEY_ENABLE", True) + monkeypatch.setattr(ext_login.dify_config, "ADMIN_API_KEY", "admin-key") + monkeypatch.setattr(ext_login.dify_config, "CONSOLE_WEB_URL", "http://console.example.com") + monkeypatch.setattr(ext_login.dify_config, "CONSOLE_API_URL", "http://api.example.com") + monkeypatch.setattr(ext_login.dify_config, "COOKIE_DOMAIN", "") + + with app.test_request_context( + "/console/api/test", + headers={ + "Authorization": "Bearer admin-key", + "Cookie": f"{COOKIE_NAME_ACCESS_TOKEN}=console-session", + "X-WORKSPACE-ID": "workspace-id", + }, + ): + result = ext_login._load_user_from_request(request, session) + + assert result is account + account.set_current_tenant_with_session.assert_called_once_with(tenant, session=session) diff --git a/api/tests/unit_tests/libs/test_token.py b/api/tests/unit_tests/libs/test_token.py index 97d156478d6..f129a1d86ce 100644 --- a/api/tests/unit_tests/libs/test_token.py +++ b/api/tests/unit_tests/libs/test_token.py @@ -92,3 +92,21 @@ def test_non_whitelisted_path_requires_csrf(): with pytest.raises(Unauthorized): token.check_csrf_token(request, "account-1") + + +def test_admin_api_key_header_bypasses_csrf_when_console_cookie_is_present(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(token.dify_config, "ADMIN_API_KEY_ENABLE", True) + monkeypatch.setattr(token.dify_config, "ADMIN_API_KEY", "admin-key") + monkeypatch.setattr(token.dify_config, "CONSOLE_WEB_URL", "http://console.example.com") + monkeypatch.setattr(token.dify_config, "CONSOLE_API_URL", "http://api.example.com") + monkeypatch.setattr(token.dify_config, "COOKIE_DOMAIN", "") + request = cast( + Request, + MockRequest( + headers={"Authorization": "Bearer admin-key"}, + cookies={COOKIE_NAME_ACCESS_TOKEN: "console-session"}, + args={}, + ), + ) + + token.check_csrf_token(request, "account-1")