151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
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_client: Optional[httpx.AsyncClient] = None
|
|
|
|
def get_http_client() -> httpx.AsyncClient:
|
|
global _http_client
|
|
if _http_client is None or _http_client.is_closed:
|
|
_http_client = httpx.AsyncClient(
|
|
timeout=httpx.Timeout(connect=5.0, read=120.0, write=60.0, pool=30.0),
|
|
follow_redirects=False,
|
|
limits=httpx.Limits(max_keepalive_connections=50, max_connections=100)
|
|
)
|
|
return _http_client
|
|
|
|
async def close_http_client():
|
|
global _http_client
|
|
if _http_client and not _http_client.is_closed:
|
|
await _http_client.aclose()
|
|
_http_client = 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}/
|
|
to strictly isolate session cookies between Hermes universes.
|
|
"""
|
|
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]:
|
|
client = get_http_client()
|
|
try:
|
|
res = await client.get(backend_url, timeout=3.0)
|
|
return {
|
|
"status": "online" if res.status_code < 500 else "degraded",
|
|
"status_code": res.status_code,
|
|
"latency_ms": int(res.elapsed.total_seconds() * 1000)
|
|
}
|
|
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:
|
|
client = get_http_client()
|
|
|
|
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 = {}
|
|
for key, value in request.headers.items():
|
|
if key.lower() not in HOP_BY_HOP_HEADERS and key.lower() != "host":
|
|
req_headers[key] = value
|
|
|
|
body = await request.body()
|
|
|
|
try:
|
|
upstream_req = client.build_request(
|
|
method=request.method,
|
|
url=target_url,
|
|
headers=req_headers,
|
|
content=body
|
|
)
|
|
|
|
upstream_res = await client.send(upstream_req, stream=True)
|
|
|
|
# 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 universe_id:
|
|
# Cloisonnement strict des cookies de session par univers
|
|
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 universe_id:
|
|
# Réécriture de la redirection si le backend renvoie vers la racine /
|
|
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"
|
|
)
|