codexLB-up 1.25.0b10__py3-none-any.whl

app/core/middleware/request_decompression.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import gzip
+import io
+import zlib
+from typing import Protocol
+
+import zstandard as zstd
+from fastapi import FastAPI
+from starlette._utils import get_route_path
+from starlette.datastructures import Headers
+from starlette.requests import ClientDisconnect, Request
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from app.core.middleware.request_body_limit import (
+    REQUEST_BODY_TOO_LARGE_MESSAGE,
+    request_body_limit_for_path,
+    request_ingress_error_response,
+)
+
+
+class _DecompressedBodyTooLarge(Exception):
+    def __init__(self, max_size: int) -> None:
+        super().__init__(f"Decompressed body exceeded {max_size} bytes")
+        self.max_size = max_size
+
+
+class _Readable(Protocol):
+    def read(self, size: int = ...) -> bytes: ...
+
+
+def _read_limited(reader: _Readable, max_size: int) -> bytes:
+    buffer = bytearray()
+    total = 0
+    chunk_size = 64 * 1024
+    while True:
+        chunk = reader.read(chunk_size)
+        if not chunk:
+            break
+        total += len(chunk)
+        if total > max_size:
+            raise _DecompressedBodyTooLarge(max_size)
+        buffer.extend(chunk)
+    return bytes(buffer)
+
+
+def _decompress_gzip(data: bytes, max_size: int) -> bytes:
+    with gzip.GzipFile(fileobj=io.BytesIO(data)) as reader:
+        return _read_limited(reader, max_size)
+
+
+def _decompress_deflate(data: bytes, max_size: int) -> bytes:
+    decompressor = zlib.decompressobj()
+    buffer = bytearray()
+    chunk_size = 64 * 1024
+    for start in range(0, len(data), chunk_size):
+        chunk = data[start : start + chunk_size]
+        # Bound output growth to avoid oversized allocations.
+        while chunk:
+            remaining = max_size - len(buffer)
+            decompressed = decompressor.decompress(chunk, max_length=remaining + 1)
+            if len(decompressed) > remaining:
+                raise _DecompressedBodyTooLarge(max_size)
+            buffer.extend(decompressed)
+            chunk = decompressor.unconsumed_tail
+    while True:
+        remaining = max_size - len(buffer)
+        drained = decompressor.decompress(b"", max_length=remaining + 1)
+        if len(drained) > remaining:
+            raise _DecompressedBodyTooLarge(max_size)
+        if not drained:
+            break
+        buffer.extend(drained)
+    if not decompressor.eof:
+        raise zlib.error("Incomplete deflate stream")
+    return bytes(buffer)
+
+
+def _decompress_zstd(data: bytes, max_size: int) -> bytes:
+    with zstd.ZstdDecompressor().stream_reader(io.BytesIO(data)) as reader:
+        return _read_limited(reader, max_size)
+
+
+def _decompress_body(data: bytes, encodings: list[str], max_size: int) -> bytes:
+    supported = {"zstd", "gzip", "deflate", "identity"}
+    if any(encoding not in supported for encoding in encodings):
+        raise ValueError("Unsupported content-encoding")
+    result = data
+    for encoding in reversed(encodings):
+        if encoding == "zstd":
+            result = _decompress_zstd(result, max_size)
+        elif encoding == "gzip":
+            result = _decompress_gzip(result, max_size)
+        elif encoding == "deflate":
+            result = _decompress_deflate(result, max_size)
+        elif encoding == "identity":
+            pass
+        if len(result) > max_size:
+            raise _DecompressedBodyTooLarge(max_size)
+    return result
+
+
+def _rewrite_scope_headers_for_body(scope: Scope, body_length: int) -> None:
+    """Drop content-encoding/content-length and declare the decompressed length."""
+    headers: list[tuple[bytes, bytes]] = []
+    for key, value in scope.get("headers", []):
+        if key.lower() in (b"content-encoding", b"content-length"):
+            continue
+        headers.append((key, value))
+    headers.append((b"content-length", str(body_length).encode("ascii")))
+    scope["headers"] = headers
+
+
+async def _drain_request_body(receive: Receive) -> bytes:
+    """Read the full request body from ``receive``.
+
+    Mirrors ``Request.body()`` semantics: a mid-body ``http.disconnect`` raises
+    ``ClientDisconnect``, and receive failures propagate. The caller's
+    ``receive`` is the body-limit middleware's limited receive, so the wire-size
+    cap (``_RequestBodyTooLarge``) propagates through this drain unchanged.
+    """
+    chunks = bytearray()
+    while True:
+        message = await receive()
+        if message["type"] == "http.disconnect":
+            raise ClientDisconnect()
+        chunks.extend(message.get("body", b""))
+        if not message.get("more_body", False):
+            return bytes(chunks)
+
+
+class RequestDecompressionMiddleware:
+    """Decompress zstd/gzip/deflate request bodies with per-layer decode budgets.
+
+    Pure ASGI replacement for the previous ``BaseHTTPMiddleware`` dispatch:
+    requests without ``Content-Encoding`` (the common case) pass straight
+    through. Encoded requests are drained through the upstream receive (the
+    body-limit middleware's limited receive, so the wire-size cap still applies
+    to the compressed bytes), decompressed under the per-path budget, and the
+    downstream app is given a replay receive that yields the decompressed body
+    once and then delegates to the original receive so ``http.disconnect`` is
+    still observed (this replaces the ``_CachedRequest`` body replay the
+    BaseHTTP wrapper used to provide).
+    """
+
+    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
+
+        content_encoding = Headers(scope=scope).get("content-encoding")
+        if not content_encoding:
+            await self.app(scope, receive, send)
+            return
+        encodings = [enc.strip().lower() for enc in content_encoding.split(",") if enc.strip()]
+        if not encodings:
+            await self.app(scope, receive, send)
+            return
+
+        max_size = request_body_limit_for_path(get_route_path(scope))
+        body = await _drain_request_body(receive)
+        try:
+            decompressed = _decompress_body(body, encodings, max_size)
+        except _DecompressedBodyTooLarge:
+            response = request_ingress_error_response(
+                Request(scope),
+                status_code=413,
+                code="payload_too_large",
+                message=REQUEST_BODY_TOO_LARGE_MESSAGE,
+            )
+            await response(scope, receive, send)
+            return
+        except ValueError:
+            response = request_ingress_error_response(
+                Request(scope),
+                status_code=400,
+                code="invalid_request",
+                message="Unsupported Content-Encoding",
+            )
+            await response(scope, receive, send)
+            return
+        except Exception:
+            response = request_ingress_error_response(
+                Request(scope),
+                status_code=400,
+                code="invalid_request",
+                message="Request body is compressed but could not be decompressed",
+            )
+            await response(scope, receive, send)
+            return
+
+        _rewrite_scope_headers_for_body(scope, len(decompressed))
+        body_replayed = False
+
+        async def replay_receive() -> Message:
+            nonlocal body_replayed
+            if not body_replayed:
+                body_replayed = True
+                return {"type": "http.request", "body": decompressed, "more_body": False}
+            return await receive()
+
+        await self.app(scope, replay_receive, send)
+
+
+def add_request_decompression_middleware(app: FastAPI) -> None:
+    app.add_middleware(RequestDecompressionMiddleware)

