RAG/backend/app/gateway.py

"""Small fixed-origin ingress gateway with explicit proxy boundaries."""

import re
from collections.abc import AsyncIterator, Mapping
from contextlib import asynccontextmanager

import httpx
from fastapi import FastAPI, Request, Response, status
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.types import Receive, Scope, Send

UPSTREAM_ORIGIN = httpx.URL("http://api:8000")
MAX_REQUEST_BODY_BYTES = 1024 * 1024
MAX_UPLOAD_BODY_BYTES = 100 * 1024 * 1024
SUPPORTED_METHODS = ["GET", "POST", "PUT", "HEAD", "OPTIONS"]
UPLOAD_CONTENT_PATH = re.compile(r"^/api/v1/document-uploads/[0-9a-fA-F-]{36}/content$")

REQUEST_HEADER_ALLOWLIST = frozenset(
    {
        "accept",
        "accept-language",
        "access-control-request-headers",
        "access-control-request-method",
        "authorization",
        "content-type",
        "if-match",
        "if-none-match",
        "idempotency-key",
        "origin",
        "range",
        "traceparent",
        "tracestate",
        "x-request-id",
    }
)


class RequestBodyTooLarge(Exception):
    """Internal control flow for a streamed body that exceeded its hard cap."""


RESPONSE_HEADER_ALLOWLIST = frozenset(
    {
        "accept-ranges",
        "access-control-allow-credentials",
        "access-control-allow-headers",
        "access-control-allow-methods",
        "access-control-allow-origin",
        "access-control-max-age",
        "allow",
        "cache-control",
        "content-disposition",
        "content-language",
        "content-range",
        "content-type",
        "etag",
        "last-modified",
        "retry-after",
        "vary",
        "www-authenticate",
        "x-request-id",
    }
)


class _UpstreamStreamingResponse(StreamingResponse):
    """Close the upstream response even if downstream headers cannot be sent."""

    def __init__(
        self,
        *,
        upstream_response: httpx.Response,
        content: AsyncIterator[bytes],
        status_code: int,
        headers: Mapping[str, str],
    ) -> None:
        self._upstream_response = upstream_response
        super().__init__(content=content, status_code=status_code, headers=headers)

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        try:
            await super().__call__(scope, receive, send)
        finally:
            await self._upstream_response.aclose()


def _allowlisted_headers(headers: Mapping[str, str], allowlist: frozenset[str]) -> dict[str, str]:
    return {name: value for name, value in headers.items() if name.lower() in allowlist}


def _upstream_url(request: Request) -> httpx.URL:
    raw_path_value: object = request.scope.get("raw_path")
    if isinstance(raw_path_value, bytes):
        raw_path = raw_path_value
    else:
        raw_path = request.url.path.encode("utf-8")
    if not raw_path.startswith(b"/"):
        raw_path = b"/" + raw_path

    query_value: object = request.scope.get("query_string")
    query = query_value if isinstance(query_value, bytes) else b""
    raw_target = raw_path + (b"?" + query if query else b"")
    # copy_with preserves the fixed scheme, host, and port even for //host-like paths.
    return UPSTREAM_ORIGIN.copy_with(raw_path=raw_target)


async def _bounded_body(request: Request) -> bytes | None:
    declared_length = request.headers.get("content-length")
    if declared_length is not None:
        try:
            parsed_length = int(declared_length)
        except ValueError:
            return None
        if parsed_length < 0 or parsed_length > MAX_REQUEST_BODY_BYTES:
            return None

    body = bytearray()
    async for chunk in request.stream():
        if len(body) + len(chunk) > MAX_REQUEST_BODY_BYTES:
            return None
        body.extend(chunk)
    return bytes(body)


def _declared_body_is_too_large(request: Request, maximum: int) -> bool:
    declared_length = request.headers.get("content-length")
    if declared_length is None:
        return False
    try:
        parsed_length = int(declared_length)
    except ValueError:
        return True
    return parsed_length < 0 or parsed_length > maximum


async def _bounded_upload_stream(request: Request) -> AsyncIterator[bytes]:
    total = 0
    async for chunk in request.stream():
        total += len(chunk)
        if total > MAX_UPLOAD_BODY_BYTES:
            raise RequestBodyTooLarge
        if chunk:
            yield chunk


def create_gateway_app(transport: httpx.AsyncBaseTransport | None = None) -> FastAPI:
    """Create a no-secret gateway; an injected transport enables hermetic tests."""

    upstream_client = httpx.AsyncClient(
        timeout=httpx.Timeout(10.0, connect=2.0),
        limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
        follow_redirects=False,
        transport=transport,
        trust_env=False,
    )
    # Drop httpx convenience defaults; Host and body framing are added by the protocol layer.
    upstream_client.headers.clear()

    @asynccontextmanager
    async def lifespan(_: FastAPI) -> AsyncIterator[None]:
        try:
            yield
        finally:
            await upstream_client.aclose()

    gateway = FastAPI(
        title="Geological RAG Gateway",
        docs_url=None,
        redoc_url=None,
        openapi_url=None,
        lifespan=lifespan,
    )

    @gateway.get("/gateway/live", include_in_schema=False)
    async def gateway_live() -> dict[str, str]:
        return {"status": "ok"}

    @gateway.api_route("/{path:path}", methods=SUPPORTED_METHODS, include_in_schema=False)
    async def proxy(request: Request, path: str) -> Response:  # noqa: ARG001
        is_upload = request.method == "PUT" and UPLOAD_CONTENT_PATH.fullmatch(request.url.path)
        if is_upload and _declared_body_is_too_large(request, MAX_UPLOAD_BODY_BYTES):
            return JSONResponse(
                status_code=status.HTTP_413_CONTENT_TOO_LARGE,
                content={"detail": "request body too large"},
            )
        if is_upload:
            body: bytes | AsyncIterator[bytes] = _bounded_upload_stream(request)
        else:
            bounded = await _bounded_body(request)
            if bounded is None:
                return JSONResponse(
                    status_code=status.HTTP_413_CONTENT_TOO_LARGE,
                    content={"detail": "request body too large"},
                )
            body = bounded

        upstream_request = upstream_client.build_request(
            request.method,
            _upstream_url(request),
            headers=_allowlisted_headers(request.headers, REQUEST_HEADER_ALLOWLIST),
            content=body,
        )
        try:
            upstream_response = await upstream_client.send(upstream_request, stream=True)
        except RequestBodyTooLarge:
            return JSONResponse(
                status_code=status.HTTP_413_CONTENT_TOO_LARGE,
                content={"detail": "request body too large"},
            )
        except httpx.RequestError:
            return JSONResponse(
                status_code=status.HTTP_502_BAD_GATEWAY,
                content={"detail": "upstream unavailable"},
            )

        async def response_body() -> AsyncIterator[bytes]:
            if upstream_response.is_stream_consumed:
                yield upstream_response.content
            else:
                async for chunk in upstream_response.aiter_raw():
                    yield chunk

        return _UpstreamStreamingResponse(
            upstream_response=upstream_response,
            content=response_body(),
            status_code=upstream_response.status_code,
            headers=_allowlisted_headers(
                upstream_response.headers,
                RESPONSE_HEADER_ALLOWLIST,
            ),
        )

    return gateway


app = create_gateway_app()