mirror of
https://github.com/langgenius/dify.git
synced 2026-04-15 18:06:36 +08:00
refactor(api): replace json.loads with Pydantic validation in security and tools layers (#34380)
This commit is contained in:
parent
05c5327f47
commit
89ce61cfea
@ -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,
|
||||
)
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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"}
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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 ""
|
||||
|
||||
|
||||
@ -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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user