fix(mcp): bound YAML integer parsing

This commit is contained in:
Sisyphus 2026-10-08 09:02:02 +08:00
parent 529cf36c15
commit 21efe15ba2
2 changed files with 20 additions and 0 deletions

View file

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

View file

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