refactor(api): replace json.loads with Pydantic validation in security and tools layers (#34380)

api/core/provider_manager.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 import contextlib
-import json
 from collections import defaultdict
 from collections.abc import Sequence
 from json import JSONDecodeError
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
 from graphon.model_runtime.entities.model_entities import ModelType
 from graphon.model_runtime.entities.provider_entities import (
@@ -15,6 +14,7 @@ from graphon.model_runtime.entities.provider_entities import (
     ProviderEntity,
 )
 from graphon.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
+from pydantic import TypeAdapter
 from sqlalchemy import select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
@@ -58,6 +58,8 @@ from services.feature_service import FeatureService
 if TYPE_CHECKING:
     from graphon.model_runtime.runtime import ModelRuntime
 
+_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
+
 
 class ProviderManager:
     """
@@ -875,8 +877,8 @@ class ProviderManager:
             return {"openai_api_key": encrypted_config}
 
         try:
-            credentials = cast(dict, json.loads(encrypted_config))
+            credentials = _credentials_adapter.validate_json(encrypted_config)
-        except JSONDecodeError:
+        except (ValueError, JSONDecodeError):
             return {}
 
         # Decrypt secret variables
@@ -1015,7 +1017,7 @@ class ProviderManager:
                 if not cached_provider_credentials:
                     provider_credentials: dict[str, Any] = {}
                     if provider_records and provider_records[0].encrypted_config:
-                        provider_credentials = json.loads(provider_records[0].encrypted_config)
+                        provider_credentials = _credentials_adapter.validate_json(provider_records[0].encrypted_config)
 
                     # Get provider credential secret variables
                     provider_credential_secret_variables = self._extract_secret_variables(
@@ -1162,8 +1164,10 @@ class ProviderManager:
 
                         if not cached_provider_model_credentials:
                             try:
-                                provider_model_credentials = json.loads(load_balancing_model_config.encrypted_config)
+                                provider_model_credentials = _credentials_adapter.validate_json(
-                            except JSONDecodeError:
+                                    load_balancing_model_config.encrypted_config
+                                )
+                            except (ValueError, JSONDecodeError):
                                 continue
 
                             # Get decoding rsa key and cipher for decrypting credentials
@@ -1176,7 +1180,7 @@ class ProviderManager:
                                 if variable in provider_model_credentials:
                                     try:
                                         provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding(
-                                            provider_model_credentials.get(variable),
+                                            provider_model_credentials.get(variable) or "",
                                             self.decoding_rsa_key,
                                             self.decoding_cipher_rsa,
                                         )

api/core/tools/entities/tool_entities.py

@@ -6,7 +6,17 @@ from collections.abc import Mapping
 from enum import StrEnum, auto
 from typing import Any, Union
 
-from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_serializer, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    TypeAdapter,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
+from typing_extensions import TypedDict
 
 from core.entities.provider_entities import ProviderConfig
 from core.plugin.entities.parameters import (
@@ -23,6 +33,14 @@ from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.constants import TOOL_SELECTOR_MODEL_IDENTITY
 
 
+class EmojiIconDict(TypedDict):
+    background: str
+    content: str
+
+
+emoji_icon_adapter: TypeAdapter[EmojiIconDict] = TypeAdapter(EmojiIconDict)
+
+
 class ToolLabelEnum(StrEnum):
     SEARCH = "search"
     IMAGE = "image"

api/core/tools/tool_manager.py

@@ -5,12 +5,14 @@ import time
 from collections.abc import Generator, Mapping
 from os import listdir, path
 from threading import Lock
-from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol, TypedDict, Union, cast
+from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol, Union, cast
 
 import sqlalchemy as sa
 from graphon.runtime import VariablePool
+from pydantic import TypeAdapter
 from sqlalchemy import select
 from sqlalchemy.orm import Session
+from typing_extensions import TypedDict
 from yarl import URL
 
 import contexts
@@ -49,9 +51,11 @@ from core.tools.entities.api_entities import ToolProviderApiEntity, ToolProvider
 from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.tool_entities import (
     ApiProviderAuthType,
+    EmojiIconDict,
     ToolInvokeFrom,
     ToolParameter,
     ToolProviderType,
+    emoji_icon_adapter,
 )
 from core.tools.errors import ToolProviderNotFoundError
 from core.tools.tool_label_manager import ToolLabelManager
@@ -72,9 +76,7 @@ class ApiProviderControllerItem(TypedDict):
     controller: ApiToolProviderController
 
 
-class EmojiIconDict(TypedDict):
+_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
-    background: str
-    content: str
 
 
 class WorkflowToolRuntimeSpec(Protocol):
@@ -885,7 +887,7 @@ class ToolManager:
             raise ValueError(f"you have not added provider {provider_name}")
 
         try:
-            credentials = json.loads(provider_obj.credentials_str) or {}
+            credentials = _credentials_adapter.validate_json(provider_obj.credentials_str) or {}
         except Exception:
             credentials = {}
 
@@ -910,7 +912,7 @@ class ToolManager:
         masked_credentials = encrypter.mask_plugin_credentials(encrypter.decrypt(credentials))
 
         try:
-            icon = json.loads(provider_obj.icon)
+            icon = emoji_icon_adapter.validate_json(provider_obj.icon)
         except Exception:
             icon = {"background": "#252525", "content": "\ud83d\ude01"}
 
@@ -973,7 +975,7 @@ class ToolManager:
             if workflow_provider is None:
                 raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
 
-            icon = json.loads(workflow_provider.icon)
+            icon = emoji_icon_adapter.validate_json(workflow_provider.icon)
             return icon
         except Exception:
             return {"background": "#252525", "content": "\ud83d\ude01"}
@@ -990,7 +992,7 @@ class ToolManager:
             if api_provider is None:
                 raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
 
-            icon = json.loads(api_provider.icon)
+            icon = emoji_icon_adapter.validate_json(api_provider.icon)
             return icon
         except Exception:
             return {"background": "#252525", "content": "\ud83d\ude01"}

api/libs/helper.py

@@ -18,8 +18,9 @@ from flask import Response, stream_with_context
 from flask_restx import fields
 from graphon.file import helpers as file_helpers
 from graphon.model_runtime.utils.encoders import jsonable_encoder
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from pydantic.functional_validators import AfterValidator
+from typing_extensions import TypedDict
 
 from configs import dify_config
 from core.app.features.rate_limiting.rate_limit import RateLimitGenerator
@@ -32,6 +33,17 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class _TokenData(TypedDict, total=False):
+    account_id: str | None
+    email: str
+    token_type: str
+    code: str
+    old_email: str
+
+
+_token_data_adapter: TypeAdapter[_TokenData] = TypeAdapter(_TokenData)
+
+
 def _stream_with_request_context(response: object) -> Any:
     """Bridge Flask's loosely-typed streaming helper without leaking casts into callers."""
     return cast(Any, stream_with_context)(response)
@@ -443,7 +455,7 @@ class TokenManager:
         if token_data_json is None:
             logger.warning("%s token %s not found with key %s", token_type, token, key)
             return None
-        token_data: dict[str, Any] | None = json.loads(token_data_json)
+        token_data = dict(_token_data_adapter.validate_json(token_data_json))
         return token_data
 
     @classmethod

api/services/tools/tools_transform_service.py

@@ -1,4 +1,3 @@
-import json
 import logging
 from collections.abc import Mapping
 from typing import Any, Union
@@ -21,6 +20,7 @@ from core.tools.entities.tool_entities import (
     ApiProviderAuthType,
     ToolParameter,
     ToolProviderType,
+    emoji_icon_adapter,
 )
 from core.tools.plugin_tool.provider import PluginToolProviderController
 from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter
@@ -53,11 +53,14 @@ class ToolTransformService:
         elif provider_type in {ToolProviderType.API, ToolProviderType.WORKFLOW}:
             try:
                 if isinstance(icon, str):
-                    return json.loads(icon)
+                    parsed = emoji_icon_adapter.validate_json(icon)
-                return icon
+                    return {"background": parsed["background"], "content": parsed["content"]}
-            except (json.JSONDecodeError, ValueError):
+                return {"background": icon["background"], "content": icon["content"]}
+            except (ValueError, ValidationError, KeyError):
                 return {"background": "#252525", "content": "\ud83d\ude01"}
         elif provider_type == ToolProviderType.MCP:
+            if isinstance(icon, Mapping):
+                return {"background": icon.get("background", ""), "content": icon.get("content", "")}
             return icon
         return ""
 

api/services/tools/workflow_tools_manage_service.py

@@ -8,7 +8,7 @@ from sqlalchemy.orm import Session
 
 from core.tools.__base.tool_provider import ToolProviderController
 from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity
-from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration
+from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration, emoji_icon_adapter
 from core.tools.tool_label_manager import ToolLabelManager
 from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils
 from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
@@ -313,7 +313,7 @@ class WorkflowToolManageService:
             "label": db_tool.label,
             "workflow_tool_id": db_tool.id,
             "workflow_app_id": db_tool.app_id,
-            "icon": json.loads(db_tool.icon),
+            "icon": emoji_icon_adapter.validate_json(db_tool.icon),
             "description": db_tool.description,
             "parameters": jsonable_encoder(db_tool.parameter_configurations),
             "output_schema": output_schema,