middleware.py

"""API-key authentication middleware for the MCP Office Documents server.

When an API key is configured (via the API_KEY environment variable), every
incoming request must carry a matching key in one of the supported HTTP
headers.  If no key is configured the middleware is not registered at all.

Supported header formats (checked in order):
  1. ``Authorization: Bearer <key>``
  2. ``Authorization: <key>``   (plain token, no scheme)
  3. ``x-api-key: <key>``
"""

from __future__ import annotations

import logging
import secrets
import time
from typing import Optional

from fastmcp.server.middleware import Middleware, MiddlewareContext
from fastmcp.server.dependencies import get_http_headers
from mcp import MCPError
from mcp.types import INVALID_REQUEST

logger = logging.getLogger(__name__)

#: What a client without a valid key is told. Stable, so a client can match it.
UNAUTHORIZED_MESSAGE = "Unauthorized – invalid or missing API key"


class ApiKeyAuthMiddleware(Middleware):
    """Reject requests that do not carry a valid API key.

    The middleware inspects the incoming HTTP headers for an API key and
    compares it against the *expected_key* provided at construction time.

    Look-up order:
      1. ``Authorization`` header with ``Bearer`` scheme
      2. ``Authorization`` header used as a plain token
      3. ``x-api-key`` header

    If the key is missing or does not match, the request is rejected with
    :data:`UNAUTHORIZED_MESSAGE` as an ``MCPError`` before any tool runs.

    It has to be an ``MCPError``, not FastMCP's ``AuthorizationError``: the MCP
    SDK (v2, FastMCP 4) treats any other exception escaping a request handler
    as a server crash. It logs a full traceback at ERROR for each one, which
    would defeat the throttled warning below, and on the 2026-07-28 protocol
    the client sees only "Internal server error" instead of the reason.
    """

    # Throttle window (seconds) for WARNING-level auth failure messages
    # to avoid log flooding from brute-force or misconfigured clients.
    _WARN_INTERVAL_SECONDS: float = 60.0

    def __init__(self, expected_key: str) -> None:
        if not expected_key:
            raise ValueError("ApiKeyAuthMiddleware requires a non-empty expected_key")
        self.expected_key = expected_key
        self._failed_attempts: int = 0
        # None means "no warning emitted yet". Do NOT use 0.0 as that sentinel:
        # it is compared against time.monotonic(), whose epoch is the host boot
        # on Linux, so on a freshly started pod "now - 0.0" is under the throttle
        # interval and the very first auth failure would be silenced for up to
        # _WARN_INTERVAL_SECONDS after boot — exactly when a misconfigured client
        # is most likely to be hammering a new deployment.
        self._last_warn_time: Optional[float] = None

    # ------------------------------------------------------------------
    # Key extraction
    # ------------------------------------------------------------------
    @staticmethod
    def _extract_key(headers: dict[str, str]) -> Optional[str]:
        """Return the API key from the first matching header, or *None*.

        Header names are compared case-insensitively (HTTP headers are
        case-insensitive per RFC 7230 §3.2).
        """
        # Normalize header names to lower-case for reliable look-up
        lower = {k.lower(): v for k, v in headers.items()}

        # 1. Authorization header
        auth = lower.get("authorization")
        if auth:
            # "Bearer <token>" – standard OAuth 2.0 scheme
            if auth.lower().startswith("bearer "):
                return auth[7:].strip()
            # Plain token (no scheme prefix)
            return auth.strip()

        # 2. x-api-key header (common API-gateway convention)
        api_key = lower.get("x-api-key")
        if api_key:
            return api_key.strip()

        return None

    # ------------------------------------------------------------------
    # Middleware hook
    # ------------------------------------------------------------------
    async def on_request(self, context: MiddlewareContext, call_next):
        # get_http_headers() strips "authorization" by default; include= opts it
        # back in so both Bearer and plain-token Authorization requests can
        # authenticate (without this, only x-api-key would ever be seen).
        headers = get_http_headers(include={"authorization"}) or {}
        api_key = self._extract_key(headers)

        if api_key is None or not secrets.compare_digest(api_key, self.expected_key):
            self._failed_attempts += 1
            now = time.monotonic()

            # Always log at DEBUG (cheap, only visible when debug is on)
            logger.debug("Auth failure on %s (attempt #%d)", context.method, self._failed_attempts)

            # Throttled WARNING: emit at most once per interval, and always on
            # the first failure regardless of how long the process has been up.
            if (self._last_warn_time is None
                    or now - self._last_warn_time >= self._WARN_INTERVAL_SECONDS):
                logger.warning(
                    "Rejected %d auth attempt(s) in the last %.0fs",
                    self._failed_attempts,
                    self._WARN_INTERVAL_SECONDS,
                )
                self._failed_attempts = 0
                self._last_warn_time = now

            raise MCPError(code=INVALID_REQUEST, message=UNAUTHORIZED_MESSAGE)

        return await call_next(context)