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, StreamingResponse
from starlette.types import Receive, Scope, Send

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


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 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, stream=True)
        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()