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",
}
METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"]


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)
    url = f"http://{config.CONTAINER_PROXY_HOST}:{port}/{path}"
    headers = {k: v for k, v in request.headers.items() if k.lower() not in HOP_HEADERS}
    body = await request.body()
    try:
        async with httpx.AsyncClient(timeout=60.0) 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}
    return Response(content=upstream.content, status_code=upstream.status_code,
                    headers=out_headers, media_type=upstream.headers.get("content-type"))


@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())