This is an automated email from the ASF dual-hosted git repository. yuqi1129 pushed a commit to branch feat/mcp-governance-task3-6 in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit a821a009ea0f1246e788aba46a7857a168bb8ac1 Author: yuqi <[email protected]> AuthorDate: Wed Jun 10 20:36:40 2026 +0800 [#11572] feat(mcp-server): per-request token isolation for HTTP transport GravitinoContext.rest_client() now extracts the Bearer token from the current HTTP request's Authorization header via fastmcp.server.dependencies.get_http_request(). A per-request token creates a fresh httpx client carrying that token, ensuring concurrent multi-principal sessions are fully isolated — one principal's token never leaks into another's Gravitino calls. Falls back to the shared default client (startup token) in stdio mode or when no Authorization header is present. Also add _extract_bearer_token() and _get_request_token() as testable helpers. tests/unit/test_per_request_token.py: 11 tests covering token extraction, stdio fallback, per-request override, and concurrent session isolation. --- mcp-server/mcp_server/core/context.py | 46 ++++++- mcp-server/mcp_server/server.py | 7 +- mcp-server/tests/unit/test_per_request_token.py | 160 ++++++++++++++++++++++++ 3 files changed, 207 insertions(+), 6 deletions(-) diff --git a/mcp-server/mcp_server/core/context.py b/mcp-server/mcp_server/core/context.py index 83350a7d10..f2cde66692 100644 --- a/mcp-server/mcp_server/core/context.py +++ b/mcp-server/mcp_server/core/context.py @@ -19,11 +19,53 @@ from mcp_server.client.factory import RESTClientFactory from mcp_server.core.setting import Setting +def _extract_bearer_token(authorization: str) -> str: + """Parse a Bearer token from a raw Authorization header value.""" + parts = authorization.split() + if len(parts) == 2 and parts[0].lower() == "bearer": + return parts[1] + return "" + + +def _get_request_token() -> str: + """Extract the Bearer token from the current HTTP request, if any. + + Returns an empty string in stdio mode or when the header is absent. + """ + try: + from fastmcp.server.dependencies import get_http_request + + authorization = get_http_request().headers.get("authorization", "") + return _extract_bearer_token(authorization) + except Exception: # noqa: BLE001 – stdio mode or missing request context + return "" + + class GravitinoContext: def __init__(self, setting: Setting): - self.gravitino_client = RESTClientFactory.create_rest_client( + self._setting = setting + # Fallback client for stdio mode or when no per-request token is present. + self._default_client = RESTClientFactory.create_rest_client( setting.metalake, setting.gravitino_uri, setting.token ) def rest_client(self): - return self.gravitino_client + """Return a REST client carrying the correct identity for this request. + + In HTTP transport mode the Bearer token from the incoming MCP request's + Authorization header takes priority over the static startup token. + This ensures concurrent sessions with different principals are fully + isolated — one principal's token never leaks into another's Gravitino calls. + + Falls back to the shared default client (static startup token) when: + - running in stdio mode (no HTTP request context), or + - the incoming request carries no Authorization header. + """ + request_token = _get_request_token() + if request_token: + return RESTClientFactory.create_rest_client( + self._setting.metalake, + self._setting.gravitino_uri, + request_token, + ) + return self._default_client diff --git a/mcp-server/mcp_server/server.py b/mcp-server/mcp_server/server.py index 8d66fec936..aea7df7f78 100644 --- a/mcp-server/mcp_server/server.py +++ b/mcp-server/mcp_server/server.py @@ -43,15 +43,14 @@ from mcp_server.tools import load_tools def _get_principal_from_request() -> str: - """Extract the principal from the current HTTP request's Authorization header. + """Derive a display principal from the current HTTP request's Authorization header. - Returns "anonymous" in stdio mode (no HTTP request) or when no token is present. + Returns "anonymous" in stdio mode or when no token is present. """ try: from fastmcp.server.dependencies import get_http_request - request = get_http_request() - authorization = request.headers.get("authorization", "") + authorization = get_http_request().headers.get("authorization", "") return audit._extract_principal(authorization) except Exception: # noqa: BLE001 – stdio mode or no request context return "anonymous" diff --git a/mcp-server/tests/unit/test_per_request_token.py b/mcp-server/tests/unit/test_per_request_token.py new file mode 100644 index 0000000000..c08cbeeb39 --- /dev/null +++ b/mcp-server/tests/unit/test_per_request_token.py @@ -0,0 +1,160 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Tests for per-request token isolation (Task 6). + +GravitinoContext.rest_client() must return a client carrying the token from +the current HTTP request, not the shared startup token, when an Authorization +header is present. This ensures concurrent multi-principal sessions are fully +isolated in HTTP transport mode. +""" + +import unittest +from unittest.mock import MagicMock, patch + +from mcp_server.core.context import ( + GravitinoContext, + _extract_bearer_token, + _get_request_token, +) +from mcp_server.core.setting import Setting + + +class TestExtractBearerToken(unittest.TestCase): + """Unit tests for the token extraction helper.""" + + def test_well_formed_bearer_header(self): + self.assertEqual(_extract_bearer_token("Bearer mytoken123"), "mytoken123") + + def test_case_insensitive_bearer(self): + self.assertEqual(_extract_bearer_token("bearer MYTOKEN"), "MYTOKEN") + + def test_empty_header_returns_empty(self): + self.assertEqual(_extract_bearer_token(""), "") + + def test_non_bearer_scheme_returns_empty(self): + self.assertEqual(_extract_bearer_token("Basic dXNlcjpwYXNz"), "") + + def test_only_scheme_no_token_returns_empty(self): + self.assertEqual(_extract_bearer_token("Bearer"), "") + + +class TestGetRequestToken(unittest.TestCase): + """Unit tests for _get_request_token() (HTTP context extraction).""" + + def test_returns_token_when_http_request_available(self): + mock_request = MagicMock() + mock_request.headers.get.return_value = "Bearer request-token-xyz" + + with patch( + "fastmcp.server.dependencies.get_http_request", + return_value=mock_request, + ): + token = _get_request_token() + + self.assertEqual(token, "request-token-xyz") + + def test_returns_empty_when_no_http_context(self): + """Simulates stdio mode where get_http_request raises LookupError.""" + with patch( + "fastmcp.server.dependencies.get_http_request", + side_effect=LookupError("no request context"), + ): + token = _get_request_token() + + self.assertEqual(token, "") + + def test_returns_empty_when_no_authorization_header(self): + mock_request = MagicMock() + mock_request.headers.get.return_value = "" + + with patch( + "fastmcp.server.dependencies.get_http_request", + return_value=mock_request, + ): + token = _get_request_token() + + self.assertEqual(token, "") + + +class TestGravitinoContextPerRequestToken(unittest.TestCase): + """GravitinoContext.rest_client() isolates per-request tokens.""" + + def _make_context(self, startup_token: str = "") -> GravitinoContext: + return GravitinoContext( + Setting( + metalake="ml", + gravitino_uri="http://localhost:8090";, + token=startup_token, + ) + ) + + def test_per_request_token_overrides_startup_token(self): + """When an HTTP request carries a token, it takes priority over the startup token.""" + ctx = self._make_context(startup_token="startup-token") + + mock_request = MagicMock() + mock_request.headers.get.return_value = "Bearer request-token" + + with patch( + "fastmcp.server.dependencies.get_http_request", + return_value=mock_request, + ): + client = ctx.rest_client() + + headers = dict(client._catalog_operation.rest_client.headers) + self.assertEqual(headers.get("authorization"), "Bearer request-token") + + def test_falls_back_to_default_client_when_no_request_token(self): + """When no per-request token exists, the shared default client (startup token) is used.""" + ctx = self._make_context(startup_token="startup-token") + + with patch( + "fastmcp.server.dependencies.get_http_request", + side_effect=LookupError, + ): + client = ctx.rest_client() + + # Must be the exact same object as the cached default client. + self.assertIs(client, ctx._default_client) + + def test_two_concurrent_requests_get_different_clients(self): + """Different request tokens must produce different client instances.""" + ctx = self._make_context() + + def make_mock(token: str) -> MagicMock: + m = MagicMock() + m.headers.get.return_value = f"Bearer {token}" + return m + + with patch( + "fastmcp.server.dependencies.get_http_request", + return_value=make_mock("alice-token"), + ): + client_alice = ctx.rest_client() + + with patch( + "fastmcp.server.dependencies.get_http_request", + return_value=make_mock("bob-token"), + ): + client_bob = ctx.rest_client() + + alice_headers = dict(client_alice._catalog_operation.rest_client.headers) + bob_headers = dict(client_bob._catalog_operation.rest_client.headers) + self.assertEqual(alice_headers.get("authorization"), "Bearer alice-token") + self.assertEqual(bob_headers.get("authorization"), "Bearer bob-token") + self.assertIsNot(client_alice, client_bob)
