import json import re import uuid from datetime import UTC, datetime from typing import Any from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.requests import Request from starlette.responses import Response from starlette.types import ASGIApp, Message, Receive, Scope, Send from app.core.config import settings as _settings _DEFAULT_REQUEST_BODY_LIMIT = 1024 * 1024 _OAUTH_REQUEST_BODY_LIMIT = 256 * 1024 _UPLOAD_REQUEST_BODY_LIMITS = ( (re.compile(r"^/api/v1/approval/asset/uploads/?$"), 60 * 1024 * 1024), (re.compile(r"^/api/v1/assets/(?:intake/bulk|bulk-update)/?$"), 6 * 1024 * 1024), (re.compile(r"^/api/v1/finance-flux/import/?$"), 21 * 1024 * 1024), (re.compile(r"^/api/v1/hr/employee-numbers/import/?$"), 6 * 1024 * 1024), ( re.compile(r"^/api/v1/auditor/vendor-assessments/[^/]+/import/?$"), 21 * 1024 * 1024, ), ( re.compile( r"^/api/v1/auditor/inbound-questionnaires/[^/]+/questions/import/?$" ), 21 * 1024 * 1024, ), ) class _RequestBodyTooLarge(BaseException): """Control-flow signal that broad ``except Exception`` blocks cannot swallow.""" def request_body_limit(path: str) -> int: if path.startswith("/api/v1/auth/oauth/"): return _OAUTH_REQUEST_BODY_LIMIT for pattern, limit in _UPLOAD_REQUEST_BODY_LIMITS: if pattern.fullmatch(path): return limit return _DEFAULT_REQUEST_BODY_LIMIT class RequestBodyLimitMiddleware: """Reject oversized fixed-length and chunked bodies before materialization.""" def __init__(self, app: ASGIApp) -> None: self.app = app async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope.get("type") != "http": await self.app(scope, receive, send) return path = scope.get("path", "") limit = request_body_limit(path) headers = {key.lower(): value for key, value in scope.get("headers", [])} raw_content_length = headers.get(b"content-length") if raw_content_length is not None: try: if int(raw_content_length) > limit: await self._send_too_large(scope, send, limit) return except ValueError: await self._send_bad_content_length(send) return received = 0 response_started = False async def bounded_receive() -> Message: nonlocal received message = await receive() if message.get("type") == "http.request": received += len(message.get("body", b"")) if received > limit: raise _RequestBodyTooLarge return message async def tracking_send(message: Message) -> None: nonlocal response_started if message.get("type") == "http.response.start": response_started = True await send(message) try: await self.app(scope, bounded_receive, tracking_send) except _RequestBodyTooLarge: if not response_started: await self._send_too_large(scope, send, limit) else: # HTTP status and headers are immutable once started. Abort the # stream instead of swallowing the signal and presenting a # truncated response as a successful application response. raise @staticmethod async def _send_bad_content_length(send: Send) -> None: payload = b'{"success":false,"errors":[{"code":"INVALID_CONTENT_LENGTH","message":"Invalid Content-Length header."}]}' await send({ "type": "http.response.start", "status": 400, "headers": [ (b"content-type", b"application/json"), (b"content-length", str(len(payload)).encode()), ], }) await send({"type": "http.response.body", "body": payload}) @staticmethod async def _send_too_large(scope: Scope, send: Send, limit: int) -> None: path = scope.get("path", "") content: dict[str, Any] if path.startswith("/api/v1/auth/oauth/"): content = { "error": "invalid_request", "error_description": f"request body exceeds the {limit}-byte limit", } else: trace_id = scope.get("state", {}).get("trace_id") content = { "success": False, "data": None, "meta": { "trace_id": trace_id, "timestamp": datetime.now(UTC).isoformat(), }, "errors": [{ "code": "PAYLOAD_TOO_LARGE", "message": f"Request body exceeds the {limit}-byte limit.", }], } payload = json.dumps(content, separators=(",", ":")).encode() await send({ "type": "http.response.start", "status": 413, "headers": [ (b"content-type", b"application/json"), (b"content-length", str(len(payload)).encode()), ], }) await send({"type": "http.response.body", "body": payload}) class TraceIdMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: raw_trace_id = request.headers.get("X-Trace-ID") try: trace_id = str(uuid.UUID(raw_trace_id)) if raw_trace_id else str(uuid.uuid4()) except (TypeError, ValueError): trace_id = str(uuid.uuid4()) request.state.trace_id = trace_id # Pre-launch P1 fix: also stamp the trace contextvar so the # leaver-blocked SQLAlchemy ``before_flush`` listener can # correlate its trap audit rows with the originating HTTP # request even when the listener fires on an unguarded ORM # write. from app.auth.auth_context import current_request_trace_id # noqa: PLC0415 current_request_trace_id.set(trace_id) # Issue-014: resolve real client IP (respects X-Forwarded-For / X-Real-IP) from app.core.dependencies import _get_real_ip request.state.client_ip = _get_real_ip(request) response = await call_next(request) response.headers["X-Trace-ID"] = trace_id return response class SecurityHeadersMiddleware(BaseHTTPMiddleware): def __init__(self, app: ASGIApp) -> None: # Keep the streaming body gate inside the existing middleware entry # so it runs after TraceIdMiddleware and its JSON errors retain trace # correlation and security headers. super().__init__(RequestBodyLimitMiddleware(app)) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: response = await call_next(request) response.headers["X-Content-Type-Options"] = "nosniff" response.headers["X-Frame-Options"] = "DENY" response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin" # ``default-src 'none'`` is right for JSON API responses but # makes the browser drop the inline