fix(api): prevent cookies from shadowing admin API keys (#39872)

Signed-off-by: kenwoodjw <blackxin55+@gmail.com>

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:

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

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)

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