hermes-hub/app/proxy.py

import httpx
import re
from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
import logging

logger = logging.getLogger("hermes_hub.proxy")

HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
    "content-length"
}

_http_transport: Optional[httpx.AsyncHTTPTransport] = None

def get_http_transport() -> httpx.AsyncHTTPTransport:
    global _http_transport
    if _http_transport is None:
        _http_transport = httpx.AsyncHTTPTransport(
            retries=1,
            limits=httpx.Limits(max_keepalive_connections=50, max_connections=100, keepalive_expiry=30.0)
        )
    return _http_transport

async def close_http_client():
    global _http_transport
    if _http_transport:
        await _http_transport.aclose()
        _http_transport = None

def rewrite_cookie_path(cookie_header: str, universe_id: str) -> str:
    """
    Rewrites the Path attribute of a Set-Cookie header to /u/{universe_id}/
    when using path-based proxying.
    """
    target_path = f"/u/{universe_id}/"
    if re.search(r'(?i)\bpath=[^;]*', cookie_header):
        return re.sub(r'(?i)\bpath=[^;]*', f'Path={target_path}', cookie_header)
    else:
        return f"{cookie_header}; Path={target_path}"

async def check_backend_health(backend_url: str) -> Dict[str, Any]:
    transport = get_http_transport()
    try:
        req = httpx.Request("GET", backend_url)
        res = await transport.handle_async_request(req)
        status_code = res.status_code
        await res.aclose()
        return {
            "status": "online" if status_code < 500 else "degraded",
            "status_code": status_code,
            "latency_ms": 100
        }
    except Exception as e:
        return {
            "status": "offline",
            "error": str(e),
            "latency_ms": None
        }

async def proxy_request(
    request: Request,
    backend_url: str,
    path: str,
    universe_id: Optional[str] = None
) -> Response:
    transport = get_http_transport()
    
    base_url = backend_url.rstrip("/")
    sub_path = path.lstrip("/")
    target_url = f"{base_url}/{sub_path}" if sub_path else base_url
    
    if request.url.query:
        target_url = f"{target_url}?{request.url.query}"
    
    req_headers: List[Tuple[bytes, bytes]] = []
    for raw_k, raw_v in request.headers.raw:
        k_str = raw_k.decode("latin-1").lower()
        if k_str not in HOP_BY_HOP_HEADERS and k_str != "host":
            req_headers.append((raw_k, raw_v))
            
    # Forward original host header info
    req_headers.append((b"x-forwarded-host", request.headers.get("host", "").encode("latin-1")))
    req_headers.append((b"x-forwarded-proto", (request.url.scheme or "http").encode("latin-1")))
    
    body = await request.body()
    
    is_path_proxied = universe_id and request.url.path.startswith(f"/u/{universe_id}")
    
    try:
        upstream_req = httpx.Request(
            method=request.method,
            url=target_url,
            headers=req_headers,
            content=body
        )
        
        # Pure stateless async request execution — NO cookie jar retention
        upstream_res = await transport.handle_async_request(upstream_req)
        
        # Build raw headers list for precise multi-header control
        raw_headers: List[Tuple[bytes, bytes]] = []
        media_type = upstream_res.headers.get("content-type")
        
        for raw_k, raw_v in upstream_res.headers.raw:
            k_str = raw_k.decode("latin-1").lower()
            v_str = raw_v.decode("latin-1")
            
            if k_str in HOP_BY_HOP_HEADERS:
                continue
                
            if k_str == "set-cookie" and is_path_proxied:
                # Path-based proxying: rewrite cookie path
                v_str = rewrite_cookie_path(v_str, universe_id)
                raw_headers.append((b"set-cookie", v_str.encode("latin-1")))
            elif k_str == "location" and is_path_proxied:
                # Path-based proxying: rewrite location redirect
                if v_str.startswith("/") and not v_str.startswith(f"/u/{universe_id}"):
                    v_str = f"/u/{universe_id}{v_str}"
                raw_headers.append((b"location", v_str.encode("latin-1")))
            else:
                raw_headers.append((raw_k, v_str.encode("latin-1")))
        
        async def stream_content() -> AsyncGenerator[bytes, None]:
            try:
                async for chunk in upstream_res.aiter_raw():
                    yield chunk
            finally:
                await upstream_res.aclose()
        
        response = StreamingResponse(
            stream_content(),
            status_code=upstream_res.status_code,
            media_type=media_type
        )
        response.raw_headers = raw_headers
        return response

    except httpx.ConnectError:
        logger.error(f"Failed to connect to backend at {target_url}")
        return Response(
            content=f"<html><head><title>Backend Unavailable</title></head><body style='background:#121212;color:#ef4444;font-family:sans-serif;padding:2rem;'><h2>Instance Hermes inaccessible</h2><p>Le backend sur <code>{backend_url}</code> ne répond pas. Vérifiez le tunnel Tailscale ou l'état du conteneur.</p></body></html>",
            status_code=502,
            media_type="text/html"
        )
    except Exception as e:
        logger.exception(f"Proxy error for {target_url}: {e}")
        return Response(
            content=f"Proxy Error: {str(e)}",
            status_code=500,
            media_type="text/plain"
        )