From 57b0d7f45fe00982369241109c130f6238a3d574 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 23 Apr 2026 15:23:53 -0700 Subject: [PATCH] fix(mcp): resolve tool_permissions dict keys by id-or-name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR feedback (greptile P1 / veria high): with the previous change, a team storing mcp_tool_permissions={"my-alias": ["read_file"]} would pass the server-access check (because the alias expanded to a concrete id in the allowed-servers list) but the per-server tool lookup still did dict.get(server_id) against the raw name-keyed dict — missing, returning None, which callers treat as "no restrictions" → all tools allowed instead of only the declared ones. Add MCPServerManager.expand_tool_permissions() that rewrites the dict so every key is a concrete server_id where possible (tool lists from keys pointing at the same server are unioned). Unresolved keys pass through unchanged so stale id-keyed restrictions still apply when the same string is used for lookup. Wire the helper into the four dict-lookup sites: get_allowed_tools_for_server (key + team paths), the agent tool lookup, and the rest_endpoints.py tool filter. Also switch expand_permission_list to pass through unresolved entries (rather than dropping them) so existing test fixtures that use bare string placeholders continue to work. The downstream access check denies unknown entries when compared to the concrete request server_id, so security posture is unchanged. Sanitize the debug log to use %r formatting so an admin-controlled identifier with newlines can't forge log entries (CodeQL log-injection warning). --- .../mcp_server/auth/user_api_key_auth_mcp.py | 53 ++++++++---- .../mcp_server/mcp_server_manager.py | 42 +++++++-- .../mcp_server/rest_endpoints.py | 10 +-- .../mcp_server/test_mcp_server_manager.py | 85 ++++++++++++++++++- 4 files changed, 161 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index db413527982..d4799e9f208 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -516,15 +516,26 @@ class MCPRequestHandler: user_api_key_auth ) - # Extract tool permissions for this server + # Extract tool permissions for this server. Dict keys may be + # server_ids OR names/aliases; normalize to server_id-keyed form + # before lookup so a name-based key does not silently drop its + # tool restrictions when server_id is the resolved uuid. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + key_tools = ( - key_obj_perm.mcp_tool_permissions.get(server_id) - if key_obj_perm and key_obj_perm.mcp_tool_permissions + global_mcp_server_manager.expand_tool_permissions( + key_obj_perm.mcp_tool_permissions + ).get(server_id) + if key_obj_perm else None ) team_tools = ( - team_obj_perm.mcp_tool_permissions.get(server_id) - if team_obj_perm and team_obj_perm.mcp_tool_permissions + global_mcp_server_manager.expand_tool_permissions( + team_obj_perm.mcp_tool_permissions + ).get(server_id) + if team_obj_perm else None ) @@ -660,8 +671,10 @@ class MCPRequestHandler: ) # servers referenced in tool permissions should also be accessible - tool_perm_servers = global_mcp_server_manager.expand_permission_list( - list((key_object_permission.mcp_tool_permissions or {}).keys()) + tool_perm_servers = list( + global_mcp_server_manager.expand_tool_permissions( + key_object_permission.mcp_tool_permissions + ).keys() ) # Combine all lists @@ -708,8 +721,10 @@ class MCPRequestHandler: ) # servers referenced in tool permissions should also be accessible - tool_perm_servers = global_mcp_server_manager.expand_permission_list( - list((object_permissions.mcp_tool_permissions or {}).keys()) + tool_perm_servers = list( + global_mcp_server_manager.expand_tool_permissions( + object_permissions.mcp_tool_permissions + ).keys() ) # Combine all lists @@ -775,8 +790,10 @@ class MCPRequestHandler: ) # servers referenced in tool permissions should also be accessible - tool_perm_servers = global_mcp_server_manager.expand_permission_list( - list((end_user_obj.object_permission.mcp_tool_permissions or {}).keys()) + tool_perm_servers = list( + global_mcp_server_manager.expand_tool_permissions( + end_user_obj.object_permission.mcp_tool_permissions + ).keys() ) # Combine all lists @@ -905,12 +922,16 @@ class MCPRequestHandler: return None mcp_tool_permissions = getattr(obj_perm, "mcp_tool_permissions", None) - if not mcp_tool_permissions: + if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict): return None - if isinstance(mcp_tool_permissions, dict): - tools = mcp_tool_permissions.get(server_id) - else: - tools = None + # Dict keys may be server_ids OR names/aliases; normalize before lookup. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + tools = global_mcp_server_manager.expand_tool_permissions( + mcp_tool_permissions + ).get(server_id) return list(tools) if tools else None except Exception as e: verbose_logger.warning( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 4173dd5abb2..251f271903b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2771,11 +2771,12 @@ class MCPServerManager: server_ids against the current region's config + DB registry union. Entries that match a server_id pass through unchanged. Entries that - don't are looked up by alias/server_name/name and replaced with every - matching server_id (duplicate names grant access to all matches). - Entries that resolve to nothing are dropped and a debug log is + match an alias/server_name/name are replaced with every matching + server_id (duplicate names grant access to all matches). Entries + that resolve to nothing pass through as-is and a debug log is emitted so admins can diagnose stale/typo permission entries — the - final access-check still denies them either way. + downstream access-check denies them when compared against the + concrete request server_id. """ if not identifiers: return [] @@ -2795,12 +2796,41 @@ class MCPServerManager: if matches: expanded.update(matches) else: + # %r quotes and escapes control chars so an admin-controlled + # identifier with newlines cannot forge log lines. verbose_logger.debug( - f"MCP permission entry '{identifier}' does not resolve " - "to any known server (config + DB union). Skipping." + "MCP permission entry %r does not resolve to any known " + "server (config + DB union). Passing through — the " + "downstream access check will deny it if it's stale.", + identifier, ) + expanded.add(identifier) return list(expanded) + def expand_tool_permissions( + self, + tool_permissions: Optional[Dict[str, List[str]]], + ) -> Dict[str, List[str]]: + """ + Rewrite an ``mcp_tool_permissions`` dict keyed by id/name/alias so + every key is a concrete server_id where possible. Tool lists from + keys that point at the same server are unioned, matching the + "duplicate names grant access to all matches" semantics of + ``expand_permission_list``. + + Required so name-based keys don't silently drop their tool + restrictions when the lookup uses the resolved server_id. Unresolved + keys pass through via ``expand_permission_list`` so stale id-keyed + restrictions still apply when the same string is used for lookup. + """ + if not tool_permissions: + return {} + result: Dict[str, List[str]] = {} + for key, tools in tool_permissions.items(): + for server_id in self.expand_permission_list([key]): + result.setdefault(server_id, []).extend(tools or []) + return result + def get_mcp_server_by_name( self, server_name: str, client_ip: Optional[str] = None ) -> Optional[MCPServer]: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 8131c040136..f1030bf0a8a 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -315,11 +315,11 @@ if MCP_AVAILABLE: and user_api_key_auth.object_permission and user_api_key_auth.object_permission.mcp_tool_permissions ): - allowed_tools_for_server = ( - user_api_key_auth.object_permission.mcp_tool_permissions.get( - server.server_id - ) - ) + # Dict keys may be server_ids OR names/aliases; normalize so lookup + # by concrete server_id resolves name-keyed restrictions too. + allowed_tools_for_server = global_mcp_server_manager.expand_tool_permissions( + user_api_key_auth.object_permission.mcp_tool_permissions + ).get(server.server_id) if ( allowed_tools_for_server is not None and len(allowed_tools_for_server) > 0 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index dbb53b9168b..6ac2b0ac4fa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2649,11 +2649,14 @@ class TestMCPServerManagerExpandPermissionList: assert manager.expand_permission_list(["public_alias"]) == ["id-1"] - def test_drops_unknown_entry_silently(self): + def test_passes_through_unknown_entry(self): + """Unresolved entries pass through unchanged (with a debug log) — + the downstream access check denies them when compared to the + concrete request server_id.""" manager = MCPServerManager() manager.config_mcp_servers["id-1"] = self._make_server("id-1", server_name="b") - assert manager.expand_permission_list(["a"]) == [] + assert manager.expand_permission_list(["a"]) == ["a"] def test_name_collision_expands_to_all_matches(self): """Two servers sharing a server_name both resolve — the documented behavior.""" @@ -2734,5 +2737,83 @@ class TestMCPServerManagerExpandPermissionList: assert usc1.expand_permission_list(["a"]) == ["hash-usc1"] +class TestMCPServerManagerExpandToolPermissions: + """Tests for tool-permission dict rewriting — the privilege-escalation guard.""" + + def _make_server(self, server_id: str, server_name: str, alias=None) -> MCPServer: + return MCPServer( + server_id=server_id, + name=alias or server_name, + alias=alias, + server_name=server_name, + url=f"https://{server_id}.example.com", + transport=MCPTransport.http, + ) + + def test_empty_or_none_returns_empty_dict(self): + manager = MCPServerManager() + assert manager.expand_tool_permissions(None) == {} + assert manager.expand_tool_permissions({}) == {} + + def test_rewrites_name_key_to_server_id(self): + """Privilege-escalation guard: a name-based key must resolve to the + concrete server_id, otherwise `.get(server_id)` misses and the tool + restriction is silently dropped (caller treats None as allow-all).""" + manager = MCPServerManager() + manager.config_mcp_servers["uuid-a"] = self._make_server( + "uuid-a", server_name="my-alias" + ) + + result = manager.expand_tool_permissions({"my-alias": ["read_file"]}) + assert result == {"uuid-a": ["read_file"]} + + def test_passes_through_existing_server_id_key(self): + manager = MCPServerManager() + manager.config_mcp_servers["uuid-a"] = self._make_server( + "uuid-a", server_name="alpha" + ) + + result = manager.expand_tool_permissions({"uuid-a": ["read_file"]}) + assert result == {"uuid-a": ["read_file"]} + + def test_unresolved_key_passes_through_unchanged(self): + """A stale id-keyed restriction (server since deleted, or just a + test-fixture placeholder) must still apply when something looks it + up by that same string — dropping would silently remove the + restriction.""" + manager = MCPServerManager() + + result = manager.expand_tool_permissions({"stale-uuid": ["read_file"]}) + assert result == {"stale-uuid": ["read_file"]} + + def test_name_collision_unions_tool_lists(self): + """Two servers sharing a server_name both match; their tool lists get + the restriction (matches the list-expansion collision semantics).""" + manager = MCPServerManager() + manager.config_mcp_servers["uuid-1"] = self._make_server( + "uuid-1", server_name="shared" + ) + manager.registry["uuid-2"] = self._make_server("uuid-2", server_name="shared") + + result = manager.expand_tool_permissions({"shared": ["read_file"]}) + assert sorted(result.keys()) == ["uuid-1", "uuid-2"] + assert result["uuid-1"] == ["read_file"] + assert result["uuid-2"] == ["read_file"] + + def test_id_and_name_keys_pointing_at_same_server_union_tools(self): + """If the admin writes both {"uuid-a": [...], "alias-a": [...]} and + both refer to the same server, the tool lists are unioned rather + than one overwriting the other.""" + manager = MCPServerManager() + manager.config_mcp_servers["uuid-a"] = self._make_server( + "uuid-a", server_name="alias-a" + ) + + result = manager.expand_tool_permissions( + {"uuid-a": ["read_file"], "alias-a": ["write_file"]} + ) + assert sorted(result["uuid-a"]) == ["read_file", "write_file"] + + if __name__ == "__main__": pytest.main([__file__])