devplacepy/devplacepy_services/base/middleware.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
import time
import uuid_utils
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response

from devplacepy_services.base.auth import validate_internal_key
from devplacepy_services.base.errors import error_response, sanitize_detail
from devplacepy_services.base.http import INTERNAL_CALLS, REQUEST_ID, current_internal_calls
from devplacepy_services.base.stats import COUNTERS

logger = logging.getLogger(__name__)

PUBLIC_PATHS = {"/health"}


class RequestStatsMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        COUNTERS.record_request()
        response = await call_next(request)
        if response.status_code >= 500:
            COUNTERS.record_error()
        return response


class InternalCallCounterMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        token_calls = INTERNAL_CALLS.set(0)
        response = await call_next(request)
        service_name = getattr(request.app.state, "service_name", "")
        if service_name in {"web", "gateway"}:
            response.headers["X-Internal-Calls"] = str(current_internal_calls())
        INTERNAL_CALLS.reset(token_calls)
        return response


class InternalAuthMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        path = request.url.path
        if path in PUBLIC_PATHS:
            return await call_next(request)
        if not validate_internal_key(request.headers.get("X-Internal-Key")):
            return error_response(401, "Unauthorized", "unauthorized")
        return await call_next(request)


class SanitizeErrorsMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        request_id = request.headers.get("X-Request-Id") or uuid_utils.uuid7().hex
        token_id = REQUEST_ID.set(request_id)
        started = time.perf_counter()
        try:
            response = await call_next(request)
        except Exception:
            COUNTERS.record_error()
            logger.exception(
                "unhandled service error",
                extra={
                    "request_id": request_id,
                    "service": getattr(request.app.state, "service_name", ""),
                    "path": request.url.path,
                },
            )
            response = error_response(500, "Internal server error", "internal_error")
        finally:
            REQUEST_ID.reset(token_id)
        duration_ms = int((time.perf_counter() - started) * 1000)
        service_name = getattr(request.app.state, "service_name", "")
        logger.info(
            "request",
            extra={
                "request_id": request_id,
                "service": service_name,
                "path": request.url.path,
                "status": response.status_code,
                "duration_ms": duration_ms,
                "internal_calls": current_internal_calls(),
            },
        )
        if response.status_code < 400:
            return response
        if not hasattr(response, "body_iterator"):
            return response
        content_type = response.headers.get("content-type", "")
        if "application/json" not in content_type:
            if response.status_code == 404:
                return error_response(404, "Not found", "not_found")
            if response.status_code == 401:
                return error_response(401, "Unauthorized", "unauthorized")
            if response.status_code >= 500:
                return error_response(500, "Internal server error", "internal_error")
            return error_response(response.status_code, "Request failed", "error")
        body = b""
        async for chunk in response.body_iterator:
            body += chunk
        try:
            parsed = json.loads(body)
        except json.JSONDecodeError:
            return error_response(response.status_code, "Request failed", "error")
        if isinstance(parsed, dict) and "error" in parsed and "code" in parsed:
            return Response(
                content=body,
                status_code=response.status_code,
                media_type="application/json",
            )
        if isinstance(parsed, dict) and isinstance(parsed.get("error"), dict) and "message" in parsed["error"]:
            return Response(
                content=body,
                status_code=response.status_code,
                media_type="application/json",
            )
        if isinstance(parsed, dict) and "detail" in parsed:
            message, code = sanitize_detail(parsed["detail"], response.status_code)
            return error_response(response.status_code, message, code)
        return error_response(response.status_code, "Request failed", "error")