diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index ab644118a47..d0f060dc07d 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -119,3 +119,57 @@ assertions: [succeeds] source: "server.py:1089" rationale: Smoke; rarely used; same auth model as tools +- id: mcp.list_tools.api_key.team_toolset_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "auth/user_api_key_auth_mcp.py:2453" + rationale: "A toolset attached to a team's object_permission narrows the team key's tools/list to exactly the tools it names (LIT-5749: stored but inert before the fix)" + fail_before_fix: proven +- id: mcp.call_tool.api_key.team_toolset_denied_outside + module: mcp + tier: P0 + operation: call_tool + auth_family: api_key + assertions: [denied_outside_toolset] + source: "mcp_server_manager.py:4925" + rationale: "A team key calling a tool on the granted server but outside the team's toolset is refused; the named tool stays callable (LIT-5749)" + fail_before_fix: proven +- id: mcp.list_tools.api_key.org_toolset_scoped + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "auth/user_api_key_auth_mcp.py:1449" + rationale: "A toolset on an org's object_permission caps a member team's key that declares no MCP grants of its own (LIT-5749: the org substitution amplifier)" + fail_before_fix: proven +- id: mcp.call_tool.api_key.org_toolset_denied_outside + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [denied_outside_toolset] + source: "mcp_server_manager.py:4925" + rationale: "An org-inherited key calling outside the org's toolset is refused while the named tool stays callable (LIT-5749)" + fail_before_fix: proven +- id: mcp.list_tools.api_key.user_toolset_scoped + module: mcp + tier: P1 + operation: list_tools + auth_family: api_key + assertions: [toolset_scoped] + source: "auth/user_api_key_auth_mcp.py:2859" + rationale: "A toolset on an internal user's row ceils the servers the user's own key grants; user-level toolsets narrow, never grant (LIT-5749)" + fail_before_fix: proven +- id: mcp.call_tool.api_key.user_toolset_denied_outside + module: mcp + tier: P1 + operation: call_tool + auth_family: api_key + assertions: [denied_outside_toolset] + source: "mcp_server_manager.py:4925" + rationale: "A key whose user row holds a toolset is refused calling outside it even though the key itself grants the whole server (LIT-5749)" + fail_before_fix: proven diff --git a/tests/e2e/mcp/datadog_mcp.py b/tests/e2e/mcp/datadog_mcp.py index d1ea53a0b3b..54c5899812b 100644 --- a/tests/e2e/mcp/datadog_mcp.py +++ b/tests/e2e/mcp/datadog_mcp.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from collections.abc import Sequence from e2e_config import datadog_mcp_url, unique_marker from lifecycle import ResourceManager @@ -35,7 +36,11 @@ def register_datadog_mcp( resources: ResourceManager, *, mcp_access_groups: list[str] | None = None, + allowed_tools: Sequence[str] | None = (SEARCH_LOGS_TOOL,), ) -> str: + """Register the Datadog remote MCP server. `allowed_tools` defaults to the + search-logs slice; pass None to expose the full core toolset (needed when a + test must prove narrowing, so a second tool has to exist to be denied).""" assert_dd_mcp_creds() name = f"e2e_dd_mcp_{unique_marker()}" server_id = client.register_server( @@ -47,7 +52,7 @@ def register_datadog_mcp( "DD-API-KEY": _dd_api_key(), "DD-APPLICATION-KEY": _dd_app_key(), }, - allowed_tools=[SEARCH_LOGS_TOOL], + allowed_tools=list(allowed_tools) if allowed_tools is not None else None, mcp_access_groups=mcp_access_groups, ) resources.defer(lambda: client.delete_server(server_id)) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 73453478e5a..8bd61e04335 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -20,7 +20,19 @@ from pydantic import BaseModel, ConfigDict, Field, RootModel from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap -from models import KeyGenerateBody, ObjectPermission +from models import ( + KeyGenerateBody, + ObjectPermission, + OrgDeleteBody, + OrgNewBody, + OrgNewResponse, + TeamDeleteBody, + TeamNewBody, + TeamNewResponse, + UserDeleteBody, + UserNewBody, + UserNewResponse, +) from proxy_client import ProxyClient McpToolArg = str | int | float | bool | list[str] | dict[str, str] @@ -46,6 +58,20 @@ class McpServerNewResponse(BaseModel): server_id: str +class McpToolsetToolSpec(BaseModel): + server_id: str + tool_name: str + + +class McpToolsetNewBody(BaseModel): + toolset_name: str + tools: list[McpToolsetToolSpec] + + +class McpToolsetNewResponse(BaseModel): + toolset_id: str + + class McpServerRow(BaseModel): server_id: str alias: str | None = None @@ -74,9 +100,7 @@ class McpToolsListResponse(BaseModel): def tool_names_for_server(self, server_id: str) -> frozenset[str]: return frozenset( - tool.name - for tool in self.tools - if tool.mcp_info is not None and tool.mcp_info.server_id == server_id + tool.name for tool in self.tools if tool.mcp_info is not None and tool.mcp_info.server_id == server_id ) def tool_name_containing(self, server_id: str, needle: str) -> str | None: @@ -225,6 +249,7 @@ class McpClient: mcp_servers: list[str] | None, mcp_access_groups: list[str] | None = None, models: list[str] | None = None, + team_id: str | None = None, ) -> str: object_permission = ( ObjectPermission(mcp_servers=mcp_servers, mcp_access_groups=mcp_access_groups) @@ -235,10 +260,114 @@ class McpClient: KeyGenerateBody( models=models if models is not None else [], user_id=user_id, + team_id=team_id, object_permission=object_permission, ) ) + def create_toolset(self, *, name: str, server_id: str, tool_names: list[str]) -> str: + """Create a named toolset over `server_id` (admin), returning its id.""" + return unwrap( + self.proxy.transport.post( + "/v1/mcp/toolset", + headers=self.proxy.transport.master, + json=McpToolsetNewBody( + toolset_name=name, + tools=[McpToolsetToolSpec(server_id=server_id, tool_name=tool) for tool in tool_names], + ), + response_type=McpToolsetNewResponse, + ) + ).toolset_id + + def delete_toolset(self, toolset_id: str) -> None: + _ = self.proxy.transport.delete( + f"/v1/mcp/toolset/{toolset_id}", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + def create_team( + self, + *, + alias: str, + object_permission: ObjectPermission | None = None, + organization_id: str | None = None, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/team/new", + headers=self.proxy.transport.master, + json=TeamNewBody( + team_alias=alias, + organization_id=organization_id, + object_permission=object_permission, + ), + response_type=TeamNewResponse, + ) + ).team_id + + def delete_team(self, team_id: str) -> None: + _ = self.proxy.transport.post( + "/team/delete", + headers=self.proxy.transport.master, + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + + def create_org( + self, + *, + alias: str, + object_permission: ObjectPermission | None = None, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/organization/new", + headers=self.proxy.transport.master, + json=OrgNewBody( + organization_alias=alias, + object_permission=object_permission, + ), + response_type=OrgNewResponse, + ) + ).organization_id + + def delete_org(self, organization_id: str) -> None: + _ = self.proxy.transport.delete( + "/organization/delete", + headers=self.proxy.transport.master, + json=OrgDeleteBody(organization_ids=[organization_id]), + response_type=NoBody, + ) + + def create_user( + self, + *, + user_email: str, + object_permission: ObjectPermission | None = None, + ) -> str: + return unwrap( + self.proxy.transport.post( + "/user/new", + headers=self.proxy.transport.master, + json=UserNewBody( + user_email=user_email, + user_role="internal_user", + object_permission=object_permission, + ), + response_type=UserNewResponse, + ) + ).user_id + + def delete_user(self, user_id: str) -> None: + _ = self.proxy.transport.post( + "/user/delete", + headers=self.proxy.transport.master, + json=UserDeleteBody(user_ids=[user_id]), + response_type=NoBody, + ) + def list_tools(self, key: str) -> Result[McpToolsListResponse]: return self.proxy.transport.get( "/mcp-rest/tools/list", @@ -316,13 +445,10 @@ class McpClient: if isinstance(last, UnknownApiError) and last.status_code == 403: return last if not _is_mcp_not_synced(last, tool_name=name): - raise AssertionError( - f"ungranted key's tools/call was not 403 access_denied: {last}" - ) + raise AssertionError(f"ungranted key's tools/call was not 403 access_denied: {last}") if time.monotonic() >= deadline: raise AssertionError( - f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; " - f"last result: {last}" + f"ungranted key never got 403 for {name!r} within {self.proxy.poll_timeout}s; last result: {last}" ) time.sleep(self.proxy.poll_interval) @@ -368,9 +494,7 @@ class McpClient: return self.proxy.transport.post( "/mcp-rest/tools/call", headers=ApiKeyHeaders(x_litellm_api_key=key), - json=McpCallToolBody( - name=name, arguments=dict(arguments), server_id=server_id - ), + json=McpCallToolBody(name=name, arguments=dict(arguments), server_id=server_id), response_type=McpCallToolResponse, ) @@ -403,9 +527,7 @@ def _is_mcp_not_synced( # Gateway: "Tool search_datadog_logs not found" (optionally inside a longer message) if tool_name is not None: - return ( - re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None - ) + return re.search(rf"\btool\s+{re.escape(tool_name)}\s+not found\b", body_l) is not None return re.search(r"\btool\s+\S+\s+not found\b", body_l) is not None diff --git a/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py new file mode 100644 index 00000000000..b6c1f66328d --- /dev/null +++ b/tests/e2e/mcp/test_mcp_toolset_enforcement_e2e.py @@ -0,0 +1,181 @@ +"""Live e2e: an MCP toolset attached to a team, org, or internal user narrows a +real MCP server to exactly the tools the toolset names (LIT-5749). + +An admin registers the Datadog remote MCP server with its full core toolset (no +`allowed_tools` cap, so more than one tool exists to be denied), creates a +toolset naming only `search_datadog_logs` on it, and attaches that toolset at +one principal level per test. A control key granted the whole server first +proves the upstream is alive and serves a second tool, so a later denial is an +authorization decision rather than a dead server. The scoped key must then see +exactly the toolset's one tool on `tools/list`, be refused (403) calling any +other tool on the same server, and still successfully call the granted tool. + +Levels covered, one per test: a team holding server + toolset, a team that +inherits its grant from an org holding server + toolset, and an internal user +whose toolset ceils the servers their own key grants. +""" + +from __future__ import annotations + +import pytest + +from datadog_mcp import SEARCH_LOGS_TOOL, register_datadog_mcp +from e2e_config import DD_SEARCH_FROM, unique_marker +from e2e_http import unwrap +from lifecycle import ResourceManager +from mcp_client import McpClient +from models import ObjectPermission + +pytestmark = pytest.mark.e2e + + +def _open_server_and_toolset(client: McpClient, resources: ResourceManager) -> tuple[str, str]: + """Register the DD server uncapped and a toolset naming only the search-logs + tool on it, returning (server_id, toolset_id).""" + server_id = register_datadog_mcp(client, resources, allowed_tools=None) + client.await_registered(server_id) + toolset_id = client.create_toolset( + name=f"e2e-search-only-{unique_marker()}", + server_id=server_id, + tool_names=[SEARCH_LOGS_TOOL], + ) + resources.defer(lambda: client.delete_toolset(toolset_id)) + return server_id, toolset_id + + +def _control_tools(client: McpClient, resources: ResourceManager, server_id: str) -> tuple[str, str]: + """A fully-granted control key's view of the server: the fully-qualified + search-logs tool name and one tool outside the toolset. Proves the upstream + is alive and actually serves something a toolset can deny.""" + control_key = client.generate_key(user_id=f"e2e-ts-control-{unique_marker()}", mcp_servers=[server_id]) + resources.defer(lambda: client.proxy.delete_key(control_key)) + granted_tool = client.await_tool(control_key, server_id, SEARCH_LOGS_TOOL) + all_tools = unwrap(client.list_tools(control_key)).tool_names_for_server(server_id) + outside = sorted(all_tools - {granted_tool}) + assert outside, ( + f"the uncapped Datadog core toolset served only {all_tools}; a toolset " + f"cannot be proven to narrow a one-tool server" + ) + return granted_tool, outside[0] + + +def _assert_toolset_ceiling( + client: McpClient, + scoped_key: str, + *, + server_id: str, + granted_tool: str, + outside_tool: str, +) -> None: + """The scoped key sees exactly the toolset's tool, is 403-refused on a tool + outside it, and can still execute the granted one.""" + _ = client.await_tool(scoped_key, server_id, SEARCH_LOGS_TOOL) + listed = unwrap(client.list_tools(scoped_key)).tool_names_for_server(server_id) + assert listed == frozenset({granted_tool}), ( + f"toolset-scoped key must list exactly {granted_tool!r}; the toolset ceiling leaked: {sorted(listed)}" + ) + + denied = client.await_call_tool_denied(scoped_key, server_id=server_id, name=outside_tool, arguments={}) + assert denied.status_code == 403 + + result = client.await_call_tool( + scoped_key, + server_id=server_id, + name=granted_tool, + arguments={"query": f"service:e2e-toolset-{unique_marker()}", "from": DD_SEARCH_FROM}, + ) + assert result.is_error is not True, f"the toolset-granted tool must stay callable; upstream said: {result.all_text}" + + +class TestMcpToolsetEnforcementPerLevel: + @pytest.mark.covers("mcp.list_tools.api_key.team_toolset_scoped") + @pytest.mark.covers("mcp.call_tool.api_key.team_toolset_denied_outside") + def test_team_toolset_narrows_team_key( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + server_id, toolset_id = _open_server_and_toolset(client, resources) + granted_tool, outside_tool = _control_tools(client, resources, server_id) + + team_id = client.create_team( + alias=f"e2e-ts-team-{unique_marker()}", + object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]), + ) + resources.defer(lambda: client.delete_team(team_id)) + team_key = client.generate_key( + user_id=f"e2e-ts-team-user-{unique_marker()}", + mcp_servers=None, + team_id=team_id, + ) + resources.defer(lambda: client.proxy.delete_key(team_key)) + + _assert_toolset_ceiling( + client, + team_key, + server_id=server_id, + granted_tool=granted_tool, + outside_tool=outside_tool, + ) + + @pytest.mark.covers("mcp.list_tools.api_key.org_toolset_scoped") + @pytest.mark.covers("mcp.call_tool.api_key.org_toolset_denied_outside") + def test_org_toolset_caps_inherited_team_key( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + server_id, toolset_id = _open_server_and_toolset(client, resources) + granted_tool, outside_tool = _control_tools(client, resources, server_id) + + org_id = client.create_org( + alias=f"e2e-ts-org-{unique_marker()}", + object_permission=ObjectPermission(mcp_servers=[server_id], mcp_toolsets=[toolset_id]), + ) + resources.defer(lambda: client.delete_org(org_id)) + # The team declares no MCP grants of its own; whatever its key can reach + # comes from the org, so the org's toolset must cap it. + team_id = client.create_team(alias=f"e2e-ts-org-team-{unique_marker()}", organization_id=org_id) + resources.defer(lambda: client.delete_team(team_id)) + team_key = client.generate_key( + user_id=f"e2e-ts-org-user-{unique_marker()}", + mcp_servers=None, + team_id=team_id, + ) + resources.defer(lambda: client.proxy.delete_key(team_key)) + + _assert_toolset_ceiling( + client, + team_key, + server_id=server_id, + granted_tool=granted_tool, + outside_tool=outside_tool, + ) + + @pytest.mark.covers("mcp.list_tools.api_key.user_toolset_scoped") + @pytest.mark.covers("mcp.call_tool.api_key.user_toolset_denied_outside") + def test_user_toolset_ceils_own_key_grant( + self, + client: McpClient, + resources: ResourceManager, + ) -> None: + server_id, toolset_id = _open_server_and_toolset(client, resources) + granted_tool, outside_tool = _control_tools(client, resources, server_id) + + # The user's row holds only the toolset; their key grants the whole + # server. The user-level toolset is a ceiling over the key's grant. + user_id = client.create_user( + user_email=f"e2e-ts-user-{unique_marker()}@example.com", + object_permission=ObjectPermission(mcp_toolsets=[toolset_id]), + ) + resources.defer(lambda: client.delete_user(user_id)) + user_key = client.generate_key(user_id=user_id, mcp_servers=[server_id]) + resources.defer(lambda: client.proxy.delete_key(user_key)) + + _assert_toolset_ceiling( + client, + user_key, + server_id=server_id, + granted_tool=granted_tool, + outside_tool=outside_tool, + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 95a02b58824..331ea3fbb55 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -52,6 +52,7 @@ class KeyMetadata(BaseModel): class ObjectPermission(BaseModel): mcp_servers: list[str] | None = None mcp_access_groups: list[str] | None = None + mcp_toolsets: list[str] | None = None class KeyGenerateBody(BaseModel): @@ -922,6 +923,7 @@ class TeamNewBody(BaseModel): team_id: str | None = None organization_id: str | None = None metadata: TeamMetadata | None = None + object_permission: ObjectPermission | None = None class TeamNewResponse(BaseModel): @@ -979,6 +981,7 @@ class UserNewBody(BaseModel): user_email: str user_role: UserRole user_id: str | None = None + object_permission: ObjectPermission | None = None class UserNewResponse(BaseModel): @@ -1029,6 +1032,7 @@ class UserListResponse(BaseModel): class OrgNewBody(BaseModel): organization_alias: str models: list[str] = [] + object_permission: ObjectPermission | None = None class OrgNewResponse(BaseModel):