266 lines
9.3 KiB
Python
266 lines
9.3 KiB
Python
import httpx
|
|
import re
|
|
import asyncio
|
|
from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple
|
|
from fastapi import Request, Response, WebSocket, WebSocketDisconnect
|
|
from fastapi.responses import StreamingResponse
|
|
import websockets
|
|
from websockets.exceptions import ConnectionClosed
|
|
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"
|
|
}
|
|
|
|
WS_HOP_BY_HOP_HEADERS = {
|
|
"connection",
|
|
"upgrade",
|
|
"sec-websocket-key",
|
|
"sec-websocket-version",
|
|
"sec-websocket-extensions",
|
|
"sec-websocket-accept",
|
|
"host",
|
|
"origin"
|
|
}
|
|
|
|
_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:
|
|
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}"
|
|
|
|
is_dsh = "dsh-vps:3080" in backend_url or universe_id == "dsh"
|
|
|
|
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 in HOP_BY_HOP_HEADERS or k_str == "host":
|
|
continue
|
|
if is_dsh and k_str == "origin":
|
|
continue
|
|
req_headers.append((raw_k, raw_v))
|
|
|
|
# If DSH backend, replicate localhost:3080 Host/Origin for internal origin check
|
|
if is_dsh:
|
|
req_headers.append((b"host", b"localhost:3080"))
|
|
req_headers.append((b"origin", b"http://localhost:3080"))
|
|
else:
|
|
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
|
|
)
|
|
|
|
upstream_res = await transport.handle_async_request(upstream_req)
|
|
|
|
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:
|
|
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:
|
|
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 réseau interne 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"
|
|
)
|
|
|
|
async def proxy_websocket(
|
|
client_ws: WebSocket,
|
|
backend_url: str,
|
|
path: str,
|
|
universe_id: Optional[str] = None
|
|
):
|
|
ws_base = backend_url.replace("https://", "wss://").replace("http://", "ws://").rstrip("/")
|
|
sub_path = path.lstrip("/")
|
|
upstream_url = f"{ws_base}/{sub_path}" if sub_path else ws_base
|
|
|
|
if client_ws.url.query:
|
|
upstream_url = f"{upstream_url}?{client_ws.url.query}"
|
|
|
|
is_dsh = "dsh-vps:3080" in backend_url or universe_id == "dsh"
|
|
|
|
upstream_headers = {}
|
|
for k, v in client_ws.headers.items():
|
|
if k.lower() not in WS_HOP_BY_HOP_HEADERS:
|
|
upstream_headers[k] = v
|
|
|
|
if is_dsh:
|
|
upstream_headers["host"] = "localhost:3080"
|
|
origin_val = "http://localhost:3080"
|
|
else:
|
|
upstream_headers["x-forwarded-host"] = client_ws.headers.get("host", "")
|
|
upstream_headers["x-forwarded-proto"] = client_ws.url.scheme or "http"
|
|
origin_val = None
|
|
|
|
if client_ws.client:
|
|
upstream_headers["x-forwarded-for"] = client_ws.client.host
|
|
|
|
subprotocols_raw = client_ws.headers.get("sec-websocket-protocol", "")
|
|
subprotocols = [s.strip() for s in subprotocols_raw.split(",") if s.strip()] or None
|
|
|
|
try:
|
|
connect_kwargs = {
|
|
"additional_headers": upstream_headers,
|
|
"subprotocols": subprotocols,
|
|
"ping_interval": 20,
|
|
"ping_timeout": 20,
|
|
"max_size": 10 * 1024 * 1024
|
|
}
|
|
if origin_val:
|
|
connect_kwargs["origin"] = origin_val
|
|
|
|
async with websockets.connect(upstream_url, **connect_kwargs) as upstream_ws:
|
|
await client_ws.accept(subprotocol=upstream_ws.subprotocol)
|
|
|
|
async def client_to_upstream():
|
|
try:
|
|
while True:
|
|
msg = await client_ws.receive()
|
|
if "text" in msg:
|
|
await upstream_ws.send(msg["text"])
|
|
elif "bytes" in msg:
|
|
await upstream_ws.send(msg["bytes"])
|
|
elif msg.get("type") == "websocket.disconnect":
|
|
break
|
|
except (WebSocketDisconnect, ConnectionClosed, asyncio.CancelledError):
|
|
pass
|
|
|
|
async def upstream_to_client():
|
|
try:
|
|
async for msg in upstream_ws:
|
|
if isinstance(msg, str):
|
|
await client_ws.send_text(msg)
|
|
elif isinstance(msg, bytes):
|
|
await client_ws.send_bytes(msg)
|
|
except (WebSocketDisconnect, ConnectionClosed, asyncio.CancelledError):
|
|
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 (WebSocketDisconnect, ConnectionClosed):
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"WebSocket proxy bridge error for {upstream_url}: {e}")
|
|
try:
|
|
await client_ws.close(code=1011)
|
|
except Exception:
|
|
pass
|