diff --git a/api/controllers/inner_api/wraps.py b/api/controllers/inner_api/wraps.py index 7c60b316e8..990f33b647 100644 --- a/api/controllers/inner_api/wraps.py +++ b/api/controllers/inner_api/wraps.py @@ -2,6 +2,7 @@ from base64 import b64encode from collections.abc import Callable from functools import wraps from hashlib import sha1 +from hmac import compare_digest from hmac import new as hmac_new from typing import ParamSpec, TypeVar @@ -22,7 +23,9 @@ def billing_inner_api_only(view: Callable[P, R]): # get header 'X-Inner-Api-Key' inner_api_key = request.headers.get("X-Inner-Api-Key") - if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY: + if not inner_api_key or not ( + dify_config.INNER_API_KEY and compare_digest(inner_api_key, dify_config.INNER_API_KEY) + ): abort(401) return view(*args, **kwargs) @@ -38,7 +41,9 @@ def enterprise_inner_api_only(view: Callable[P, R]): # get header 'X-Inner-Api-Key' inner_api_key = request.headers.get("X-Inner-Api-Key") - if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY: + if not inner_api_key or not ( + dify_config.INNER_API_KEY and compare_digest(inner_api_key, dify_config.INNER_API_KEY) + ): abort(401) return view(*args, **kwargs) @@ -72,7 +77,7 @@ def enterprise_inner_api_user_auth(view: Callable[P, R]): signature = hmac_new(inner_api_key.encode("utf-8"), data_to_sign.encode("utf-8"), sha1) signature_base64 = b64encode(signature.digest()).decode("utf-8") - if signature_base64 != token: + if not compare_digest(signature_base64, token): return view(*args, **kwargs) kwargs["user"] = db.session.get(EndUser, user_id) @@ -90,7 +95,9 @@ def plugin_inner_api_only(view: Callable[P, R]): # get header 'X-Inner-Api-Key' inner_api_key = request.headers.get("X-Inner-Api-Key") - if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY_FOR_PLUGIN: + if not inner_api_key or not ( + dify_config.INNER_API_KEY_FOR_PLUGIN and compare_digest(inner_api_key, dify_config.INNER_API_KEY_FOR_PLUGIN) + ): abort(404) return view(*args, **kwargs) diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index 02e50a90fc..2d46ec60e6 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -1,3 +1,4 @@ +import hmac import json import flask_login @@ -32,7 +33,7 @@ def load_user_from_request(request_from_flask_login): # 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: + if admin_api_key and hmac.compare_digest(admin_api_key, auth_token): workspace_id = request.headers.get("X-WORKSPACE-ID") if workspace_id: tenant_account_join = db.session.execute( diff --git a/api/libs/token.py b/api/libs/token.py index a34db70764..777df2dbd7 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 @@ -191,7 +192,7 @@ def check_csrf_token(request: Request, user_id: str): # 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: + if auth_token and dify_config.ADMIN_API_KEY and hmac.compare_digest(auth_token, dify_config.ADMIN_API_KEY): return def _unauthorized():