You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by ru...@apache.org on 2026/06/18 17:28:00 UTC
(superset) branch master updated: fix(mcp): add safeguards to ensure all MCP tools are wrapped with mcp_auth_hook (#40412)
This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 76e2418f1ea fix(mcp): add safeguards to ensure all MCP tools are wrapped with mcp_auth_hook (#40412)
76e2418f1ea is described below
commit 76e2418f1eab7efd915f7cf335336e5c18bb7ed3
Author: Abdul Rehman <76...@users.noreply.github.com>
AuthorDate: Thu Jun 18 22:27:53 2026 +0500
fix(mcp): add safeguards to ensure all MCP tools are wrapped with mcp_auth_hook (#40412)
Co-authored-by: Evan Rusackas <ev...@preset.io>
---
superset/core/mcp/core_mcp_injection.py | 13 +-
superset/mcp_service/app.py | 66 +++++++
superset/mcp_service/auth.py | 4 +
.../unit_tests/mcp_service/test_auth_safeguards.py | 211 +++++++++++++++++++++
4 files changed, 290 insertions(+), 4 deletions(-)
diff --git a/superset/core/mcp/core_mcp_injection.py b/superset/core/mcp/core_mcp_injection.py
index abfbeac236a..d143857a685 100644
--- a/superset/core/mcp/core_mcp_injection.py
+++ b/superset/core/mcp/core_mcp_injection.py
@@ -151,9 +151,11 @@ def create_tool_decorator(
return wrapped_func
except Exception as e:
+ # Fail-fast: don't silently serve unprotected tools. Returning the
+ # unwrapped function here would bypass mcp_auth_hook, leaving the
+ # tool exposed without auth or context isolation (see #39395).
logger.error("Failed to register MCP tool %s: %s", name or func.__name__, e)
- # Return the original function so extension doesn't break
- return func
+ raise
# If called as @tool (without parentheses)
if callable(func_or_name):
@@ -247,11 +249,14 @@ def create_prompt_decorator(
return wrapped_func
except Exception as e:
+ # Fail-fast: don't silently serve unprotected prompts. Returning
+ # the unwrapped function here would bypass mcp_auth_hook, leaving
+ # the prompt exposed without auth or context isolation (see
+ # #39395 — same invariant as create_tool_decorator).
logger.error(
"Failed to register MCP prompt %s: %s", name or func.__name__, e
)
- # Return the original function so extension doesn't break
- return func
+ raise
# If called as @prompt (without parentheses)
if callable(func_or_name):
diff --git a/superset/mcp_service/app.py b/superset/mcp_service/app.py
index a45e1578d7e..ac76545e5a0 100644
--- a/superset/mcp_service/app.py
+++ b/superset/mcp_service/app.py
@@ -758,6 +758,69 @@ from superset.mcp_service.user.tool import ( # noqa: F401, E402
list_users,
)
+#: Tool names exempt from the mcp_auth_hook protection check. Adding a tool
+#: here is a security-significant choice — review carefully. Entries are tools
+#: that intentionally run without authentication; ``generate_bug_report`` is
+#: public so users can collect diagnostics even when auth itself is broken.
+#: Frozen so accidental post-init mutation (``ALLOWED_UNPROTECTED.add(...)``)
+#: raises ``AttributeError`` rather than silently widening the security
+#: allowlist after the startup assertion has already run.
+ALLOWED_UNPROTECTED: frozenset[str] = frozenset({"generate_bug_report"})
+
+
+def assert_all_tools_protected(mcp_instance: FastMCP) -> None:
+ """Fail loudly at startup if any registered tool bypassed ``mcp_auth_hook``.
+
+ The fresh-app-context-per-call fix in #39385 only protects tools that
+ actually go through ``mcp_auth_hook``. This catches all three known bypass
+ paths (see #39395):
+
+ * ``@tool(protect=False)`` — the wrapper is skipped entirely.
+ * Silent fallback in ``create_tool_decorator`` (now fail-fast, but a future
+ regression could reintroduce it).
+ * Direct ``mcp.add_tool()`` calls that skip the decorator.
+
+ Raises:
+ RuntimeError: if any tool's underlying function lacks the
+ ``_mcp_auth_protected`` marker set by ``mcp_auth_hook``.
+ """
+ # FastMCP 3.x exposes components keyed as ``"<kind>:<name>@..."`` (tools,
+ # prompts, resources) in the local provider's component dict. Tool values
+ # are ``FunctionTool`` objects with ``.name`` and ``.fn`` attributes.
+ tools_checked = 0
+ for key, component in mcp_instance.local_provider._components.items():
+ # Prompts and resources are intentionally skipped here. They use the
+ # same ``mcp_auth_hook`` (via ``create_prompt_decorator`` and the
+ # resource-level ``@mcp_auth_hook`` convention documented in
+ # ``mcp_service/CLAUDE.md``) but their bypass surface is different —
+ # ``protect=False`` on a prompt would need its own ``assert_all_
+ # prompts_protected`` check. Tracked as a follow-up per @aminghadersohi.
+ if not key.startswith("tool:"):
+ continue
+ tools_checked += 1
+ name = getattr(component, "name", None) or key
+ fn = getattr(component, "fn", None)
+ if name in ALLOWED_UNPROTECTED:
+ continue
+ if not getattr(fn, "_mcp_auth_protected", False):
+ raise RuntimeError(
+ f"SECURITY: MCP tool '{name}' registered without mcp_auth_hook. "
+ f"All tools must use @tool() with protect=True or be explicitly "
+ f"allowlisted in ALLOWED_UNPROTECTED."
+ )
+
+ # Defense against silent FastMCP API drift: if the private
+ # ``local_provider._components`` attribute or the ``"tool:"`` key prefix
+ # changes in a future FastMCP release, this loop would match nothing and
+ # vacuously return success. Log a warning so the regression is visible in
+ # the startup logs and routine ops review.
+ if tools_checked == 0:
+ logger.warning(
+ "assert_all_tools_protected inspected 0 tools — FastMCP internal "
+ "API (local_provider._components, 'tool:' key prefix) may have "
+ "changed. Review and update the iteration in app.py."
+ )
+
def _remove_disabled_tools(disabled_tools: set[str]) -> None:
"""Remove tools listed in MCP_DISABLED_TOOLS from the global MCP instance.
@@ -890,5 +953,8 @@ def init_fastmcp_server(
# Apply any additional configuration
_apply_config(mcp, config)
+ # Final invariant: every tool must have gone through mcp_auth_hook.
+ assert_all_tools_protected(mcp)
+
logger.info("Configured FastMCP instance: %s (auth=%s)", name, auth is not None)
return mcp
diff --git a/superset/mcp_service/auth.py b/superset/mcp_service/auth.py
index 68da0c68261..da225837a2e 100644
--- a/superset/mcp_service/auth.py
+++ b/superset/mcp_service/auth.py
@@ -1081,4 +1081,8 @@ def mcp_auth_hook(tool_func: F) -> F: # noqa: C901
if "ctx" in new_wrapper.__annotations__:
del new_wrapper.__annotations__["ctx"]
+ # Mark this wrapper as protected so a startup assertion can verify every
+ # registered tool went through mcp_auth_hook (see issue #39395).
+ new_wrapper._mcp_auth_protected = True # type: ignore[attr-defined]
+
return new_wrapper # type: ignore[return-value]
diff --git a/tests/unit_tests/mcp_service/test_auth_safeguards.py b/tests/unit_tests/mcp_service/test_auth_safeguards.py
new file mode 100644
index 00000000000..de5d59906af
--- /dev/null
+++ b/tests/unit_tests/mcp_service/test_auth_safeguards.py
@@ -0,0 +1,211 @@
+# 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 MCP tool auth safeguards (#39395).
+
+The race-condition fix in #39385 only protects tools that go through
+``mcp_auth_hook``. These tests pin the three structural safeguards that
+enforce that invariant:
+
+1. ``mcp_auth_hook`` stamps ``_mcp_auth_protected = True`` on the wrapper.
+2. ``assert_all_tools_protected`` raises if any tool lacks the marker.
+3. ``create_tool_decorator`` fails fast instead of returning the unwrapped
+ function when registration errors.
+"""
+
+import logging
+from types import SimpleNamespace
+from typing import Any, Callable
+from unittest.mock import patch
+
+import pytest
+
+from superset.core.mcp.core_mcp_injection import (
+ create_prompt_decorator,
+ create_tool_decorator,
+)
+from superset.mcp_service.app import assert_all_tools_protected
+from superset.mcp_service.auth import mcp_auth_hook
+
+
+def _fake_mcp_with_tools(tools_by_name: dict[str, Any]) -> SimpleNamespace:
+ """Build a minimal FastMCP stand-in exposing ``local_provider._components``.
+
+ FastMCP 3.x keys components as ``"<kind>:<name>@"`` — tools, prompts and
+ resources all share the dict. ``assert_all_tools_protected`` filters to the
+ ``tool:`` entries, so the fake mirrors that shape.
+ """
+ components = {f"tool:{name}@": tool for name, tool in tools_by_name.items()}
+ return SimpleNamespace(
+ local_provider=SimpleNamespace(_components=components),
+ )
+
+
+def test_mcp_auth_hook_stamps_protection_marker() -> None:
+ """``mcp_auth_hook`` must mark its wrapper so the startup assertion can
+ verify the tool went through it."""
+
+ def sample_tool(request: dict[str, Any]) -> dict[str, Any]:
+ """A tool that does nothing useful — only the wrapping matters."""
+ return request
+
+ wrapped = mcp_auth_hook(sample_tool)
+ assert getattr(wrapped, "_mcp_auth_protected", False) is True
+
+
+def test_assert_all_tools_protected_passes_when_every_tool_is_marked() -> None:
+ """Happy path: every tool went through ``mcp_auth_hook`` — no raise."""
+
+ def protected_fn() -> None:
+ pass
+
+ protected_fn._mcp_auth_protected = True # type: ignore[attr-defined]
+ mcp = _fake_mcp_with_tools(
+ {
+ "list_charts": SimpleNamespace(name="list_charts", fn=protected_fn),
+ "get_chart_info": SimpleNamespace(name="get_chart_info", fn=protected_fn),
+ }
+ )
+
+ # Should not raise.
+ assert_all_tools_protected(mcp)
+
+
+def test_assert_all_tools_protected_raises_on_unprotected_tool() -> None:
+ """An unwrapped tool (the three bypass paths in #39395) must blow up at
+ startup rather than silently serve unprotected traffic."""
+
+ def unprotected_fn() -> None:
+ pass # no ``_mcp_auth_protected`` attribute
+
+ mcp = _fake_mcp_with_tools(
+ {"sneaky_tool": SimpleNamespace(name="sneaky_tool", fn=unprotected_fn)}
+ )
+
+ with pytest.raises(RuntimeError, match="sneaky_tool.*without mcp_auth_hook"):
+ assert_all_tools_protected(mcp)
+
+
+def test_assert_all_tools_protected_respects_allowlist() -> None:
+ """Tools explicitly listed in ``ALLOWED_UNPROTECTED`` are skipped — the
+ allowlist is the only legitimate path to a public tool."""
+
+ def unprotected_fn() -> None:
+ pass
+
+ mcp = _fake_mcp_with_tools(
+ {
+ "public_health_probe": SimpleNamespace(
+ name="public_health_probe", fn=unprotected_fn
+ )
+ }
+ )
+
+ # Replace the module-level frozenset so the function under test sees a
+ # deterministic allowlist regardless of which real tools exist.
+ with patch(
+ "superset.mcp_service.app.ALLOWED_UNPROTECTED",
+ frozenset({"public_health_probe"}),
+ ):
+ # Should not raise — the tool is allowlisted.
+ assert_all_tools_protected(mcp)
+
+
+def test_create_tool_decorator_fails_fast_on_registration_error() -> None:
+ """When tool registration fails at any point after ``mcp_auth_hook`` wraps
+ the function, ``create_tool_decorator`` must re-raise instead of returning
+ the unwrapped function — the silent-return path was the bug #39395 closes."""
+
+ def sample_tool() -> None:
+ pass
+
+ # Make Tool.from_function raise — simulates any registration-time failure
+ # in the path after mcp_auth_hook returns. Previously this would log an
+ # error and return ``func`` unwrapped; now it must re-raise.
+ decorator: Callable[..., Any] = create_tool_decorator()
+ with patch("fastmcp.tools.Tool.from_function", side_effect=RuntimeError("boom")):
+ with pytest.raises(RuntimeError, match="boom"):
+ decorator(sample_tool)
+
+
+def test_create_prompt_decorator_fails_fast_on_registration_error() -> None:
+ """``create_prompt_decorator`` must propagate registration errors instead
+ of returning the unwrapped function. Prompts can expose system instructions
+ and sensitive context to LLM clients — the same bypass risk as tools
+ (raised by @aminghadersohi on PR #40412)."""
+
+ def sample_prompt() -> str:
+ return "hi"
+
+ decorator: Callable[..., Any] = create_prompt_decorator()
+ with patch(
+ "fastmcp.prompts.Prompt.from_function", side_effect=RuntimeError("boom")
+ ):
+ with pytest.raises(RuntimeError, match="boom"):
+ decorator(sample_prompt)
+
+
+def test_assert_all_tools_protected_skips_non_tool_components() -> None:
+ """``local_provider._components`` mixes tools, prompts and resources under
+ different key prefixes. Only ``tool:`` entries should be checked — prompts
+ and resources have their own auth model and must not trip the assertion."""
+
+ def protected_fn() -> None:
+ pass
+
+ protected_fn._mcp_auth_protected = True # type: ignore[attr-defined]
+
+ def unprotected_fn() -> None:
+ pass
+
+ # Hand-craft a components dict directly so prompt/resource keys are present
+ # (the helper only builds ``tool:`` entries).
+ components = {
+ "tool:list_charts@": SimpleNamespace(name="list_charts", fn=protected_fn),
+ "prompt:create_chart_guided@": SimpleNamespace(
+ name="create_chart_guided", fn=unprotected_fn
+ ),
+ "resource:chart://configs@": SimpleNamespace(
+ name="chart://configs", fn=unprotected_fn
+ ),
+ }
+ mcp = SimpleNamespace(
+ local_provider=SimpleNamespace(_components=components),
+ )
+
+ # Should not raise — the unprotected prompt/resource entries are skipped
+ # because their keys don't start with ``tool:``.
+ assert_all_tools_protected(mcp)
+
+
+def test_assert_all_tools_protected_warns_when_no_tools_found(
+ caplog: pytest.LogCaptureFixture,
+) -> None:
+ """Defense against silent FastMCP API drift: if ``_components`` ever
+ contains no ``tool:`` entries (e.g. FastMCP renamed the attribute or
+ changed the key prefix), the function must log a warning rather than
+ return success vacuously."""
+ mcp = SimpleNamespace(
+ local_provider=SimpleNamespace(_components={}),
+ )
+
+ with caplog.at_level(logging.WARNING, logger="superset.mcp_service.app"):
+ assert_all_tools_protected(mcp)
+
+ assert any(
+ "inspected 0 tools" in rec.message and "FastMCP internal" in rec.message
+ for rec in caplog.records
+ ), "expected a 'no tools inspected' warning when components dict is empty"