From 9051ed55eebd989e72882f0fd4cc3bab895620ac Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 15:45:52 -0700 Subject: [PATCH] test(mcp): cover scoped gateway JWT admission and rejection --- tests/e2e/coverage_registry/mcp.yaml | 90 +++++++++++++ tests/e2e/mcp/mcp_client.py | 12 +- tests/e2e/mcp/test_mcp_jwt_auth_e2e.py | 178 +++++++++++++++++++++++++ tests/e2e/models.py | 7 + 4 files changed, 281 insertions(+), 6 deletions(-) create mode 100644 tests/e2e/mcp/test_mcp_jwt_auth_e2e.py diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index 85ace835144..ca04593df01 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -127,3 +127,93 @@ assertions: [toolset_scoped] source: "user_api_key_auth_mcp.py:2137" rationale: "A key granted a toolset lists exactly the toolset's tools: the rest of the server's catalog stays hidden and every stored name resolves" + +- id: mcp.list_tools.bearer.jwt_valid_scoped + module: mcp + tier: P0 + operation: list_tools + auth_family: bearer + assertions: [valid_scoped] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.call_tool.bearer.jwt_valid_scoped + module: mcp + tier: P0 + operation: call_tool + auth_family: bearer + assertions: [valid_scoped] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.list_tools.bearer.jwt_invalid_signature_denied + module: mcp + tier: P0 + operation: list_tools + auth_family: bearer + assertions: [invalid_signature_denied] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.call_tool.bearer.jwt_invalid_signature_denied + module: mcp + tier: P0 + operation: call_tool + auth_family: bearer + assertions: [invalid_signature_denied] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.list_tools.bearer.jwt_expired_denied + module: mcp + tier: P0 + operation: list_tools + auth_family: bearer + assertions: [expired_denied] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.call_tool.bearer.jwt_expired_denied + module: mcp + tier: P0 + operation: call_tool + auth_family: bearer + assertions: [expired_denied] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.list_tools.bearer.jwt_inactive_user_denied + module: mcp + tier: P0 + operation: list_tools + auth_family: bearer + assertions: [inactive_user_denied] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.call_tool.bearer.jwt_inactive_user_denied + module: mcp + tier: P0 + operation: call_tool + auth_family: bearer + assertions: [inactive_user_denied] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.list_tools.bearer.jwt_header_precedence + module: mcp + tier: P0 + operation: list_tools + auth_family: bearer + assertions: [header_precedence] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" + +- id: mcp.call_tool.bearer.jwt_header_precedence + module: mcp + tier: P0 + operation: call_tool + auth_family: bearer + assertions: [header_precedence] + source: "test_mcp_jwt_auth_e2e.py" + rationale: "LIT-4506 preserves scoped gateway JWT admission independently of upstream credentials" diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 56f7fffba29..cbc79ee2c2a 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -17,12 +17,11 @@ from collections.abc import Mapping from dataclasses import dataclass from typing import Literal -from pydantic import BaseModel, ConfigDict, Field, RootModel - from e2e_config import settle_propagation -from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap +from e2e_http import AuthHeaders, Headers, NoBody, Result, Success, UnknownApiError, unwrap from models import KeyGenerateBody, McpServerListResponse, McpServerRow, ObjectPermission from proxy_client import ProxyClient +from pydantic import BaseModel, ConfigDict, Field, RootModel McpToolArg = str | int | float | bool | list[str] | dict[str, str] McpToolArguments = Mapping[str, McpToolArg] @@ -254,10 +253,10 @@ class McpClient: ) ) - def list_tools(self, key: str) -> Result[McpToolsListResponse]: + def list_tools(self, key: str, *, headers: AuthHeaders | None = None) -> Result[McpToolsListResponse]: return self.proxy.transport.get( "/mcp-rest/tools/list", - headers=ApiKeyHeaders(x_litellm_api_key=key), + headers=headers if headers is not None else ApiKeyHeaders(x_litellm_api_key=key), params=NoBody(), response_type=McpToolsListResponse, ) @@ -393,10 +392,11 @@ class McpClient: server_id: str, name: str, arguments: McpToolArguments, + headers: AuthHeaders | None = None, ) -> Result[McpCallToolResponse]: return self.proxy.transport.post( "/mcp-rest/tools/call", - headers=ApiKeyHeaders(x_litellm_api_key=key), + headers=headers if headers is not None else ApiKeyHeaders(x_litellm_api_key=key), json=McpCallToolBody( name=name, arguments=dict(arguments), server_id=server_id ), diff --git a/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py b/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py new file mode 100644 index 00000000000..0b25ffd1f69 --- /dev/null +++ b/tests/e2e/mcp/test_mcp_jwt_auth_e2e.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Final, Literal + +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 AuthHeaders, UnauthorizedError, unwrap +from idp import SHORT_LIVED_CLIENT_ID, Identity, Keycloak, token_claims +from lifecycle import ResourceManager +from management.management_client import build_client as build_management_client +from mcp_client import McpClient +from models import KeyGenerateBody, ObjectPermission, UserScimMetadata, UserUpdateBody + +pytestmark = pytest.mark.e2e + + +@dataclass(frozen=True, slots=True) +class GrantedMcp: + server_id: str + other_server_id: str + tool: str + token: str = field(repr=False) + + +def _allowed(client: McpClient, access: GrantedMcp, headers: AuthHeaders) -> None: + listing: Final = unwrap(client.list_tools(access.token, headers=headers)) + assert listing.tool_names_for_server(access.server_id) == frozenset((access.tool,)) + assert listing.tool_names_for_server(access.other_server_id) == frozenset() + result: Final = unwrap( + client.call_tool( + access.token, + server_id=access.server_id, + name=access.tool, + arguments={"query": f"lit4506-{unique_marker()}", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000}, + headers=headers, + ) + ) + assert result.is_error is False, result + assert result.all_text.strip(), "The successful control must return a real tool result" + + +def _denied(client: McpClient, access: GrantedMcp, headers: AuthHeaders, reason: str) -> None: + listing: Final = client.list_tools(access.token, headers=headers) + assert isinstance(listing, UnauthorizedError), listing + assert reason in listing.body.lower(), listing.body + execution: Final = client.call_tool( + access.token, + server_id=access.server_id, + name=access.tool, + arguments={"query": f"lit4506-{unique_marker()}", "from": DD_SEARCH_FROM, "to": "now", "max_tokens": 1000}, + headers=headers, + ) + assert isinstance(execution, UnauthorizedError), execution + assert reason in execution.body.lower(), execution.body + + +@pytest.fixture +def granted_mcp(client: McpClient, resources: ResourceManager, jwt_identity: Identity, idp: Keycloak) -> GrantedMcp: + server_id: Final = register_datadog_mcp(client, resources) + other_server_id: Final = register_datadog_mcp(client, resources) + client.await_registered(server_id) + client.await_registered(other_server_id) + management: Final = build_management_client(client.proxy) + management.update_user( + UserUpdateBody( + user_id=jwt_identity.user_id, + user_role="internal_user", + object_permission=ObjectPermission(mcp_servers=[server_id]), + ) + ) + stored: Final = management.user_info(jwt_identity.user_id) + assert stored.user_info.user_id == jwt_identity.user_id + assert stored.user_info.user_role == "internal_user" + token: Final = idp.access_token(jwt_identity) + assert token_claims(token).sub == jwt_identity.user_id + tool: Final = client.await_tool(token, server_id, SEARCH_LOGS_TOOL) + return GrantedMcp(server_id, other_server_id, tool, token) + + +class TestMcpGatewayJwt: + @pytest.mark.covers("mcp.list_tools.bearer.jwt_valid_scoped", "mcp.call_tool.bearer.jwt_valid_scoped") + @pytest.mark.parametrize("header", ("authorization", "x_litellm_api_key")) + def test_valid_scoped_jwt_lists_and_calls( + self, client: McpClient, granted_mcp: GrantedMcp, header: Literal["authorization", "x_litellm_api_key"] + ) -> None: + headers: Final = ( + AuthHeaders(authorization=f"Bearer {granted_mcp.token}") + if header == "authorization" + else AuthHeaders.model_validate({"x-litellm-api-key": granted_mcp.token}) + ) + _allowed(client, granted_mcp, headers) + + @pytest.mark.covers( + "mcp.list_tools.bearer.jwt_invalid_signature_denied", "mcp.call_tool.bearer.jwt_invalid_signature_denied" + ) + def test_tampered_signature_denies_both_operations(self, client: McpClient, granted_mcp: GrantedMcp) -> None: + valid: Final = AuthHeaders(authorization=f"Bearer {granted_mcp.token}") + _allowed(client, granted_mcp, valid) + header, payload, signature = granted_mcp.token.split(".") + flipped: Final = "A" if signature[10] != "A" else "B" + tampered: Final = f"{header}.{payload}.{signature[:10]}{flipped}{signature[11:]}" + _denied(client, granted_mcp, AuthHeaders(authorization=f"Bearer {tampered}"), "signature") + _allowed(client, granted_mcp, valid) + + @pytest.mark.covers("mcp.list_tools.bearer.jwt_expired_denied", "mcp.call_tool.bearer.jwt_expired_denied") + def test_expired_signed_jwt_denies_both_operations( + self, client: McpClient, granted_mcp: GrantedMcp, jwt_identity: Identity, idp: Keycloak + ) -> None: + valid: Final = AuthHeaders(authorization=f"Bearer {granted_mcp.token}") + _allowed(client, granted_mcp, valid) + expiring: Final = idp.access_token(jwt_identity, client_id=SHORT_LIVED_CLIENT_ID) + claims: Final = token_claims(expiring) + assert claims.sub == jwt_identity.user_id + delay: Final = claims.exp - time.time() + 1 + assert delay <= 5, "Short-lived token configuration or IdP clock drifted" + time.sleep(max(0, delay)) + _denied(client, granted_mcp, AuthHeaders(authorization=f"Bearer {expiring}"), "expired") + _allowed(client, granted_mcp, valid) + + @pytest.mark.covers( + "mcp.list_tools.bearer.jwt_inactive_user_denied", "mcp.call_tool.bearer.jwt_inactive_user_denied" + ) + def test_deactivated_user_cannot_reuse_warm_jwt( + self, client: McpClient, granted_mcp: GrantedMcp, jwt_identity: Identity + ) -> None: + headers: Final = AuthHeaders(authorization=f"Bearer {granted_mcp.token}") + _allowed(client, granted_mcp, headers) + management: Final = build_management_client(client.proxy) + management.update_user( + UserUpdateBody( + user_id=jwt_identity.user_id, user_role="internal_user", metadata=UserScimMetadata(scim_active=False) + ) + ) + stored: Final = management.user_info(jwt_identity.user_id) + assert stored.user_info.metadata is not None and stored.user_info.metadata.scim_active is False + _denied(client, granted_mcp, headers, "deactivated") + management.update_user( + UserUpdateBody( + user_id=jwt_identity.user_id, user_role="internal_user", metadata=UserScimMetadata(scim_active=True) + ) + ) + _allowed(client, granted_mcp, headers) + + @pytest.mark.covers("mcp.list_tools.bearer.jwt_header_precedence", "mcp.call_tool.bearer.jwt_header_precedence") + def test_explicit_gateway_header_wins_without_fallback( + self, client: McpClient, granted_mcp: GrantedMcp, resources: ResourceManager + ) -> None: + _allowed( + client, + granted_mcp, + AuthHeaders.model_validate( + {"x-litellm-api-key": granted_mcp.token, "authorization": "Bearer invalid-secondary-token"} + ), + ) + _denied( + client, + granted_mcp, + AuthHeaders.model_validate( + {"x-litellm-api-key": "sk-invalid-primary-token", "authorization": f"Bearer {granted_mcp.token}"} + ), + "key", + ) + sibling_key: Final = client.proxy.generate_key( + KeyGenerateBody(object_permission=ObjectPermission(mcp_servers=[granted_mcp.other_server_id])) + ) + resources.defer(lambda: client.proxy.delete_key(sibling_key)) + sibling_tool: Final = client.await_tool(sibling_key, granted_mcp.other_server_id, SEARCH_LOGS_TOOL) + sibling: Final = GrantedMcp(granted_mcp.other_server_id, granted_mcp.server_id, sibling_tool, sibling_key) + _allowed( + client, + sibling, + AuthHeaders.model_validate( + {"x-litellm-api-key": sibling_key, "authorization": f"Bearer {granted_mcp.token}"} + ), + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9f49c5974d0..eaf6de0bdd1 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1229,9 +1229,15 @@ class UserNewResponse(BaseModel): user_id: str +class UserScimMetadata(BaseModel): + scim_active: bool | None = None + + class UserUpdateBody(BaseModel): user_id: str user_role: UserRole + object_permission: ObjectPermission | None = None + metadata: UserScimMetadata | None = None class UserInfoParams(BaseModel): @@ -1242,6 +1248,7 @@ class UserData(BaseModel): user_id: str | None = None user_email: str | None = None user_role: str | None = None + metadata: UserScimMetadata | None = None class UserInfoResponse(BaseModel):