ai-flow-python/app/middleware/api_key.py

import logging
import time
from typing import Callable

from fastapi import Request, Response
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from app.config import settings

logger = logging.getLogger(__name__)


class ApiKeyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        method = request.method
        path = request.url.path
        client_ip = request.client.host if request.client else "unknown"

        t0 = time.time()

        if path == "/health":
            response = await call_next(request)
            return response

        logger.info("--> %s %s  (client=%s)", method, path, client_ip)

        if settings.api_auth_enabled and path != "/health":
            header_value = request.headers.get(settings.api_auth_header_name, "")
            prefix = settings.api_auth_prefix
            if prefix and header_value.startswith(prefix + " "):
                token = header_value[len(prefix) + 1:]
            else:
                token = header_value

            if token not in settings.api_auth_keys:
                elapsed = time.time() - t0
                logger.warning("<-- %s %s  401 Unauthorized  (%.3fs, client=%s)",
                                method, path, elapsed, client_ip)
                return JSONResponse(
                    status_code=401,
                    content={"error": settings.api_auth_error_message},
                )

        response = await call_next(request)
        elapsed = time.time() - t0
        logger.info("<-- %s %s  %d  (%.3fs)", method, path, response.status_code, elapsed)
        return response