proxy.py

"""
Reverse proxy for virtual ports.

Each zone can run web servers on internal ports (e.g. 3000, 8080).
Since HF Spaces only exposes port 7860, this module proxies
/port/{zone}/{port}/... → localhost:{port} with proper path rewriting.

Port mappings are stored in zones_meta.json under each zone's "ports" key.
"""

import json
import re
from pathlib import Path

import httpx
from fastapi import Request, WebSocket, WebSocketDisconnect
from fastapi.responses import Response

# Port range allowed for proxying (unprivileged ports only)
MIN_PORT = 1024
MAX_PORT = 65535

# Reuse a single async HTTP client
_client: httpx.AsyncClient | None = None


def _get_client() -> httpx.AsyncClient:
    global _client
    if _client is None:
        _client = httpx.AsyncClient(
            timeout=httpx.Timeout(30.0, connect=5.0),
            follow_redirects=False,
            limits=httpx.Limits(max_connections=50),
        )
    return _client


def _meta_path() -> Path:
    from zones import ZONES_META
    return ZONES_META


def _load_meta() -> dict:
    p = _meta_path()
    if p.exists():
        return json.loads(p.read_text(encoding="utf-8"))
    return {}


def _save_meta(meta: dict):
    _meta_path().write_text(json.dumps(meta, indent=2, default=str), encoding="utf-8")


def _validate_port(port: int):
    if not (MIN_PORT <= port <= MAX_PORT):
        raise ValueError(f"Port must be between {MIN_PORT} and {MAX_PORT}")


def _validate_zone(meta: dict, zone_name: str):
    if zone_name not in meta:
        raise ValueError(f"Zone '{zone_name}' does not exist")


# ── Port CRUD ─────────────────────────────────

def list_ports(zone_name: str) -> list[dict]:
    """List all port mappings for a zone."""
    meta = _load_meta()
    _validate_zone(meta, zone_name)
    ports = meta[zone_name].get("ports", [])
    return ports


def add_port(zone_name: str, port: int, label: str = "") -> dict:
    """Add a port mapping to a zone."""
    _validate_port(port)
    meta = _load_meta()
    _validate_zone(meta, zone_name)

    ports = meta[zone_name].setdefault("ports", [])

    # Check duplicate
    for p in ports:
        if p["port"] == port:
            raise ValueError(f"Port {port} already mapped in zone '{zone_name}'")

    entry = {"port": port, "label": label or f"Port {port}"}
    ports.append(entry)
    _save_meta(meta)
    return entry


def remove_port(zone_name: str, port: int):
    """Remove a port mapping from a zone."""
    meta = _load_meta()
    _validate_zone(meta, zone_name)

    ports = meta[zone_name].get("ports", [])
    before = len(ports)
    meta[zone_name]["ports"] = [p for p in ports if p["port"] != port]
    if len(meta[zone_name]["ports"]) == before:
        raise ValueError(f"Port {port} not found in zone '{zone_name}'")
    _save_meta(meta)


# ── HTTP Reverse Proxy ────────────────────────

# Headers that should not be forwarded
_HOP_HEADERS = frozenset({
    "connection", "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "transfer-encoding", "upgrade",
})


async def proxy_http(request: Request, zone_name: str, port: int, subpath: str = "") -> Response:
    """Proxy an HTTP request to localhost:{port}."""
    _validate_port(port)
    meta = _load_meta()
    _validate_zone(meta, zone_name)

    # Verify port is registered for this zone
    ports = meta[zone_name].get("ports", [])
    if not any(p["port"] == port for p in ports):
        return Response(content="Port not mapped", status_code=404)

    target_url = f"http://127.0.0.1:{port}/{subpath}"
    if request.url.query:
        target_url += f"?{request.url.query}"

    # Build headers, filtering out hop-by-hop
    headers = {}
    for key, value in request.headers.items():
        if key.lower() not in _HOP_HEADERS and key.lower() != "host":
            headers[key] = value
    headers["host"] = f"127.0.0.1:{port}"
    headers["x-forwarded-for"] = request.client.host if request.client else "127.0.0.1"
    headers["x-forwarded-proto"] = request.url.scheme
    # Tell the target app its real base path
    headers["x-forwarded-prefix"] = f"/port/{zone_name}/{port}"

    body = await request.body()
    client = _get_client()

    try:
        resp = await client.request(
            method=request.method,
            url=target_url,
            headers=headers,
            content=body,
        )
    except httpx.ConnectError:
        return Response(
            content=f"Cannot connect to port {port}. Make sure your server is running.",
            status_code=502,
            media_type="text/plain",
        )
    except httpx.TimeoutException:
        return Response(content=f"Timeout connecting to port {port}", status_code=504, media_type="text/plain")

    # Build response headers
    resp_headers = {}
    for key, value in resp.headers.items():
        if key.lower() not in _HOP_HEADERS and key.lower() != "content-encoding":
            resp_headers[key] = value

    return Response(
        content=resp.content,
        status_code=resp.status_code,
        headers=resp_headers,
    )


# ── WebSocket Reverse Proxy ──────────────────

async def proxy_ws(websocket: WebSocket, zone_name: str, port: int, subpath: str = ""):
    """Proxy a WebSocket connection to localhost:{port}."""
    _validate_port(port)
    meta = _load_meta()
    _validate_zone(meta, zone_name)

    ports = meta[zone_name].get("ports", [])
    if not any(p["port"] == port for p in ports):
        await websocket.close(code=4004, reason="Port not mapped")
        return

    await websocket.accept()

    target_url = f"ws://127.0.0.1:{port}/{subpath}"

    try:
        async with httpx.AsyncClient() as client:
            async with client.stream("GET", target_url.replace("ws://", "http://")) as _:
                pass
    except Exception:
        pass

    # Use raw websockets for the backend connection
    import asyncio
    import websockets

    try:
        async with websockets.connect(target_url) as backend_ws:
            async def client_to_backend():
                try:
                    while True:
                        msg = await websocket.receive()
                        if msg.get("type") == "websocket.disconnect":
                            break
                        if "text" in msg:
                            await backend_ws.send(msg["text"])
                        elif "bytes" in msg:
                            await backend_ws.send(msg["bytes"])
                except (WebSocketDisconnect, Exception):
                    pass

            async def backend_to_client():
                try:
                    async for message in backend_ws:
                        if isinstance(message, str):
                            await websocket.send_text(message)
                        else:
                            await websocket.send_bytes(message)
                except (WebSocketDisconnect, Exception):
                    pass

            await asyncio.gather(client_to_backend(), backend_to_client())
    except Exception:
        try:
            await websocket.send_text(json.dumps({"error": f"Cannot connect WebSocket to port {port}"}))
        except Exception:
            pass
    finally:
        try:
            await websocket.close()
        except Exception:
            pass