devplacepy/devplacepy/routers/proxy.py

# retoor <retoor@molodetz.nl>

import asyncio
import json
import logging

import httpx
import websockets
from fastapi import APIRouter, Request, WebSocket
from starlette.responses import Response

from devplacepy import config
from devplacepy.services.containers import store
from devplacepy.utils import not_found

logger = logging.getLogger(__name__)
router = APIRouter()

HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "host",
    "content-length",
    "content-encoding",
}
METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"]


def _forward_headers(request: Request, prefix: str) -> dict:
    headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_HEADERS}
    headers["X-Forwarded-Prefix"] = prefix
    headers["X-Script-Name"] = prefix
    headers["X-Forwarded-Host"] = request.headers.get(
        "host", request.url.hostname or ""
    )
    headers["X-Forwarded-Proto"] = request.headers.get(
        "x-forwarded-proto", request.url.scheme
    )
    headers["Accept-Encoding"] = "identity"
    return headers


def _inject_base(body: bytes, prefix: str) -> bytes:
    lowered = body.lower()
    if b"<base" in lowered:
        return body
    tag = f'<base href="{prefix}/">'.encode()
    head = lowered.find(b"<head")
    anchor = (
        lowered.find(b">", head)
        if head != -1
        else lowered.find(b">", lowered.find(b"<html"))
    )
    if anchor == -1:
        return tag + body
    return body[: anchor + 1] + tag + body[anchor + 1 :]


def _rewrite_location(value: str, prefix: str) -> str:
    if value.startswith("/") and not value.startswith("//"):
        return prefix + value
    return value


def _resolve(slug: str):
    instance = store.find_instance_by_ingress(slug)
    if not instance or instance.get("status") != store.ST_RUNNING:
        return None, None
    ports = json.loads(instance.get("ports_json") or "[]")
    if not ports:
        return instance, None
    ingress_port = int(instance.get("ingress_port") or 0)
    if ingress_port:
        for port in ports:
            if int(port["container"]) == ingress_port:
                return instance, int(port["host"])
        return instance, None
    return instance, int(ports[0]["host"])


@router.api_route("/{slug}", methods=METHODS)
@router.api_route("/{slug}/{path:path}", methods=METHODS)
async def proxy_http(request: Request, slug: str, path: str = ""):
    instance, port = _resolve(slug)
    if instance is None:
        raise not_found("No running container is published at this address")
    if port is None:
        return Response(
            "the published container has no reachable port", status_code=502
        )
    prefix = f"/p/{slug}"
    url = f"http://{config.CONTAINER_PROXY_HOST}:{port}/{path}"
    headers = _forward_headers(request, prefix)
    body = await request.body()
    try:
        async with httpx.AsyncClient(timeout=60.0, follow_redirects=False) as client:
            upstream = await client.request(
                request.method,
                url,
                params=request.query_params,
                headers=headers,
                content=body,
            )
    except httpx.HTTPError as exc:
        return Response(f"upstream error: {exc}", status_code=502)
    out_headers = {
        k: v
        for k, v in upstream.headers.items()
        if k.lower() not in HOP_HEADERS and k.lower() != "set-cookie"
    }
    if "location" in out_headers:
        out_headers["location"] = _rewrite_location(out_headers["location"], prefix)
    content_type = upstream.headers.get("content-type", "")
    content = upstream.content
    if "text/html" in content_type.lower():
        content = _inject_base(content, prefix)
    response = Response(
        content=content,
        status_code=upstream.status_code,
        headers=out_headers,
        media_type=content_type or None,
    )
    for cookie in upstream.headers.get_list("set-cookie"):
        response.headers.append("set-cookie", cookie)
    return response


@router.websocket("/{slug}")
@router.websocket("/{slug}/{path:path}")
async def proxy_ws(websocket: WebSocket, slug: str, path: str = ""):
    instance, port = _resolve(slug)
    if instance is None or port is None:
        await websocket.close(code=1011)
        return
    upstream_url = f"ws://{config.CONTAINER_PROXY_HOST}:{port}/{path}"
    if websocket.url.query:
        upstream_url += f"?{websocket.url.query}"
    await websocket.accept()
    try:
        async with websockets.connect(
            upstream_url, open_timeout=10, max_size=None
        ) as upstream:
            await _pump(websocket, upstream)
    except Exception as exc:
        logger.debug("ws proxy %s failed: %s", slug, exc)
        try:
            await websocket.close(code=1011)
        except Exception:
            pass


async def _pump(client_ws: WebSocket, upstream) -> None:
    async def client_to_upstream():
        try:
            while True:
                message = await client_ws.receive()
                if message["type"] == "websocket.disconnect":
                    break
                if message.get("text") is not None:
                    await upstream.send(message["text"])
                elif message.get("bytes") is not None:
                    await upstream.send(message["bytes"])
        except Exception:
            pass
        finally:
            await upstream.close()

    async def upstream_to_client():
        try:
            async for message in upstream:
                if isinstance(message, (bytes, bytearray)):
                    await client_ws.send_bytes(bytes(message))
                else:
                    await client_ws.send_text(message)
        except Exception:
            pass
        finally:
            try:
                await client_ws.close()
            except Exception:
                pass

    await asyncio.gather(client_to_upstream(), upstream_to_client())