class UiProxyHandler(tornado.websocket.WebSocketHandler):
"""Proxies one Process-backed UI; a websocket upgrade switches to frame relaying."""
SUPPORTED_METHODS: tuple[str, ...] = (
"GET",
"HEAD",
"POST",
"PUT",
"PATCH",
"DELETE",
"OPTIONS",
) # pyright: ignore[reportIncompatibleVariableOverride]
ui_name: str = ""
public: bool = False
auth: Authentication | None = None
_upstream_ws: tornado.websocket.WebSocketClientConnection | None = None
_client_subprotocols: list[str] | None = None
_relay: asyncio.Task[None] | None = None
@override
def initialize(self, ui_name: str, public: bool = False) -> None: # pyright: ignore[reportIncompatibleMethodOverride]
self.ui_name = ui_name
self.public = public
@override
def prepare(self) -> None:
self.auth = authenticate(self) if self.public else require_view(self, self.ui_name)
def _upstream_base(self) -> str | None:
processes: UiProcesses = self.application.settings["ui_processes"]
return processes.upstream(self.ui_name)
def _upstream_headers(self) -> tornado.httputil.HTTPHeaders:
headers = tornado.httputil.HTTPHeaders()
for name, value in self.request.headers.get_all():
lower = name.lower()
if lower in _NOT_FORWARDED or lower.startswith(("x-woodglue-", "sec-websocket-")):
continue
headers.add(name, value)
cookie = strip_cookie(self.request.headers.get("Cookie", ""), SESSION_COOKIE)
if cookie:
headers["Cookie"] = cookie
headers["X-Forwarded-Proto"] = public_scheme(self)
headers["X-Forwarded-Host"] = self.request.host
if self.request.remote_ip:
headers["X-Forwarded-For"] = self.request.remote_ip
if self.auth is not None:
store: AuthStore | None = self.application.settings.get("auth_store")
entry = self.application.settings["config"].uis.get(self.ui_name)
grants = entry.grants if entry is not None else None
for name, value in identity_headers(self.auth, store, self.ui_name, grants).items():
headers[name] = value
return headers
# HTTP
@override
async def get(self, *args: Any, **kwargs: Any) -> None:
if self.request.headers.get("Upgrade", "").lower() == "websocket":
await super().get(*args, **kwargs)
else:
await self._proxy_http()
@override
def head(self, *args: Any) -> Awaitable[None]: # pyright: ignore[reportIncompatibleMethodOverride]
return self._proxy_http()
@override
def post(self, *args: Any) -> Awaitable[None]: # pyright: ignore[reportIncompatibleMethodOverride]
return self._proxy_http()
@override
def put(self, *args: Any) -> Awaitable[None]: # pyright: ignore[reportIncompatibleMethodOverride]
return self._proxy_http()
@override
def patch(self, *args: Any) -> Awaitable[None]: # pyright: ignore[reportIncompatibleMethodOverride]
return self._proxy_http()
@override
def delete(self, *args: Any) -> Awaitable[None]: # pyright: ignore[reportIncompatibleMethodOverride]
return self._proxy_http()
@override
def options(self, *args: Any) -> Awaitable[None]: # pyright: ignore[reportIncompatibleMethodOverride]
return self._proxy_http()
async def _proxy_http(self) -> None:
base = self._upstream_base()
if base is None:
self.set_status(503)
self.finish(f"UI '{self.ui_name}' is starting or unavailable; try again shortly")
return
method = self.request.method or "GET"
relay = _ResponseRelay(self, method)
request = tornado.httpclient.HTTPRequest(
f"{base}{self.request.uri}",
method=method,
headers=self._upstream_headers(),
body=self.request.body if method not in ("GET", "HEAD", "OPTIONS") else None,
follow_redirects=False,
decompress_response=False,
allow_nonstandard_methods=True,
connect_timeout=10,
# No overall limit: event streams and long polls stay open as long as the UI wants.
request_timeout=0,
header_callback=relay.on_header_line,
streaming_callback=relay.on_chunk,
)
try:
resp = await tornado.httpclient.AsyncHTTPClient().fetch(request, raise_error=False)
failed = resp.code == 599
except OSError as exc:
logger.warning("UI '%s' upstream failed: %s", self.ui_name, exc)
failed = True
if failed and not relay.started:
self.set_status(502)
self.finish(f"UI '{self.ui_name}' did not answer")
return
self.finish()
# Websockets
@override
def select_subprotocol(self, subprotocols: list[str]) -> str | None:
self._client_subprotocols = subprotocols
return subprotocols[0] if subprotocols else None
@override
async def open(self, *args: Any, **kwargs: Any) -> None:
base = self._upstream_base()
if base is None:
self.close(1013, "UI is starting or unavailable")
return
ws_base = "ws" + base[len("http") :]
request = tornado.httpclient.HTTPRequest(
f"{ws_base}{self.request.uri}", headers=self._upstream_headers()
)
try:
self._upstream_ws = await tornado.websocket.websocket_connect(
request, subprotocols=self._client_subprotocols
)
except (OSError, tornado.httpclient.HTTPClientError) as exc:
logger.warning("UI '%s' websocket upstream failed: %s", self.ui_name, exc)
self.close(1011, "UI did not accept the connection")
return
if self.auth is not None and self.auth.session_secret:
sockets: UiSockets = self.application.settings["ui_sockets"]
sockets.add(self.auth.session_secret, self)
self._relay = asyncio.get_running_loop().create_task(self._relay_upstream())
async def _relay_upstream(self) -> None:
upstream = self._upstream_ws
assert upstream is not None
while True:
message = await upstream.read_message()
if message is None:
self.close()
return
try:
await self.write_message(message, binary=isinstance(message, bytes))
except tornado.websocket.WebSocketClosedError:
return
@override
async def on_message(self, message: str | bytes) -> None:
if self._upstream_ws is not None:
await self._upstream_ws.write_message(message, binary=isinstance(message, bytes))
@override
def on_close(self) -> None:
if self._upstream_ws is not None:
self._upstream_ws.close()
if self.auth is not None and self.auth.session_secret:
sockets: UiSockets = self.application.settings["ui_sockets"]
sockets.remove(self.auth.session_secret, self)