4b1ff38e51
hermes proxy attaches the operator's subscription credential to every forwarded request but accepted any Host header and any Origin. A loopback or specific-IP bind now serves only its own Host names, and a request whose Origin is not the proxy's own is refused with 403. Non-browser clients send no Origin and are unaffected; a wildcard bind keeps accepting any Host. (cherry picked from commit 3cf3870cfc2a795e439b1a8c0c53424f79c24505)
255 lines
12 KiB
Python
255 lines
12 KiB
Python
"""HTTP server that forwards OpenAI-compatible requests to a configured upstream.
|
|
|
|
A credential-attaching forwarder: request/response bodies are never mediated, logged, or rewritten.
|
|
The one shim: after a *clean* upstream EOF, a ``text/event-stream`` response that carried a terminal
|
|
``finish_reason`` or ``lastOne: true`` but omitted ``data: [DONE]`` gets a single ``[DONE]`` frame.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import signal
|
|
from typing import Optional
|
|
from urllib.parse import urlsplit
|
|
|
|
try:
|
|
import aiohttp
|
|
from aiohttp import web
|
|
AIOHTTP_AVAILABLE = True
|
|
except ImportError:
|
|
aiohttp = None # type: ignore[assignment]
|
|
web = None # type: ignore[assignment]
|
|
AIOHTTP_AVAILABLE = False
|
|
|
|
from hermes_cli.proxy.adapters.base import UpstreamAdapter, UpstreamCredential
|
|
from hermes_cli.proxy.sse_done import DONE_SSE_FRAME, SseDoneTracker, content_type_is_sse
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Stripped when forwarding upstream: ``host``/``content-length`` are recomputed by aiohttp,
|
|
# ``authorization`` is replaced with our bearer; everything else passes through.
|
|
_HOP_BY_HOP_HEADERS = frozenset({
|
|
"host", "content-length", "connection", "keep-alive", "proxy-authenticate",
|
|
"proxy-authorization", "te", "trailers", "transfer-encoding", "upgrade", "authorization",
|
|
})
|
|
# aiohttp recomputes Content-Encoding/Content-Length on stream — let it.
|
|
_RESPONSE_DROP_HEADERS = _HOP_BY_HOP_HEADERS | {"content-encoding", "content-length"}
|
|
|
|
DEFAULT_PORT = 8645
|
|
DEFAULT_HOST = "127.0.0.1"
|
|
# Mirrors api_server's MAX_REQUEST_BYTES (10 MB); client_max_size bounds every read path,
|
|
# including chunked bodies.
|
|
MAX_REQUEST_BYTES = 10_000_000
|
|
|
|
|
|
def _require_aiohttp() -> None:
|
|
if not AIOHTTP_AVAILABLE:
|
|
raise RuntimeError("aiohttp is required for `hermes proxy`. Run `hermes setup` to install it.")
|
|
|
|
|
|
def _json_error(status: int, message: str, code: str = "proxy_error") -> "web.Response":
|
|
"""OpenAI-style error JSON response."""
|
|
body = {"error": {"message": message, "type": code, "code": code}}
|
|
return web.json_response(body, status=status)
|
|
|
|
|
|
_LOOPBACK_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"})
|
|
_WILDCARD_HOSTS = frozenset({"", "0.0.0.0", "::"})
|
|
|
|
|
|
def _authority_hostname(authority: str) -> str:
|
|
"""Lowercased hostname of a Host-header authority; ``""`` when malformed (fails closed)."""
|
|
if not authority or any(c in authority for c in "@/\\?# \t"):
|
|
return ""
|
|
try:
|
|
parts = urlsplit("//" + authority)
|
|
parts.port # noqa: B018 — raises ValueError on a malformed port
|
|
except ValueError:
|
|
return ""
|
|
return (parts.hostname or "").lower()
|
|
|
|
|
|
def _local_request_error(host_header: str, origin: Optional[str], bound_host: str) -> Optional[str]:
|
|
"""Why a request must be refused, or None. The proxy attaches the operator's subscription
|
|
credential to whatever reaches it, so a web page in the operator's browser must not be able
|
|
to drive it: a loopback/specific-IP bind accepts only its own Host names (a DNS-rebound
|
|
hostname fails), and any Origin other than the proxy's own is a cross-site browser request.
|
|
Non-browser clients send no Origin and are unaffected. A wildcard bind is an explicit LAN
|
|
opt-in, so only the Origin rule applies there."""
|
|
hostname = _authority_hostname(host_header)
|
|
bound = bound_host.strip("[]").lower()
|
|
if not hostname or (bound not in _WILDCARD_HOSTS and hostname not in _LOOPBACK_HOSTS | {bound}):
|
|
return "host_not_allowed"
|
|
if origin is not None and origin.strip().lower() != f"http://{host_header.strip().lower()}":
|
|
return "origin_not_allowed"
|
|
return None
|
|
|
|
|
|
def _filter_headers(headers, drop: frozenset = _HOP_BY_HOP_HEADERS) -> dict:
|
|
"""Strip hop-by-hop (+ auth) headers; ``drop`` widens the set for upstream responses."""
|
|
return {key: value for key, value in headers.items() if key.lower() not in drop}
|
|
|
|
|
|
async def _open_upstream(request: "web.Request", rel_path: str, body: bytes, cred: UpstreamCredential):
|
|
"""Send the request upstream with ``cred``; returns ``(session, response)`` or
|
|
``(error_response, None)``."""
|
|
upstream_url = f"{cred.base_url.rstrip('/')}{rel_path}"
|
|
if request.query_string: # preserved verbatim
|
|
upstream_url = f"{upstream_url}?{request.query_string}"
|
|
fwd_headers = _filter_headers(request.headers)
|
|
fwd_headers["Authorization"] = f"{cred.token_type} {cred.bearer}"
|
|
logger.debug("proxy: forwarding %s %s -> %s (body=%d bytes)", request.method, rel_path, upstream_url, len(body))
|
|
try:
|
|
session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None, sock_connect=15, sock_read=300))
|
|
except Exception as exc: # pragma: no cover - aiohttp setup issue
|
|
return _json_error(500, f"proxy session init failed: {exc}"), None
|
|
try:
|
|
upstream_resp = await session.request(
|
|
request.method, upstream_url, data=body if body else None, headers=fwd_headers, allow_redirects=False
|
|
)
|
|
except RuntimeError as exc:
|
|
await session.close()
|
|
return _json_error(500, str(exc)), None
|
|
except aiohttp.ClientError as exc:
|
|
await session.close()
|
|
logger.warning("proxy: upstream connection failed: %s", exc)
|
|
return _json_error(502, f"upstream connection failed: {exc}", code="upstream_unreachable"), None
|
|
except asyncio.TimeoutError:
|
|
await session.close()
|
|
return _json_error(504, "upstream request timed out", code="upstream_timeout"), None
|
|
except Exception:
|
|
await session.close()
|
|
raise
|
|
return session, upstream_resp
|
|
|
|
|
|
async def _stream_back(request: "web.Request", session, upstream_resp) -> "web.StreamResponse":
|
|
"""Relay status + filtered headers, then the body chunk-by-chunk, appending a missing SSE
|
|
``[DONE]`` only after a clean EOF."""
|
|
resp = web.StreamResponse(
|
|
status=upstream_resp.status, headers=_filter_headers(upstream_resp.headers, _RESPONSE_DROP_HEADERS)
|
|
)
|
|
await resp.prepare(request)
|
|
done_tracker: Optional[SseDoneTracker] = None
|
|
if content_type_is_sse(upstream_resp.headers):
|
|
done_tracker = SseDoneTracker()
|
|
try:
|
|
async for chunk in upstream_resp.content.iter_any():
|
|
if chunk:
|
|
if done_tracker is not None:
|
|
done_tracker.feed(chunk)
|
|
await resp.write(chunk)
|
|
if done_tracker is not None and done_tracker.should_append_done():
|
|
try:
|
|
await resp.write(DONE_SSE_FRAME)
|
|
except Exception as exc: # client hung up at EOF — harmless
|
|
logger.debug("proxy: DONE append skipped: %s", exc)
|
|
except (aiohttp.ClientError, asyncio.CancelledError, OSError) as exc:
|
|
if done_tracker is not None:
|
|
done_tracker.mark_interrupted()
|
|
logger.warning("proxy: streaming interrupted: %s", exc)
|
|
finally:
|
|
upstream_resp.release()
|
|
await session.close()
|
|
await resp.write_eof()
|
|
return resp
|
|
|
|
|
|
def create_app(adapter: UpstreamAdapter, bound_host: str = DEFAULT_HOST) -> "web.Application":
|
|
"""Build the aiohttp application bound to a specific upstream adapter.
|
|
|
|
Every adapter method is synchronous and blocking (the Nous adapter takes the 15s cross-process
|
|
``_auth_store_lock()`` and may POST a token refresh; xAI rotates its key pool under a lock),
|
|
so all three are run via ``asyncio.to_thread`` — a contended lock or refresh must never freeze
|
|
the single loop and every other in-flight streaming completion.
|
|
"""
|
|
_require_aiohttp()
|
|
|
|
@web.middleware
|
|
async def local_only(request: "web.Request", handler):
|
|
refusal = _local_request_error(request.headers.get("Host", ""), request.headers.get("Origin"), bound_host)
|
|
if refusal:
|
|
return _json_error(403, "Request refused: only local, non-browser clients may use this proxy.", code=refusal)
|
|
return await handler(request)
|
|
|
|
app = web.Application(client_max_size=MAX_REQUEST_BYTES, middlewares=[local_only])
|
|
# AppKey: forward-compat with aiohttp versions that strip bare-string keys.
|
|
app[web.AppKey("adapter", UpstreamAdapter)] = adapter
|
|
|
|
async def handle_health(request: "web.Request") -> "web.Response":
|
|
authenticated = await asyncio.to_thread(adapter.is_authenticated)
|
|
return web.json_response({"status": "ok", "upstream": adapter.display_name, "authenticated": authenticated})
|
|
|
|
async def handle_proxy(request: "web.Request") -> "web.StreamResponse":
|
|
rel_path = "/" + request.match_info.get("tail", "").lstrip("/")
|
|
if rel_path not in adapter.allowed_paths:
|
|
allowed = ", ".join(sorted(adapter.allowed_paths))
|
|
return _json_error(
|
|
404, f"Path /v1{rel_path} is not forwarded by this proxy. Allowed: {allowed}", code="path_not_allowed"
|
|
)
|
|
try:
|
|
cred = await asyncio.to_thread(adapter.get_credential)
|
|
except Exception as exc:
|
|
logger.warning("proxy: credential resolution failed: %s", exc)
|
|
return _json_error(401, str(exc), code="upstream_auth_failed")
|
|
# Body read into memory once (chat/embeddings payloads are small); switch to streaming
|
|
# if large multipart uploads ever need forwarding.
|
|
body = await request.read()
|
|
session, upstream_resp = await _open_upstream(request, rel_path, body, cred)
|
|
if upstream_resp is None:
|
|
return session
|
|
if upstream_resp.status in {401, 429}:
|
|
# One-shot retry with a refreshed/rotated credential (Nous: unconditional refresh
|
|
# POST under the auth lock; xAI: pool rotation).
|
|
try:
|
|
retry_cred = await asyncio.to_thread(
|
|
adapter.get_retry_credential, failed_credential=cred, status_code=upstream_resp.status
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("proxy: retry credential resolution failed: %s", exc)
|
|
retry_cred = None
|
|
if retry_cred is not None:
|
|
upstream_resp.release()
|
|
await session.close()
|
|
session, upstream_resp = await _open_upstream(request, rel_path, body, retry_cred)
|
|
if upstream_resp is None:
|
|
return session
|
|
return await _stream_back(request, session, upstream_resp)
|
|
|
|
app.router.add_get("/health", handle_health) # never goes upstream
|
|
app.router.add_route("*", "/v1/{tail:.*}", handle_proxy) # forwards if the path is allowed
|
|
return app
|
|
|
|
|
|
async def run_server(
|
|
adapter: UpstreamAdapter,
|
|
host: str = DEFAULT_HOST,
|
|
port: int = DEFAULT_PORT,
|
|
shutdown_event: Optional[asyncio.Event] = None,
|
|
) -> None:
|
|
"""Run the proxy in the current event loop until shutdown_event is set."""
|
|
_require_aiohttp()
|
|
app = create_app(adapter, bound_host=host)
|
|
runner = web.AppRunner(app, access_log=None)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, host=host, port=port)
|
|
await site.start()
|
|
logger.info("proxy: listening on http://%s:%d/v1 -> %s", host, port, adapter.display_name)
|
|
stop_event = shutdown_event or asyncio.Event()
|
|
if shutdown_event is None: # we own the loop's lifetime → wire signal handlers
|
|
loop = asyncio.get_running_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
try:
|
|
loop.add_signal_handler(sig, stop_event.set) # windows-footgun: ok
|
|
except NotImplementedError:
|
|
pass # Windows / restricted envs — Ctrl+C still raises KeyboardInterrupt
|
|
try:
|
|
await stop_event.wait()
|
|
finally:
|
|
logger.info("proxy: shutting down")
|
|
await runner.cleanup()
|
|
|
|
|
|
__all__ = ["create_app", "run_server", "DEFAULT_HOST", "DEFAULT_PORT", "AIOHTTP_AVAILABLE"]
|