diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index db70c7d77de..0ef5059e62e 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -7222,11 +7222,30 @@ class MCPServerManager: if tool_name in self.tool_name_to_mcp_server_name_mapping: server_name: Final = self.tool_name_to_mcp_server_name_mapping[tool_name] normalised_lookup: Final = normalize_server_name(server_name) - if normalised_lookup in prefix_to_server: - return prefix_to_server[normalised_lookup] - for server in registry_servers: - if normalize_server_name(server.name) == normalised_lookup: - return server + registry_prefix_hit: Final = prefix_to_server.get(normalised_lookup) + registry_candidate: Final = ( + registry_prefix_hit + if registry_prefix_hit is not None + else next( + (server for server in registry_servers if normalize_server_name(server.name) == normalised_lookup), + None, + ) + ) + # The bare-name mapping is keyed off each server's own tool + # names, so any server may claim an arbitrary bare-name key. A + # name that carries a prefix-qualified wire form of a *different* + # registered server must not dispatch through the bare-name + # branch: an attacker server advertising e.g. "github-delete_repo" + # otherwise shadows the prefix-qualified github tool (and bypasses + # its per-tool ACL) (R2-23). + if registry_candidate is not None: + foreign_prefixes: Final = tuple( + known_prefix + for known_prefix, mapped_server in prefix_to_server.items() + if mapped_server is not registry_candidate + ) + if match_known_server_prefix(tool_name, foreign_prefixes) is None: + return registry_candidate # If not found and tool name is prefixed, extract the prefix and # match against any known form. diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_shadow_gate.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_shadow_gate.py new file mode 100644 index 00000000000..3e76de8c597 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_shadow_gate.py @@ -0,0 +1,51 @@ +# bare-name mapping shadow gate (LD fork R2-23): a bare-name tool dispatch +# must not be hijackable by a foreign server advertising another server's +# prefix-qualified tool names. + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._types import MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def _manager_with_github_and_evil(): + manager = MCPServerManager() + github = MCPServer( + server_id="srv-github", + name="github", + alias="github", + url="https://github-upstream.example/mcp", + transport=MCPTransport.http, + ) + evil = MCPServer( + server_id="srv-evil", + name="evil", + alias="evil", + url="https://evil-upstream.example/mcp", + transport=MCPTransport.http, + ) + manager.registry[github.server_id] = github + manager.registry[evil.server_id] = evil + return manager, github + + +def test_foreign_server_cannot_shadow_another_server_tool_name(): + """Server 'evil' lists a tool literally named 'github-delete_repo'; its + registration overwrites github's prefix-qualified routing row + ('github-delete_repo' -> 'evil'). The bare-name branch must refuse a + candidate when the requested name carries a *different* registered + server's known prefix, so the poisoned row cannot dispatch to evil.""" + manager, github = _manager_with_github_and_evil() + routing = manager.tool_name_to_mcp_server_name_mapping + routing["github-list_prs"] = "github" + routing["github-delete_repo"] = "evil" # evil overwrote github's row + routing["delete_repo"] = "evil" + routing["list_prs"] = "github" + try: + # the poisoned prefix-qualified row must not dispatch to evil + assert manager.get_mcp_server_from_tool_name("github-delete_repo") is None + assert manager.get_mcp_server_from_tool_name("github-list_prs") is github + # a legit bare-name row keeps dispatching to its owner + assert manager.get_mcp_server_from_tool_name("list_prs") is github + finally: + for key in ("github-list_prs", "github-delete_repo", "delete_repo", "list_prs"): + routing.pop(key, None)