app/core/middleware/request_id.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from uuid import uuid4
+
+from fastapi import FastAPI
+from starlette.datastructures import MutableHeaders
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from app.core.utils.request_id import (
+    clear_request_id,
+    clear_request_scope_id,
+    reset_request_id,
+    reset_request_scope_id,
+    set_request_id,
+    set_request_scope_id,
+)
+
+
+def _inbound_request_id(scope: Scope) -> str | None:
+    """Return the first ``x-request-id`` value, else the first ``request-id`` value."""
+    x_request_id: bytes | None = None
+    request_id: bytes | None = None
+    for name, value in scope.get("headers", []):
+        lowered = name.lower()
+        if x_request_id is None and lowered == b"x-request-id":
+            x_request_id = value
+        elif request_id is None and lowered == b"request-id":
+            request_id = value
+    inbound = x_request_id or request_id
+    if not inbound:
+        return None
+    return inbound.decode("latin-1")
+
+
+class RequestIdMiddleware:
+    """Bind request-id contextvars around the request and echo ``x-request-id``.
+
+    Pure ASGI: the contextvars are set and reset in the same task that runs the
+    downstream app, so they stay visible for the whole response stream (with
+    ``BaseHTTPMiddleware`` they were reset when the dispatch returned, before
+    the body finished streaming).
+    """
+
+    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
+
+        request_id = _inbound_request_id(scope) or str(uuid4())
+        request_id_token = set_request_id(request_id)
+        request_scope_token = set_request_scope_id(str(uuid4()))
+
+        async def send_with_request_id(message: Message) -> None:
+            if message["type"] == "http.response.start":
+                headers = MutableHeaders(raw=message.setdefault("headers", []))
+                headers.setdefault("x-request-id", request_id)
+            await send(message)
+
+        try:
+            await self.app(scope, receive, send_with_request_id)
+        finally:
+            reset_request_scope_id(request_scope_token)
+            reset_request_id(request_id_token)
+            clear_request_scope_id()
+            clear_request_id()
+
+
+def add_request_id_middleware(app: FastAPI) -> None:
+    app.add_middleware(RequestIdMiddleware)

