Api-Gateway/app/proxy.py

import re
import time
from functools import lru_cache

import httpx
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session

from app import config, ratelimit, security
from app.database import get_db
from app.models import ApiKey, Endpoint, RequestLog, Service, utcnow

router = APIRouter()

# Hop-by-hop headers must not be forwarded either direction
_HOP_BY_HOP = {
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade", "host", "content-length",
}

# One client per TLS-verification mode (verify is a client-level setting in httpx).
_clients: dict[bool, httpx.AsyncClient] = {}


async def get_client(verify_tls: bool = True) -> httpx.AsyncClient:
    client = _clients.get(verify_tls)
    if client is None:
        client = _clients[verify_tls] = httpx.AsyncClient(
            follow_redirects=False, verify=verify_tls)
    return client


async def close_client() -> None:
    for client in _clients.values():
        await client.aclose()
    _clients.clear()


@lru_cache(maxsize=1024)
def _pattern_to_regex(pattern: str) -> re.Pattern:
    """Endpoint path pattern -> regex. '{param}' = one segment, '*' = any
    within a segment, '**' = any depth."""
    regex = ""
    for token in re.split(r"(\{[^}]*\}|\*\*|\*)", pattern):
        if not token:
            continue
        if token == "**":
            regex += ".*"
        elif token == "*":
            regex += "[^/]*"
        elif token.startswith("{") and token.endswith("}"):
            regex += "[^/]+"
        else:
            regex += re.escape(token)
    return re.compile("^" + regex + "$")


def match_endpoints(endpoints: list[Endpoint], method: str, path: str) -> list[Endpoint]:
    """All endpoints of a service that match this request, most specific
    (exact method) first."""
    matches = [
        e for e in endpoints
        if e.method in ("*", method) and _pattern_to_regex(e.path).match(path)
    ]
    return sorted(matches, key=lambda e: e.method == "*")


def _error(status: int, code: str, message: str) -> JSONResponse:
    return JSONResponse(status_code=status, content={"error": code, "message": message})


# Payload capture for transaction inspection
MAX_STORED_BODY = 64 * 1024
_TEXTUAL_TYPES = ("application/json", "application/xml",
                  "application/x-www-form-urlencoded", "text/", "+json", "+xml")


def _readable(body: bytes, content_type: str) -> str:
    """Body as storable text: textual payloads are kept (truncated at 64 KB),
    binary ones are summarized."""
    if not body:
        return ""
    ct = (content_type or "").lower()
    if not any(t in ct for t in _TEXTUAL_TYPES):
        return f"<{len(body)} bytes of {ct or 'unknown content type'}>"
    text = body[:MAX_STORED_BODY].decode("utf-8", "replace")
    if len(body) > MAX_STORED_BODY:
        text += "\n… (truncated)"
    return text


@router.api_route(
    "/{slug}/{path:path}",
    methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"],
)
@router.api_route(
    "/{slug}",
    methods=["GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"],
)
async def gateway(slug: str, request: Request, path: str = "",
                  db: Session = Depends(get_db)):
    start = time.perf_counter()

    # --- authenticate ---
    plain_key = request.headers.get(config.API_KEY_HEADER)
    if not plain_key:
        return _error(401, "missing_api_key", f"Provide your API key in the {config.API_KEY_HEADER} header.")

    api_key = (
        db.query(ApiKey)
        .filter(ApiKey.key_hash == security.hash_api_key(plain_key))
        .one_or_none()
    )
    if api_key is None or not api_key.is_active or not api_key.user.is_active:
        return _error(403, "invalid_api_key", "API key is unknown or has been revoked.")

    # --- resolve service ---
    service = db.query(Service).filter(Service.slug == slug).one_or_none()
    if service is None or not service.is_active:
        return _error(404, "unknown_service", f"No active service registered under '{slug}'.")

    # --- read the request payload once (forwarded and logged) ---
    body = await request.body()
    stored_request = _readable(body, request.headers.get("content-type", ""))

    def deny(status: int, code: str, message: str, endpoint: Endpoint | None) -> JSONResponse:
        response = _error(status, code, message)
        _log(db, api_key, service, endpoint, request, request_path, status, start,
             stored_request, response.body.decode())
        return response

    # --- resolve endpoint & authorize ---
    request_path = "/" + path
    matched = match_endpoints(service.endpoints, request.method, request_path)
    if not matched:
        return deny(404, "unknown_endpoint",
                    f"No endpoint of '{slug}' matches {request.method} {request_path}.", None)

    granted_ids = {e.id for e in api_key.endpoints}
    endpoint = next((e for e in matched if e.id in granted_ids), None)
    if endpoint is None:
        return deny(403, "access_denied",
                    f"This API key has no access to {request.method} {request_path} on '{slug}'.",
                    matched[0])

    # --- rate limit ---
    if not ratelimit.check(api_key.id, api_key.rate_limit_per_minute):
        return deny(429, "rate_limited",
                    f"Rate limit of {api_key.rate_limit_per_minute} requests/minute exceeded.",
                    endpoint)

    # --- proxy upstream ---
    upstream_url = service.base_url.rstrip("/") + request_path
    headers = {
        k: v for k, v in request.headers.items()
        if k.lower() not in _HOP_BY_HOP and k.lower() != config.API_KEY_HEADER.lower()
    }
    headers["x-forwarded-for"] = request.client.host if request.client else ""
    headers["x-forwarded-proto"] = request.url.scheme

    client = await get_client(service.verify_tls)
    try:
        upstream = await client.request(
            request.method,
            upstream_url,
            params=request.query_params,
            content=body,
            headers=headers,
            timeout=service.timeout_seconds or config.PROXY_DEFAULT_TIMEOUT,
        )
    except httpx.TimeoutException:
        return deny(504, "upstream_timeout", "The upstream API did not respond in time.", endpoint)
    except httpx.HTTPError:
        return deny(502, "upstream_unreachable", "Could not reach the upstream API.", endpoint)

    _log(db, api_key, service, endpoint, request, request_path, upstream.status_code, start,
         stored_request, _readable(upstream.content, upstream.headers.get("content-type", "")))

    response_headers = {
        k: v for k, v in upstream.headers.items() if k.lower() not in _HOP_BY_HOP
    }
    return Response(
        content=upstream.content,
        status_code=upstream.status_code,
        headers=response_headers,
        media_type=upstream.headers.get("content-type"),
    )


def _log(db: Session, api_key: ApiKey, service: Service, endpoint: Endpoint | None,
         request: Request, request_path: str, status: int, start: float,
         request_body: str = "", response_body: str = "") -> None:
    api_key.last_used_at = utcnow()
    db.add(RequestLog(
        api_key_id=api_key.id,
        service_id=service.id,
        endpoint_id=endpoint.id if endpoint else None,
        method=request.method,
        path=request_path,
        query_string=request.url.query,
        status_code=status,
        latency_ms=round((time.perf_counter() - start) * 1000, 2),
        client_ip=request.client.host if request.client else "",
        request_body=request_body,
        response_body=response_body,
    ))
    db.commit()