aminghadersohi commented on code in PR #38747:
URL: https://github.com/apache/superset/pull/38747#discussion_r2961839019


##########
superset/mcp_service/auth.py:
##########
@@ -172,62 +172,121 @@ def load_user_with_relationships(
     return query.first()
 
 
-def get_user_from_request() -> User:
+def _resolve_user_from_jwt_context(app: Any) -> User | None:
     """
-    Get the current user for the MCP tool request.
+    Resolve the current user from the MCP SDK's per-request JWT context.
 
-    Priority order:
-    1. g.user if already set (by Preset workspace middleware)
-    2. MCP_DEV_USERNAME from configuration (for development/testing)
+    Uses FastMCP's ``get_access_token()`` which returns the JWT AccessToken
+    for the current async task via a ContextVar — safe across concurrent
+    requests, unlike ``g.user`` which can be stale.
+
+    The username is extracted from token claims using a configurable resolver
+    (``MCP_USER_RESOLVER`` config) or the default ``default_user_resolver()``.
 
     Returns:
-        User object with roles and groups eagerly loaded
+        User object with relationships loaded, or None if no JWT context
+        (i.e. no token present — caller should fall through to next source).
 
     Raises:
-        ValueError: If user cannot be authenticated or found
+        ValueError: If JWT resolves a username that doesn't exist in the DB
+            (fail closed — do NOT fall through to weaker auth sources).
     """
-    from flask import current_app
+    try:
+        from fastmcp.server.dependencies import get_access_token
+    except ImportError:
+        logger.debug("fastmcp.server.dependencies not available, skipping JWT 
context")
+        return None
 
-    # First check if user is already set by Preset workspace middleware
-    if hasattr(g, "user") and g.user:
-        return g.user
+    access_token = get_access_token()
+    if access_token is None:
+        return None
 
-    # Fall back to configured username for development/single-user deployments
-    username = current_app.config.get("MCP_DEV_USERNAME")
+    # Use configurable resolver or default
+    from superset.mcp_service.mcp_config import default_user_resolver
+
+    resolver = app.config.get("MCP_USER_RESOLVER", default_user_resolver)
+    username = resolver(app, access_token)
 
     if not username:
-        auth_enabled = current_app.config.get("MCP_AUTH_ENABLED", False)
-        jwt_configured = bool(
-            current_app.config.get("MCP_JWKS_URI")
-            or current_app.config.get("MCP_JWT_PUBLIC_KEY")
-            or current_app.config.get("MCP_JWT_SECRET")
-        )
-        details = []
-        details.append(
-            f"g.user was not set by JWT middleware "
-            f"(MCP_AUTH_ENABLED={auth_enabled}, "
-            f"JWT keys configured={jwt_configured})"
-        )
-        details.append("MCP_DEV_USERNAME is not configured")
-        raise ValueError(
-            "No authenticated user found. Tried:\n"
-            + "\n".join(f"  - {d}" for d in details)
-            + "\n\nEither pass a valid JWT bearer token or configure "
-            "MCP_DEV_USERNAME for development."
+        logger.warning(
+            "JWT context present but no username could be extracted from 
claims"
         )
+        return None

Review Comment:
   Agreed, fixed in 508869d. Now raises `ValueError` instead of returning 
`None` when JWT is present but no username can be extracted.



##########
superset/mcp_service/auth.py:
##########
@@ -172,62 +172,121 @@ def load_user_with_relationships(
     return query.first()
 
 
-def get_user_from_request() -> User:
+def _resolve_user_from_jwt_context(app: Any) -> User | None:
     """
-    Get the current user for the MCP tool request.
+    Resolve the current user from the MCP SDK's per-request JWT context.
 
-    Priority order:
-    1. g.user if already set (by Preset workspace middleware)
-    2. MCP_DEV_USERNAME from configuration (for development/testing)
+    Uses FastMCP's ``get_access_token()`` which returns the JWT AccessToken
+    for the current async task via a ContextVar — safe across concurrent
+    requests, unlike ``g.user`` which can be stale.
+
+    The username is extracted from token claims using a configurable resolver
+    (``MCP_USER_RESOLVER`` config) or the default ``default_user_resolver()``.
 
     Returns:
-        User object with roles and groups eagerly loaded
+        User object with relationships loaded, or None if no JWT context
+        (i.e. no token present — caller should fall through to next source).
 
     Raises:
-        ValueError: If user cannot be authenticated or found
+        ValueError: If JWT resolves a username that doesn't exist in the DB
+            (fail closed — do NOT fall through to weaker auth sources).
     """
-    from flask import current_app
+    try:
+        from fastmcp.server.dependencies import get_access_token
+    except ImportError:
+        logger.debug("fastmcp.server.dependencies not available, skipping JWT 
context")
+        return None
 
-    # First check if user is already set by Preset workspace middleware
-    if hasattr(g, "user") and g.user:
-        return g.user
+    access_token = get_access_token()
+    if access_token is None:
+        return None
 
-    # Fall back to configured username for development/single-user deployments
-    username = current_app.config.get("MCP_DEV_USERNAME")
+    # Use configurable resolver or default
+    from superset.mcp_service.mcp_config import default_user_resolver
+
+    resolver = app.config.get("MCP_USER_RESOLVER", default_user_resolver)
+    username = resolver(app, access_token)
 
     if not username:
-        auth_enabled = current_app.config.get("MCP_AUTH_ENABLED", False)
-        jwt_configured = bool(
-            current_app.config.get("MCP_JWKS_URI")
-            or current_app.config.get("MCP_JWT_PUBLIC_KEY")
-            or current_app.config.get("MCP_JWT_SECRET")
-        )
-        details = []
-        details.append(
-            f"g.user was not set by JWT middleware "
-            f"(MCP_AUTH_ENABLED={auth_enabled}, "
-            f"JWT keys configured={jwt_configured})"
-        )
-        details.append("MCP_DEV_USERNAME is not configured")
-        raise ValueError(
-            "No authenticated user found. Tried:\n"
-            + "\n".join(f"  - {d}" for d in details)
-            + "\n\nEither pass a valid JWT bearer token or configure "
-            "MCP_DEV_USERNAME for development."
+        logger.warning(
+            "JWT context present but no username could be extracted from 
claims"
         )
+        return None
 
-    # Use helper function to load user with all required relationships
     user = load_user_with_relationships(username)

Review Comment:
   Good catch, fixed in 508869d. Added `if not user and "@" in username: user = 
load_user_with_relationships(email=username)` before the fail-closed raise.



##########
superset/mcp_service/auth.py:
##########
@@ -339,6 +398,9 @@ def _get_app_context_manager() -> 
AbstractContextManager[None]:
         @functools.wraps(tool_func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
             with _get_app_context_manager():
+                # Clear stale g.user to prevent user impersonation across
+                # tool calls when no per-request middleware refreshes it.
+                g.pop("user", None)

Review Comment:
   Valid point, fixed in 508869d. Now only clears `g.user` when 
`has_request_context()` is False, preserving middleware-set identity in 
request-scoped contexts.



##########
superset/mcp_service/mcp_config.py:
##########
@@ -318,11 +318,28 @@ def create_default_mcp_auth_factory(app: Flask) -> 
Optional[Any]:
         return None
 
 
-def default_user_resolver(app: Any, access_token: Any) -> Optional[str]:
-    """Extract username from JWT token claims."""
-    if hasattr(access_token, "subject"):
+def default_user_resolver(app: Any, access_token: Any) -> str | None:
+    """Extract username from JWT token claims.
+
+    Checks the ``claims`` dict first (FastMCP's AccessToken format),
+    then falls back to legacy attribute access for backward compatibility.
+    """
+    # FastMCP AccessToken stores JWT claims in a dict
+    claims = getattr(access_token, "claims", None)
+    if isinstance(claims, dict) and claims:
+        username = (
+            claims.get("sub")
+            or claims.get("preferred_username")
+            or claims.get("email")
+            or claims.get("username")

Review Comment:
   Agreed, fixed in 508869d. Reordered to `preferred_username > username > 
email > sub` since OIDC `sub` is often an opaque ID.



##########
tests/unit_tests/mcp_service/test_auth_user_resolution.py:
##########
@@ -0,0 +1,330 @@
+# 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 user resolution priority and stale g.user prevention."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from flask import g
+
+from superset.mcp_service.auth import (
+    _resolve_user_from_jwt_context,
+    get_user_from_request,
+    mcp_auth_hook,
+)
+from superset.mcp_service.mcp_config import default_user_resolver
+
+
+def _make_mock_user(username: str = "testuser") -> MagicMock:
+    """Create a mock User with required attributes."""
+    user = MagicMock()
+    user.username = username
+    user.roles = []
+    user.groups = []
+    return user
+
+
+def _make_access_token(
+    claims: dict[str, str] | None = None, **kwargs: str
+) -> MagicMock:
+    """Create a mock AccessToken matching FastMCP's format."""
+    token = MagicMock()
+    token.claims = claims or {}
+    token.client_id = kwargs.get("client_id", "")
+    token.scopes = kwargs.get("scopes", [])
+    # Remove auto-created attributes so getattr fallbacks work correctly
+    for attr in ("subject", "payload"):
+        if attr not in kwargs:
+            delattr(token, attr)
+    for attr in kwargs:
+        setattr(token, attr, kwargs[attr])
+    return token
+
+
+# -- _resolve_user_from_jwt_context --
+
+
+def test_jwt_context_resolves_correct_user(app) -> None:
+    """JWT context with valid claims resolves the correct DB user."""
+    mock_user = _make_mock_user("alice")
+    token = _make_access_token(claims={"sub": "alice"})
+
+    with app.app_context():
+        with (
+            patch("fastmcp.server.dependencies.get_access_token", 
return_value=token),
+            patch(
+                "superset.mcp_service.auth.load_user_with_relationships",
+                return_value=mock_user,
+            ),
+        ):
+            result = _resolve_user_from_jwt_context(app)
+
+    assert result is not None
+    assert result.username == "alice"
+
+
+def test_jwt_context_returns_none_when_no_token(app) -> None:
+    """No JWT token present returns None (fall through to next source)."""
+    with app.app_context():
+        with patch("fastmcp.server.dependencies.get_access_token", 
return_value=None):
+            result = _resolve_user_from_jwt_context(app)
+
+    assert result is None
+
+
+def test_jwt_context_raises_for_unknown_user(app) -> None:
+    """JWT resolves a username not in DB — raises ValueError (fail closed)."""
+    token = _make_access_token(claims={"sub": "nonexistent"})
+
+    with app.app_context():
+        with (
+            patch("fastmcp.server.dependencies.get_access_token", 
return_value=token),
+            patch(
+                "superset.mcp_service.auth.load_user_with_relationships",
+                return_value=None,
+            ),
+        ):
+            with pytest.raises(ValueError, match="not found in Superset 
database"):
+                _resolve_user_from_jwt_context(app)
+
+
+def test_jwt_context_returns_none_when_no_username_in_claims(app) -> None:
+    """JWT present but claims have no extractable username — returns None."""
+    token = _make_access_token(claims={"iss": "some-issuer"})
+
+    with app.app_context():
+        with patch("fastmcp.server.dependencies.get_access_token", 
return_value=token):
+            result = _resolve_user_from_jwt_context(app)
+
+    assert result is None
+
+
+def test_jwt_context_uses_custom_resolver(app) -> None:
+    """Custom MCP_USER_RESOLVER config is used when set."""
+    mock_user = _make_mock_user("custom_user")
+    token = _make_access_token(claims={"custom_field": "custom_user"})
+    custom_resolver = MagicMock(return_value="custom_user")
+
+    with app.app_context():
+        app.config["MCP_USER_RESOLVER"] = custom_resolver
+        try:
+            with (
+                patch(
+                    "fastmcp.server.dependencies.get_access_token", 
return_value=token
+                ),
+                patch(
+                    "superset.mcp_service.auth.load_user_with_relationships",
+                    return_value=mock_user,
+                ),
+            ):
+                result = _resolve_user_from_jwt_context(app)
+        finally:
+            app.config.pop("MCP_USER_RESOLVER", None)
+
+    assert result is not None
+    assert result.username == "custom_user"
+    custom_resolver.assert_called_once_with(app, token)
+
+
+# -- get_user_from_request priority order --
+
+
+def test_jwt_takes_priority_over_stale_g_user(app) -> None:
+    """Core regression test: JWT user wins over stale g.user."""
+    stale_user = _make_mock_user("stale_bob")
+    jwt_user = _make_mock_user("jwt_alice")
+    token = _make_access_token(claims={"sub": "jwt_alice"})
+
+    with app.app_context():
+        g.user = stale_user
+        with (
+            patch("fastmcp.server.dependencies.get_access_token", 
return_value=token),
+            patch(
+                "superset.mcp_service.auth.load_user_with_relationships",
+                return_value=jwt_user,
+            ),
+        ):
+            result = get_user_from_request()
+
+    assert result.username == "jwt_alice"
+
+
+def test_dev_username_fallback_when_no_jwt(app) -> None:
+    """MCP_DEV_USERNAME used when no JWT context available."""
+    mock_user = _make_mock_user("dev_admin")
+
+    with app.app_context():
+        app.config["MCP_DEV_USERNAME"] = "dev_admin"
+        try:
+            with (
+                patch(
+                    "fastmcp.server.dependencies.get_access_token", 
return_value=None
+                ),
+                patch(
+                    "superset.mcp_service.auth.load_user_with_relationships",
+                    return_value=mock_user,
+                ),
+            ):
+                result = get_user_from_request()
+        finally:
+            app.config.pop("MCP_DEV_USERNAME", None)
+
+    assert result.username == "dev_admin"
+
+
+def test_g_user_fallback_when_no_jwt_and_no_dev_username(app) -> None:
+    """g.user used as last-resort fallback (Preset middleware 
compatibility)."""
+    preset_user = _make_mock_user("preset_user")
+
+    with app.app_context():
+        app.config.pop("MCP_DEV_USERNAME", None)
+        g.user = preset_user
+        with patch("fastmcp.server.dependencies.get_access_token", 
return_value=None):
+            result = get_user_from_request()
+
+    assert result.username == "preset_user"
+
+
+def test_raises_when_no_auth_source(app) -> None:
+    """ValueError raised when no auth source is available."""
+    with app.app_context():
+        app.config.pop("MCP_DEV_USERNAME", None)
+        g.pop("user", None)
+        with patch("fastmcp.server.dependencies.get_access_token", 
return_value=None):
+            with pytest.raises(ValueError, match="No authenticated user 
found"):
+                get_user_from_request()
+
+
+def test_dev_username_not_found_raises(app) -> None:
+    """MCP_DEV_USERNAME configured but user not in DB raises ValueError."""
+    with app.app_context():
+        app.config["MCP_DEV_USERNAME"] = "ghost"
+        try:
+            with (
+                patch(
+                    "fastmcp.server.dependencies.get_access_token", 
return_value=None
+                ),
+                patch(
+                    "superset.mcp_service.auth.load_user_with_relationships",
+                    return_value=None,
+                ),
+            ):
+                with pytest.raises(ValueError, match="not found"):
+                    get_user_from_request()
+        finally:
+            app.config.pop("MCP_DEV_USERNAME", None)
+
+
+# -- g.user clearing in mcp_auth_hook --
+
+
+def test_mcp_auth_hook_clears_stale_g_user(app) -> None:
+    """mcp_auth_hook clears g.user before setting up user context."""
+    stale_user = _make_mock_user("stale")
+    fresh_user = _make_mock_user("fresh")
+
+    def dummy_tool():
+        """Dummy tool."""
+        return g.user.username
+
+    wrapped = mcp_auth_hook(dummy_tool)
+
+    with app.app_context():
+        g.user = stale_user
+        with patch(
+            "superset.mcp_service.auth.get_user_from_request",
+            return_value=fresh_user,
+        ):
+            result = wrapped()
+
+    assert result == "fresh"
+
+
+def test_mcp_auth_hook_clears_stale_g_user_async(app) -> None:
+    """mcp_auth_hook clears g.user before setting up user context (async)."""
+    import asyncio
+
+    stale_user = _make_mock_user("stale")
+    fresh_user = _make_mock_user("fresh")
+
+    async def dummy_tool():
+        """Dummy tool."""
+        return g.user.username
+
+    wrapped = mcp_auth_hook(dummy_tool)
+
+    with app.app_context():
+        g.user = stale_user
+        with patch(
+            "superset.mcp_service.auth.get_user_from_request",
+            return_value=fresh_user,
+        ):
+            result = asyncio.get_event_loop().run_until_complete(wrapped())

Review Comment:
   Fixed in 508869d. Switched to `asyncio.run()`.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to