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

This commit is contained in:
Dream 2026-04-07 08:11:51 -04:00 committed by GitHub
parent 05c5327f47
commit 89ce61cfea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 64 additions and 25 deletions

View File

@ -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,
)

View File

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

View File

@ -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"}

View File

@ -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

View File

@ -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)
return icon
except (json.JSONDecodeError, ValueError):
parsed = emoji_icon_adapter.validate_json(icon)
return {"background": parsed["background"], "content": parsed["content"]}
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 ""

View File

@ -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,