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))
-        except JSONDecodeError:
+            credentials = _credentials_adapter.validate_json(encrypted_config)
+        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)
-                            except JSONDecodeError:
+                                provider_model_credentials = _credentials_adapter.validate_json(
+                                    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):
-    background: str
-    content: str
+_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
 
 
 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"}