next-tjwater-server/app/mcp/auth.py

"""Verify resource-bound Keycloak tokens without passing them to business APIs."""

import asyncio
import time
from uuid import UUID

import httpx
from jose import JWTError, jwt
from mcp.server.auth.provider import AccessToken
from starlette.responses import JSONResponse

from app.mcp.config import McpSettings

CATALOG_SCOPE = "mcp.catalog.read"
QUERY_SCOPE = "mcp.query"
CONTEXT_SCOPE = "mcp.context.bind"


class McpTokenVerifier:
    def __init__(self, settings: McpSettings, client: httpx.AsyncClient | None = None):
        self.settings = settings
        self.client = client
        self._keys: list[dict] = []
        self._expires = 0.0
        self._last_refresh = 0.0
        self._lock = asyncio.Lock()

    async def _load_keys(self, kid: str) -> list[dict]:
        now = time.monotonic()
        if self._expires > now and any(k.get("kid") == kid for k in self._keys):
            return self._keys
        async with self._lock:
            now = time.monotonic()
            if self._expires > now and any(k.get("kid") == kid for k in self._keys):
                return self._keys
            if self._expires > now and now - self._last_refresh < 20:
                return self._keys
            url = (
                self.settings.issuer_uri.rstrip("/") + "/protocol/openid-connect/certs"
            )
            self._last_refresh = now
            if self.client:
                response = await self.client.get(url)
            else:
                async with httpx.AsyncClient(
                    timeout=5, trust_env=False, follow_redirects=False
                ) as client:
                    response = await client.get(url)
            response.raise_for_status()
            try:
                payload = response.json()
            except ValueError:
                raise httpx.RequestError("invalid OAuth key set response") from None
            keys = payload.get("keys") if isinstance(payload, dict) else None
            if not isinstance(keys, list) or len(keys) > 100:
                raise httpx.RequestError("invalid OAuth key set")
            self._keys = [
                k
                for k in keys
                if isinstance(k, dict)
                and k.get("kty") == "RSA"
                and k.get("use", "sig") == "sig"
                and k.get("alg", "RS256") == "RS256"
            ]
            self._expires = now + 300
            return self._keys

    async def verify_token(self, token: str) -> AccessToken | None:
        if not token or len(token) > 16384:
            return None
        try:
            header = jwt.get_unverified_header(token)
            kid = header.get("kid")
            if (
                header.get("alg") != "RS256"
                or not isinstance(kid, str)
                or not 0 < len(kid) <= 128
            ):
                return None
            keys = await self._load_keys(kid)
            key = next((k for k in keys if k.get("kid") == kid), None)
            if key is None:
                return None
            claims = jwt.decode(
                token,
                key,
                algorithms=["RS256"],
                audience=self.settings.audience,
                issuer=self.settings.issuer_uri,
                options={"require_exp": True, "require_iat": True, "require_sub": True},
            )
            UUID(claims["sub"])
            client_id = claims.get("azp")
            if not isinstance(client_id, str) or not client_id:
                return None
            if (
                not isinstance(claims.get("iat"), (int, float))
                or not time.time() - 900 < claims["iat"] <= time.time() + 30
                or claims["exp"] > claims["iat"] + 900
            ):
                return None
            return AccessToken(
                token=token,
                client_id=client_id,
                subject=claims["sub"],
                scopes=str(claims.get("scope", "")).split(),
                expires_at=int(claims["exp"]),
                resource=self.settings.resource_uri,
                claims=claims,
            )
        except (JWTError, ValueError, KeyError, TypeError):
            return None
        # JWKS/network failures propagate as unavailable, never as empty data.


def challenge(
    settings: McpSettings, scope: str, *, insufficient: bool = False
) -> dict[str, str]:
    from mcp.server.auth.routes import build_resource_metadata_url

    metadata = build_resource_metadata_url(settings.resource_uri)
    suffix = ', error="insufficient_scope"' if insufficient else ""
    return {
        "WWW-Authenticate": f'Bearer resource_metadata="{metadata}", scope="{scope}"{suffix}'
    }


class QueryScopeMiddleware:
    """HTTP scope challenges before JSON-RPC dispatch; SDK validates mirrored headers."""

    def __init__(self, app, verifier: McpTokenVerifier, settings: McpSettings):
        self.app, self.verifier, self.settings = app, verifier, settings

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http" or scope.get("path") != "/mcp":
            return await self.app(scope, receive, send)
        headers = {
            k.decode("latin1").lower(): v.decode("latin1")
            for k, v in scope.get("headers", [])
        }
        auth = headers.get("authorization", "")
        try:
            token = (
                await self.verifier.verify_token(auth[7:])
                if auth.lower().startswith("bearer ")
                else None
            )
        except (httpx.HTTPError, ValueError):
            return await JSONResponse(
                {"error": "authentication_unavailable"}, status_code=503
            )(scope, receive, send)
        required = (
            QUERY_SCOPE
            if headers.get("mcp-method") in {"tools/call", "resources/read"}
            else CATALOG_SCOPE
        )
        if token is None:
            return await JSONResponse(
                {"error": "invalid_token"},
                status_code=401,
                headers=challenge(self.settings, required),
            )(scope, receive, send)
        if required not in token.scopes:
            return await JSONResponse(
                {"error": "insufficient_scope"},
                status_code=403,
                headers=challenge(self.settings, required, insufficient=True),
            )(scope, receive, send)
        await self.app(scope, receive, send)