RAG/backend/app/gateway.py

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

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

import httpx
from fastapi import FastAPI, Request, Response, status
from fastapi.responses import JSONResponse

UPSTREAM_ORIGIN = httpx.URL("http://api:8000")
MAX_REQUEST_BODY_BYTES = 1024 * 1024
SUPPORTED_METHODS = ["GET", "POST", "HEAD", "OPTIONS"]

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

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",
    }
)


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 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
        body = await _bounded_body(request)
        if body is None:
            return JSONResponse(
                status_code=status.HTTP_413_CONTENT_TOO_LARGE,
                content={"detail": "request body too large"},
            )

        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)
        except httpx.RequestError:
            return JSONResponse(
                status_code=status.HTTP_502_BAD_GATEWAY,
                content={"detail": "upstream unavailable"},
            )

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

    return gateway


app = create_gateway_app()