homelable/mcp/app/auth.py

import hmac
import json
from starlette.types import ASGIApp, Receive, Scope, Send

from .config import settings

_BYPASS_PATHS = {"/health", "/register"}


class ApiKeyMiddleware:
    """Pure ASGI middleware — compatible with SSE/streaming responses.

    BaseHTTPMiddleware buffers the full response body and breaks SSE streams.
    This implementation operates at the ASGI scope level and never touches
    the response stream.
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path: str = scope.get("path", "")

        if path in _BYPASS_PATHS or path.startswith("/.well-known/"):
            await self.app(scope, receive, send)
            return

        headers = dict(scope.get("headers", []))
        key = headers.get(b"x-api-key", b"").decode()
        expected = settings.mcp_api_key

        if not key or not hmac.compare_digest(key.encode(), expected.encode()):
            body = json.dumps({"detail": "Invalid or missing X-API-Key"}).encode()
            await send({"type": "http.response.start", "status": 401,
                        "headers": [(b"content-type", b"application/json"),
                                    (b"content-length", str(len(body)).encode())]})
            await send({"type": "http.response.body", "body": body, "more_body": False})
            return

        await self.app(scope, receive, send)