app/core/middleware/required_capability_http.py

@@ -0,0 +1,116 @@
+from __future__ import annotations
+
+import logging
+import time
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from starlette._utils import get_route_path
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from app.core.auth.dependencies import validate_required_proxy_api_key_authorization
+from app.core.clients.proxy import CODEX_LB_REQUIRED_CAPABILITY_HEADER
+from app.core.errors import openai_error
+from app.core.exceptions import ProxyAuthError, ProxyRequiredCapabilityTransportError
+from app.core.runtime_logging import log_error_response
+from app.modules.proxy.images_observability import (
+    IMAGE_ROUTE_STARTED_AT_STATE,
+    record_images_route_observability,
+)
+
+logger = logging.getLogger(__name__)
+
+_REQUIRED_CAPABILITY_HEADER_BYTES = CODEX_LB_REQUIRED_CAPABILITY_HEADER.lower().encode("latin-1")
+
+_JSON_BODY_DENY_PATHS = frozenset(
+    {
+        "/backend-api/codex/responses",
+        "/backend-api/codex/responses/compact",
+        "/backend-api/codex/images/generations",
+        "/v1/responses",
+        "/v1/responses/compact",
+        "/v1/chat/completions",
+        "/v1/embeddings",
+        "/v1/images/generations",
+        "/v1/reset-credit",
+        "/v1/warmup",
+        "/api/codex/rate-limit-reset-credits/consume",
+    }
+)
+
+
+def _is_pre_body_deny_path(path: str) -> bool:
+    normalized = path.rstrip("/")
+    return normalized in _JSON_BODY_DENY_PATHS or normalized.startswith("/v1/warmup/")
+
+
+def _has_required_capability_header(scope: Scope) -> bool:
+    for name, _value in scope.get("headers", []):
+        if name.lower() == _REQUIRED_CAPABILITY_HEADER_BYTES:
+            return True
+    return False
+
+
+class RequiredCapabilityHttpMiddleware:
+    """Deny capability-marked POSTs on JSON-body proxy paths before the body is read.
+
+    Pure ASGI with cheap synchronous guards first; the ``Request`` object is
+    only constructed on the cold deny path.
+    """
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if (
+            scope["type"] != "http"
+            or scope["method"] != "POST"
+            or not _has_required_capability_header(scope)
+            or not _is_pre_body_deny_path(get_route_path(scope))
+        ):
+            await self.app(scope, receive, send)
+            return
+
+        request = Request(scope)
+        try:
+            await validate_required_proxy_api_key_authorization(request.headers.get("authorization"))
+        except ProxyAuthError as exc:
+            response = _capability_error_response(request, exc)
+        else:
+            response = _capability_error_response(request, ProxyRequiredCapabilityTransportError())
+        await response(scope, receive, send)
+
+
+def add_required_capability_http_middleware(app: FastAPI) -> None:
+    app.add_middleware(RequiredCapabilityHttpMiddleware)
+
+
+def _capability_error_response(
+    request: Request,
+    exc: ProxyAuthError | ProxyRequiredCapabilityTransportError,
+) -> JSONResponse:
+    log_error_response(
+        logger,
+        request,
+        exc.status_code,
+        exc.code,
+        exc.message,
+        category="openai_error_response",
+    )
+    path = get_route_path(request.scope)
+    if path.rstrip("/").endswith("/images/generations"):
+        started_at = getattr(request.state, IMAGE_ROUTE_STARTED_AT_STATE, None)
+        if not isinstance(started_at, float):
+            started_at = time.perf_counter()
+        record_images_route_observability(
+            route="generations",
+            model=None,
+            stream=False,
+            status=exc.status_code,
+            outcome="auth_error" if isinstance(exc, ProxyAuthError) else "invalid_request",
+            started_at=started_at,
+        )
+    return JSONResponse(
+        status_code=exc.status_code,
+        content=openai_error(exc.code, exc.message, error_type=exc.error_type),
+    )

app/core/middleware/trusted_proxy_headers.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TypeAlias, cast
+
+from fastapi import FastAPI
+from starlette.types import ASGIApp, Receive, Scope, Send
+from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
+
+from app.core.config.settings import get_settings
+from app.core.socket_peer import _capture_raw_socket_peer
+
+_UvicornASGIApp: TypeAlias = Callable[..., Awaitable[None]]
+
+
+class TrustedProxyHeadersMiddleware:
+    """Preserve the raw socket peer, then apply Uvicorn's proxy projection."""
+
+    def __init__(self, app: ASGIApp) -> None:
+        trusted_hosts = get_settings().forwarded_allow_ips
+        self._proxy_headers = cast(
+            ASGIApp,
+            ProxyHeadersMiddleware(
+                cast(_UvicornASGIApp, app),
+                trusted_hosts="127.0.0.1" if trusted_hosts is None else trusted_hosts,
+            ),
+        )
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] in {"http", "websocket"}:
+            _capture_raw_socket_peer(scope)
+        await self._proxy_headers(scope, receive, send)
+
+
+def add_trusted_proxy_headers_middleware(app: FastAPI) -> None:
+    app.add_middleware(TrustedProxyHeadersMiddleware)

app/core/multipart.py

@@ -0,0 +1,237 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, AsyncIterator, Collection
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass
+
+from python_multipart.exceptions import MultipartParseError
+from starlette.datastructures import FormData, Headers, UploadFile
+from starlette.exceptions import HTTPException
+from starlette.formparsers import MultiPartException, MultiPartParser
+from starlette.requests import Request
+
+from app.core.multipart_fields import raise_upload_validation
+
+__all__ = [
+    "ACCOUNT_IMPORT_MULTIPART_POLICY",
+    "IMAGE_EDITS_MULTIPART_POLICY",
+    "TRANSCRIPTION_MULTIPART_POLICY",
+    "MultipartPayloadTooLarge",
+    "MultipartPolicy",
+    "bounded_multipart_form",
+    "read_bounded_upload",
+]
+
+_MIB = 1024 * 1024
+_PAYLOAD_TOO_LARGE_MESSAGE = "Multipart payload exceeds the allowed size."
+_INVALID_MULTIPART_MESSAGE = "Invalid multipart request."
+
+
+@dataclass(frozen=True, slots=True)
+class MultipartPolicy:
+    max_body_bytes: int
+    max_file_bytes: int
+    max_aggregate_file_bytes: int
+    max_files: int
+    max_fields: int
+    max_text_part_bytes: int
+
+    def __post_init__(self) -> None:
+        limits = (
+            self.max_body_bytes,
+            self.max_file_bytes,
+            self.max_aggregate_file_bytes,
+            self.max_files,
+            self.max_fields,
+            self.max_text_part_bytes,
+        )
+        if any(limit < 0 for limit in limits):
+            raise ValueError("Multipart limits must be non-negative")
+
+
+ACCOUNT_IMPORT_MULTIPART_POLICY = MultipartPolicy(
+    max_body_bytes=2 * _MIB,
+    max_file_bytes=_MIB,
+    max_aggregate_file_bytes=_MIB,
+    max_files=1,
+    max_fields=0,
+    max_text_part_bytes=0,
+)
+
+TRANSCRIPTION_MULTIPART_POLICY = MultipartPolicy(
+    max_body_bytes=32 * _MIB,
+    max_file_bytes=25_000_000,
+    max_aggregate_file_bytes=25_000_000,
+    max_files=1,
+    max_fields=32,
+    max_text_part_bytes=256 * 1024,
+)
+
+IMAGE_EDITS_MULTIPART_POLICY = MultipartPolicy(
+    max_body_bytes=64 * _MIB,
+    max_file_bytes=49_999_999,
+    max_aggregate_file_bytes=49_999_999,
+    max_files=17,
+    max_fields=32,
+    max_text_part_bytes=256 * 1024,
+)
+
+
+def _normalize_param(param: str | None) -> str | None:
+    if param == "image[]":
+        return "image"
+    return param
+
+
+class MultipartPayloadTooLarge(MultiPartException):
+    def __init__(self, *, param: str | None = None) -> None:
+        super().__init__(_PAYLOAD_TOO_LARGE_MESSAGE)
+        self.param = _normalize_param(param)
+
+
+async def _counted_request_stream(request: Request, max_body_bytes: int) -> AsyncGenerator[bytes, None]:
+    streamed_bytes = 0
+    async for chunk in request.stream():
+        streamed_bytes += len(chunk)
+        if streamed_bytes > max_body_bytes:
+            raise MultipartPayloadTooLarge()
+        yield chunk
+
+
+def _usable_content_length(headers: Headers) -> int | None:
+    raw_length = headers.get("content-length")
+    if raw_length is None:
+        return None
+    try:
+        declared_length = int(raw_length)
+    except ValueError:
+        return None
+    return declared_length if declared_length >= 0 else None
+
+
+def _is_multipart_form_data(headers: Headers) -> bool:
+    content_type = headers.get("content-type", "")
+    media_type = content_type.split(";", 1)[0].strip().lower()
+    return media_type == "multipart/form-data"
+
+
+class _BoundedMultiPartParser(MultiPartParser):
+    def __init__(
+        self,
+        headers: Headers,
+        stream: AsyncGenerator[bytes, None],
+        policy: MultipartPolicy,
+        typed_upload_fields: Collection[str],
+    ) -> None:
+        super().__init__(
+            headers,
+            stream,
+            max_files=policy.max_files,
+            max_fields=policy.max_fields,
+            max_part_size=policy.max_text_part_bytes,
+        )
+        self._policy = policy
+        self._typed_upload_fields = frozenset(typed_upload_fields)
+        self._current_file_bytes = 0
+        self._aggregate_file_bytes = 0
+
+    def on_part_begin(self) -> None:
+        super().on_part_begin()
+        self._current_file_bytes = 0
+
+    def on_headers_finished(self) -> None:
+        try:
+            super().on_headers_finished()
+        except MultiPartException:
+            if self._current_fields > self.max_fields and self._current_part.field_name in self._typed_upload_fields:
+                raise_upload_validation(self._current_part.field_name)
+            raise
+
+    def on_part_data(self, data: bytes, start: int, end: int) -> None:
+        part_bytes = end - start
+        current_part = self._current_part
+        if current_part.file is None:
+            if len(current_part.data) + part_bytes > self._policy.max_text_part_bytes:
+                raise MultipartPayloadTooLarge()
+        else:
+            param = _normalize_param(current_part.field_name)
+            if self._current_file_bytes + part_bytes > self._policy.max_file_bytes:
+                raise MultipartPayloadTooLarge(param=param)
+            if self._aggregate_file_bytes + part_bytes > self._policy.max_aggregate_file_bytes:
+                raise MultipartPayloadTooLarge(param=param)
+            self._current_file_bytes += part_bytes
+            self._aggregate_file_bytes += part_bytes
+        super().on_part_data(data, start, end)
+
+    def close_created_files(self) -> None:
+        files = self._files_to_close_on_error
+        self._files_to_close_on_error = []
+        for file in files:
+            with suppress(OSError):
+                file.close()
+
+
+def _multipart_bad_request(exc: MultiPartException | MultipartParseError | OSError) -> HTTPException:
+    if isinstance(exc, MultiPartException):
+        detail = exc.message
+    else:
+        detail = _INVALID_MULTIPART_MESSAGE
+    return HTTPException(status_code=400, detail=detail)
+
+
+async def _parse_with_cleanup(parser: _BoundedMultiPartParser) -> FormData:
+    try:
+        return await parser.parse()
+    except BaseException:
+        parser.close_created_files()
+        raise
+
+
+@asynccontextmanager
+async def bounded_multipart_form(
+    request: Request,
+    policy: MultipartPolicy,
+    *,
+    typed_upload_fields: Collection[str] = (),
+) -> AsyncIterator[FormData]:
+    if not _is_multipart_form_data(request.headers):
+        empty_form = FormData()
+        try:
+            yield empty_form
+        finally:
+            await empty_form.close()
+        return
+
+    declared_length = _usable_content_length(request.headers)
+    if declared_length is not None and declared_length > policy.max_body_bytes:
+        raise MultipartPayloadTooLarge()
+
+    parser = _BoundedMultiPartParser(
+        request.headers,
+        _counted_request_stream(request, policy.max_body_bytes),
+        policy,
+        typed_upload_fields,
+    )
+    try:
+        try:
+            form = await _parse_with_cleanup(parser)
+        except MultipartPayloadTooLarge:
+            raise
+        except (MultiPartException, MultipartParseError, OSError) as exc:
+            raise _multipart_bad_request(exc) from exc
+
+        try:
+            yield form
+        finally:
+            await form.close()
+    finally:
+        parser.close_created_files()
+
+
+async def read_bounded_upload(upload: UploadFile, max_bytes: int, param: str | None) -> bytes:
+    if max_bytes < 0:
+        raise ValueError("Upload limit must be non-negative")
+    data = await upload.read(max_bytes + 1)
+    if len(data) > max_bytes:
+        raise MultipartPayloadTooLarge(param=param)
+    return data