mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
fix(mcp): bound YAML integer parsing
This commit is contained in:
parent
529cf36c15
commit
21efe15ba2
2 changed files with 20 additions and 0 deletions
|
|
@ -28,6 +28,7 @@ from litellm.proxy._experimental.mcp_server.exceptions import (
|
|||
# which include '/'. Sanitize here so the same regex passes everywhere downstream.
|
||||
_OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
|
||||
_OPENAPI_TOOL_NAME_MAX_LEN: Final = 128
|
||||
_MAX_YAML_INT_LENGTH: Final = 1024
|
||||
|
||||
|
||||
def sanitize_openapi_tool_name(raw_name: str) -> str:
|
||||
|
|
@ -171,6 +172,7 @@ def _parse_openapi_spec(text: str) -> Mapping[str, Any]:
|
|||
parsed: Any = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
import yaml
|
||||
from yaml.nodes import ScalarNode
|
||||
|
||||
class _NoMergeSafeLoader(yaml.SafeLoader):
|
||||
def compose_node(self, parent: object, index: object) -> object:
|
||||
|
|
@ -184,6 +186,16 @@ def _parse_openapi_spec(text: str) -> Mapping[str, Any]:
|
|||
raise yaml.YAMLError("YAML merge keys are not supported")
|
||||
super().flatten_mapping(node)
|
||||
|
||||
def construct_yaml_int(self, node: ScalarNode) -> int:
|
||||
if len(node.value) > _MAX_YAML_INT_LENGTH:
|
||||
raise yaml.YAMLError("YAML integer is too long")
|
||||
return super().construct_yaml_int(node)
|
||||
|
||||
_NoMergeSafeLoader.add_constructor(
|
||||
"tag:yaml.org,2002:int",
|
||||
_NoMergeSafeLoader.construct_yaml_int,
|
||||
)
|
||||
|
||||
parsed = yaml.load(text, Loader=_NoMergeSafeLoader)
|
||||
|
||||
if not isinstance(parsed, dict):
|
||||
|
|
|
|||
|
|
@ -347,6 +347,14 @@ components:
|
|||
)
|
||||
|
||||
|
||||
def test_parse_openapi_spec_rejects_oversized_yaml_integer() -> None:
|
||||
import yaml
|
||||
|
||||
oversized_integer: Final = "1" + ":11" * gen._MAX_YAML_INT_LENGTH
|
||||
with pytest.raises(yaml.YAMLError, match="YAML integer is too long"):
|
||||
gen._parse_openapi_spec(f"openapi: 3.0.0\nx-expensive: {oversized_integer}\n")
|
||||
|
||||
|
||||
def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock:
|
||||
"""Utility to create a mocked async httpx client for the given method.
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue