src/spotify_mcp_server/server_auth.py

"""Scalekit authentication and hosted runtime configuration."""

from __future__ import annotations

import asyncio
import contextvars
import logging
import os
import re
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, replace
from typing import Any, Protocol
from urllib.parse import urlparse

import jwt
from mcp.server.auth.provider import AccessToken
from mcp.server.auth.settings import AuthSettings
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
from mcp.server.transport_security import TransportSecuritySettings

from spotify_mcp_server.spotify.config import Settings

logger = logging.getLogger(__name__)
_subject_context: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "scalekit_subject", default=None
)

# Scalekit signs with RS256. Pinning the list here keeps an attacker from selecting the
# algorithm, and lets the JWKS client reject a token before it can influence a key fetch.
_ALLOWED_ALGORITHMS: tuple[str, ...] = ("RS256",)

# Shortest gap between two cache-bypassing JWKS refreshes. Long enough that forged key ids
# cannot amplify traffic to Scalekit, short enough that a genuine key rotation is picked up
# on the next request rather than waiting out the JWK set cache.
_JWKS_MIN_REFRESH_INTERVAL = 30.0


class ScalekitClaimsClient(Protocol):
    def validate_access_token_and_get_claims(self, token: str, options: Any) -> dict[str, Any]: ...


@dataclass(frozen=True, slots=True)
class HostedSettings:
    """Required environment for the authenticated public deployment."""

    scalekit_environment_url: str
    scalekit_resource_id: str
    mcp_server_url: str
    database_url: str
    token_encryption_key: str
    spotify: Settings
    allowed_subjects: frozenset[str]
    host: str = "0.0.0.0"
    port: int = 7860

    @classmethod
    def from_env(cls) -> HostedSettings:
        environment_url = _required("SCALEKIT_ENVIRONMENT_URL").rstrip("/")
        resource_id = _required("SCALEKIT_RESOURCE_ID")
        server_url = _required("MCP_SERVER_URL").rstrip("/")
        _validate_hosted_urls(environment_url, server_url, resource_id)

        parsed_server = urlparse(server_url)
        callback_url = f"{parsed_server.scheme}://{parsed_server.netloc}/spotify/callback"
        spotify = replace(
            Settings.from_env(),
            client_id=_required("SPOTIFY_CLIENT_ID"),
            redirect_uri=callback_url,
        )
        allowed = frozenset(
            value.strip()
            for value in os.environ.get("MCP_ALLOWED_SUBJECTS", "").split(",")
            if value.strip()
        )
        return cls(
            scalekit_environment_url=environment_url,
            scalekit_resource_id=resource_id,
            mcp_server_url=server_url,
            database_url=_required("DATABASE_URL"),
            token_encryption_key=_required("TOKEN_ENCRYPTION_KEY"),
            spotify=spotify,
            allowed_subjects=allowed,
            port=int(os.environ.get("MCP_PORT", os.environ.get("PORT", "7860"))),
        )

    @property
    def authorization_server_url(self) -> str:
        return f"{self.scalekit_environment_url}/resources/{self.scalekit_resource_id}"

    @property
    def public_origin(self) -> str:
        parsed = urlparse(self.mcp_server_url)
        return f"{parsed.scheme}://{parsed.netloc}"

    def auth_settings(self) -> AuthSettings:
        return AuthSettings(
            issuer_url=self.authorization_server_url,
            resource_server_url=self.mcp_server_url,
            required_scopes=[],
        )

    def transport_security(self) -> TransportSecuritySettings:
        hostname = urlparse(self.mcp_server_url).hostname
        assert hostname is not None
        return TransportSecuritySettings(
            enable_dns_rebinding_protection=True,
            allowed_hosts=[hostname, f"{hostname}:*", "127.0.0.1:*", "localhost:*"],
            allowed_origins=[self.public_origin],
        )


class ScalekitTokenVerifier:
    """Validate Scalekit JWTs and expose their stable user subject to MCP handlers."""

    def __init__(
        self,
        settings: HostedSettings,
        *,
        client: ScalekitClaimsClient | None = None,
    ) -> None:
        self.settings = settings
        if client is None:
            self.client = _ScalekitJWTClaimsClient(settings.scalekit_environment_url)
            self.options = _ValidationOptions(
                issuer=settings.scalekit_environment_url,
                audience=[settings.mcp_server_url],
            )
        else:
            self.client = client
            self.options = _ValidationOptions(
                issuer=settings.scalekit_environment_url,
                audience=[settings.mcp_server_url],
            )

    async def verify_token(self, token: str) -> AccessToken | None:
        try:
            claims = await asyncio.to_thread(
                self.client.validate_access_token_and_get_claims,
                token,
                self.options,
            )
            subject = claims.get("sub")
            if not isinstance(subject, str) or not subject:
                return None
            return AccessToken(
                token=token,
                client_id=_client_id(claims),
                scopes=_scopes(claims),
                expires_at=_integer_claim(claims, "exp"),
                resource=self.settings.mcp_server_url,
                subject=subject,
                claims=claims,
            )
        except Exception:
            logger.warning("Scalekit rejected an MCP bearer token")
            return None


class AuthenticatedSubjectMiddleware:
    """Make the authenticated subject available to static resource handlers."""

    async def __call__(
        self,
        ctx: ServerRequestContext[Any, Any],
        call_next: CallNext,
    ) -> HandlerResult:
        request = ctx.request
        user = getattr(request, "user", None)
        access_token = getattr(user, "access_token", None)
        subject = getattr(access_token, "subject", None)
        token = _subject_context.set(subject if isinstance(subject, str) else None)
        try:
            return await call_next(ctx)
        finally:
            _subject_context.reset(token)


def current_authenticated_subject() -> str:
    subject = _subject_context.get()
    if not subject:
        raise RuntimeError("Authenticated Scalekit user subject is unavailable")
    return subject


@dataclass(frozen=True, slots=True)
class _ValidationOptions:
    issuer: str
    audience: list[str]
    required_scopes: list[str] | None = None


class _ThrottledJWKSClient:
    """Resolve Scalekit signing keys while rate-limiting cache-bypassing JWKS refreshes.

    `PyJWKClient.get_signing_key_from_jwt` refetches the whole JWK set whenever a token's
    `kid` misses the cache, and that happens before any signature is checked. Left alone it
    turns one unauthenticated request into one outbound request to Scalekit, each holding a
    worker thread for the duration of a blocking fetch. Refreshes are throttled instead, so a
    stream of forged `kid`s cannot amplify traffic to the issuer or starve the thread pool.
    """

    def __init__(
        self,
        uri: str,
        *,
        client: Any | None = None,
        min_refresh_interval: float = _JWKS_MIN_REFRESH_INTERVAL,
        now: Callable[[], float] = time.monotonic,
    ) -> None:
        self._client = client or jwt.PyJWKClient(uri)
        self._min_refresh_interval = min_refresh_interval
        self._now = now
        # verify_token runs under asyncio.to_thread, so refreshes race across worker threads.
        self._lock = threading.Lock()
        self._last_refresh: float | None = None

    def get_signing_key_from_jwt(self, token: str) -> Any:
        header = jwt.get_unverified_header(token)
        algorithm = header.get("alg")
        if algorithm not in _ALLOWED_ALGORITHMS:
            raise jwt.InvalidAlgorithmError(f"unsupported token algorithm: {algorithm!r}")
        kid = header.get("kid")
        if not isinstance(kid, str) or not kid:
            # Scalekit always publishes a key id; without one there is nothing to match, and
            # falling through would spend a refresh on a token that can never resolve.
            raise jwt.PyJWKClientError("token header is missing a key id")

        key = self._client.match_kid(self._client.get_signing_keys(), kid)
        if key is not None:
            return key
        if not self._claim_refresh():
            raise jwt.PyJWKClientError(f'Unable to find a signing key that matches: "{kid}"')

        key = self._client.match_kid(self._client.get_signing_keys(refresh=True), kid)
        if key is None:
            raise jwt.PyJWKClientError(f'Unable to find a signing key that matches: "{kid}"')
        return key

    def _claim_refresh(self) -> bool:
        """Take the right to refresh, or report that another refresh happened too recently."""

        with self._lock:
            now = self._now()
            last = self._last_refresh
            if last is not None and now - last < self._min_refresh_interval:
                return False
            self._last_refresh = now
            return True


class _ScalekitJWTClaimsClient:
    """Validate ScaleKit access tokens against its public signing keys."""

    def __init__(self, environment_url: str, *, jwks_client: Any | None = None) -> None:
        self.jwks_client = jwks_client or _ThrottledJWKSClient(f"{environment_url}/keys")

    def validate_access_token_and_get_claims(
        self,
        token: str,
        options: _ValidationOptions,
    ) -> dict[str, Any]:
        signing_key = self.jwks_client.get_signing_key_from_jwt(token)
        claims = jwt.decode(
            token,
            key=signing_key.key,
            algorithms=list(_ALLOWED_ALGORITHMS),
            issuer=options.issuer,
            audience=options.audience,
            options={"require": ["exp", "iss", "sub", "aud"]},
        )
        if options.required_scopes:
            missing = set(options.required_scopes) - set(_scopes(claims))
            if missing:
                raise jwt.InvalidTokenError("required scope is missing")
        return claims


def _required(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        raise ValueError(f"{name} is required in hosted mode")
    return value


def _validate_hosted_urls(environment_url: str, server_url: str, resource_id: str) -> None:
    for name, value in (
        ("SCALEKIT_ENVIRONMENT_URL", environment_url),
        ("MCP_SERVER_URL", server_url),
    ):
        parsed = urlparse(value)
        if parsed.scheme != "https" or not parsed.hostname or parsed.query or parsed.fragment:
            raise ValueError(f"{name} must be a clean HTTPS URL")
    if urlparse(server_url).path != "/mcp":
        raise ValueError("MCP_SERVER_URL must end with /mcp")
    if not re.fullmatch(r"res_[A-Za-z0-9]+", resource_id):
        raise ValueError("SCALEKIT_RESOURCE_ID is invalid")


def _client_id(claims: dict[str, Any]) -> str:
    for name in ("client_id", "azp"):
        value = claims.get(name)
        if isinstance(value, str) and value:
            return value
    return "scalekit-mcp-client"


def _scopes(claims: dict[str, Any]) -> list[str]:
    scopes = claims.get("scopes")
    if isinstance(scopes, list):
        return [scope for scope in scopes if isinstance(scope, str) and scope]
    scope = claims.get("scope")
    if isinstance(scope, str):
        return [value for value in scope.split() if value]
    return []


def _integer_claim(claims: dict[str, Any], name: str) -> int | None:
    value = claims.get(name)
    return value if isinstance(value, int) else None