AgentEvalTool/backend/agenteval/web/routers/proxy.py

"""Reverse proxy for OpenClaw, rewriting headers for iframe embedding.

OpenClaw is configured with trusted-proxy auth mode, so the proxy sends
an x-forwarded-user header instead of a Bearer token. All connection
parameters come from :mod:`agenteval.config.settings` — no hardcoded hosts.
"""

import asyncio
import json

import httpx
import websockets
from fastapi import APIRouter, Request, WebSocket
from fastapi.responses import StreamingResponse

from agenteval.config import get_settings


def get_http_upstream() -> str:
    return get_settings().openclaw_upstream


def get_ws_upstream() -> str:
    return get_settings().openclaw_ws_upstream


def get_proxy_user() -> str:
    return get_settings().openclaw_proxy_user


def get_auth_token() -> str:
    return get_settings().openclaw_auth_token


_client = httpx.AsyncClient(follow_redirects=True, timeout=60.0)

router = APIRouter()


def _rewrite_headers(headers: httpx.Headers) -> dict[str, str]:
    """Strip iframe-blocking headers and rewrite CSP."""
    result: dict[str, str] = {}
    for key, value in headers.items():
        lower = key.lower()
        if lower in ("x-frame-options", "transfer-encoding", "content-encoding", "content-length"):
            continue
        if lower == "content-security-policy":
            value = value.replace("frame-ancestors 'none'", "frame-ancestors 'self'")
            value = value.replace("script-src 'self'", "script-src 'self' 'unsafe-inline'")
        result[key] = value
    return result


@router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
async def http_proxy(request: Request, path: str):
    url = f"{get_http_upstream()}/{path}"
    if request.query_params:
        url += f"?{request.query_params}"

    body = await request.body()
    headers = dict(request.headers)
    headers.pop("host", None)
    headers.pop("referer", None)
    headers.pop("origin", None)
    headers["x-forwarded-user"] = get_proxy_user()
    headers["authorization"] = f"Bearer {get_auth_token()}"

    resp = await _client.request(
        method=request.method,
        url=url,
        content=body,
        headers=headers,
    )

    return StreamingResponse(
        content=iter([resp.content]),
        status_code=resp.status_code,
        headers=_rewrite_headers(resp.headers),
    )


async def ws_bridge(client_ws: WebSocket, upstream_url: str):
    """Bidirectional WebSocket bridge between client and upstream."""
    await client_ws.accept()

    settings = get_settings()
    auth_token = settings.openclaw_auth_token
    connect_kwargs: dict = {
        "additional_headers": {
            "x-forwarded-user": settings.openclaw_proxy_user,
            "authorization": f"Bearer {auth_token}",
        },
        "max_size": 10 * 1024 * 1024,
    }
    if settings.openclaw_ws_origin:
        connect_kwargs["origin"] = settings.openclaw_ws_origin

    try:
        async with websockets.connect(upstream_url, **connect_kwargs) as upstream_ws:

            async def client_to_upstream():
                try:
                    while True:
                        msg = await client_ws.receive_text()
                        try:
                            payload = json.loads(msg)
                            if (
                                isinstance(payload, dict)
                                and payload.get("method") == "connect.authenticate"
                                and isinstance(payload.get("payload"), dict)
                            ):
                                params = payload["payload"].setdefault("params", {})
                                params.setdefault("auth", {})
                                params["auth"].setdefault("token", auth_token)
                                msg = json.dumps(payload)
                        except (ValueError, TypeError, AttributeError):
                            pass
                        await upstream_ws.send(msg)
                except Exception:
                    pass

            async def upstream_to_client():
                try:
                    async for msg in upstream_ws:
                        if isinstance(msg, str):
                            await client_ws.send_text(msg)
                        else:
                            await client_ws.send_bytes(msg)
                except Exception:
                    pass

            done, pending = await asyncio.wait(
                [asyncio.create_task(client_to_upstream()), asyncio.create_task(upstream_to_client())],
                return_when=asyncio.FIRST_COMPLETED,
            )
            for task in pending:
                task.cancel()
    except Exception:
        pass
    finally:
        try:
            await client_ws.close()
        except Exception:
            pass