| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import http.client |
| 5 | from http.cookies import SimpleCookie |
| 6 | from urllib.parse import parse_qs, quote, unquote, urlsplit |
| 7 | |
| 8 | from flask.sessions import SecureCookieSessionInterface |
| 9 | from starlette.requests import Request |
| 10 | from starlette.responses import JSONResponse, PlainTextResponse, RedirectResponse, Response |
| 11 | from starlette.types import Receive, Scope, Send |
| 12 | from starlette.websockets import WebSocket |
| 13 | from wsproto import ConnectionType, WSConnection |
| 14 | from wsproto.events import ( |
| 15 | AcceptConnection, |
| 16 | BytesMessage, |
| 17 | CloseConnection, |
| 18 | Ping, |
| 19 | RejectConnection, |
| 20 | Request as WebSocketRequest, |
| 21 | TextMessage, |
| 22 | ) |
| 23 | |
| 24 | from helpers import login, virtual_desktop |
| 25 | |
| 26 | |
| 27 | HOP_BY_HOP_HEADERS = { |
| 28 | "connection", |
| 29 | "content-length", |
| 30 | "cookie", |
| 31 | "host", |
| 32 | "keep-alive", |
| 33 | "proxy-authenticate", |
| 34 | "proxy-authorization", |
| 35 | "te", |
| 36 | "trailer", |
| 37 | "transfer-encoding", |
| 38 | "upgrade", |
| 39 | } |
| 40 | |
| 41 | |
| 42 | XPRA_MENU_CUSTOM_PATCH = b""" |
| 43 | ;(function () { |
| 44 | function a0DesktopElement(selector) { |
| 45 | return document.querySelector(selector); |
| 46 | } |
| 47 | |
| 48 | window.noWindowList = function noWindowList() { |
| 49 | const openWindows = a0DesktopElement("#open_windows"); |
| 50 | if (openWindows) openWindows.remove(); |
| 51 | }; |
| 52 | |
| 53 | const originalAddWindowListItem = window.addWindowListItem; |
| 54 | if (typeof originalAddWindowListItem === "function" && !originalAddWindowListItem.__a0SafeWindowList) { |
| 55 | const safeAddWindowListItem = function addWindowListItem(...args) { |
| 56 | if (!a0DesktopElement("#open_windows_list")) return undefined; |
| 57 | return originalAddWindowListItem.apply(this, args); |
| 58 | }; |
| 59 | safeAddWindowListItem.__a0SafeWindowList = true; |
| 60 | window.addWindowListItem = safeAddWindowListItem; |
| 61 | } |
| 62 | }()); |
| 63 | """ |
| 64 | |
| 65 | XPRA_WINDOW_OFFSET_WARNING = b'&&this.warn("window does not fit in canvas, offsets: ",x,y)' |
| 66 | XPRA_WINDOW_OFFSET_WARNING_PATCH = b'&&false&&this.warn("window does not fit in canvas, offsets: ",x,y)' |
| 67 | XPRA_WINDOW_SCRIPT = b'src="js/Window.js"' |
| 68 | XPRA_WINDOW_SCRIPT_PATCH = b'src="js/Window.js?a0_desktop_patch=20260506"' |
| 69 | |
| 70 | |
| 71 | class VirtualDesktopGateway: |
| 72 | def __init__(self, flask_app=None, mount_path: str = "/desktop") -> None: |
| 73 | self.flask_app = flask_app |
| 74 | self.mount_path = "/" + mount_path.strip("/") |
| 75 | |
| 76 | async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: |
| 77 | if scope["type"] == "websocket": |
| 78 | await self.websocket(scope, receive, send) |
| 79 | return |
| 80 | if scope["type"] == "http": |
| 81 | await self.http(scope, receive, send) |
| 82 | return |
| 83 | await PlainTextResponse("Unsupported scope", status_code=500)(scope, receive, send) |
| 84 | |
| 85 | async def http(self, scope: Scope, receive: Receive, send: Send) -> None: |
| 86 | if not self.is_authorized(scope): |
| 87 | await PlainTextResponse("Authentication required", status_code=401)(scope, receive, send) |
| 88 | return |
| 89 | |
| 90 | path = self.relative_path(scope) |
| 91 | if path in {"", "/"}: |
| 92 | await RedirectResponse(f"{self.mount_path}/health")(scope, receive, send) |
| 93 | return |
| 94 | if path == "/health": |
| 95 | await JSONResponse(virtual_desktop.collect_status())(scope, receive, send) |
| 96 | return |
| 97 | if path == "/resize": |
| 98 | await self.resize(scope, receive, send) |
| 99 | return |
| 100 | |
| 101 | session_request = self.session_request(path) |
| 102 | if not session_request: |
| 103 | await PlainTextResponse("Desktop session not found.", status_code=404)(scope, receive, send) |
| 104 | return |
| 105 | token, upstream_path = session_request |
| 106 | if upstream_path in {"", "/"}: |
| 107 | await RedirectResponse(self.session_index_url(token))(scope, receive, send) |
| 108 | return |
| 109 | await self.proxy_http(scope, receive, send, token, upstream_path) |
| 110 | |
| 111 | async def resize(self, scope: Scope, receive: Receive, send: Send) -> None: |
| 112 | query = self.query(scope) |
| 113 | payload: dict[str, object] = {} |
| 114 | if scope.get("method") == "POST": |
| 115 | try: |
| 116 | payload = await Request(scope, receive).json() |
| 117 | except Exception: |
| 118 | payload = {} |
| 119 | token = str(payload.get("token") or query.get("token", [""])[0]) |
| 120 | width = payload.get("width") or query.get("width", [0])[0] |
| 121 | height = payload.get("height") or query.get("height", [0])[0] |
| 122 | try: |
| 123 | result = virtual_desktop.resize_session(token, int(float(width)), int(float(height))) |
| 124 | except (TypeError, ValueError): |
| 125 | result = {"ok": False, "error": "Invalid virtual desktop size."} |
| 126 | await JSONResponse(result, status_code=200 if result.get("ok") else 400)(scope, receive, send) |
| 127 | |
| 128 | async def proxy_http( |
| 129 | self, |
| 130 | scope: Scope, |
| 131 | receive: Receive, |
| 132 | send: Send, |
| 133 | token: str, |
| 134 | upstream_path: str, |
| 135 | ) -> None: |
| 136 | endpoint = virtual_desktop.proxy_for_token(token) |
| 137 | if not endpoint: |
| 138 | await PlainTextResponse("Desktop session not found.", status_code=404)(scope, receive, send) |
| 139 | return |
| 140 | |
| 141 | body = await Request(scope, receive).body() |
| 142 | try: |
| 143 | status, headers, content = await asyncio.to_thread( |
| 144 | self.fetch_http, |
| 145 | endpoint, |
| 146 | upstream_path, |
| 147 | scope.get("query_string", b"").decode("latin-1"), |
| 148 | str(scope.get("method") or "GET"), |
| 149 | self.proxy_request_headers(scope), |
| 150 | body, |
| 151 | ) |
| 152 | content = self.proxy_response_content(upstream_path, content) |
| 153 | await Response( |
| 154 | content, |
| 155 | status_code=status, |
| 156 | headers=self.proxy_response_headers(headers, token), |
| 157 | )(scope, receive, send) |
| 158 | except (http.client.HTTPException, OSError, asyncio.TimeoutError): |
| 159 | await PlainTextResponse("Desktop proxy is unavailable.", status_code=502)(scope, receive, send) |
| 160 | |
| 161 | async def websocket(self, scope: Scope, receive: Receive, send: Send) -> None: |
| 162 | websocket = WebSocket(scope, receive=receive, send=send) |
| 163 | if not self.is_authorized(scope): |
| 164 | await websocket.close(code=1008) |
| 165 | return |
| 166 | |
| 167 | session_request = self.session_request(self.relative_path(scope)) |
| 168 | if not session_request: |
| 169 | await websocket.close(code=1008) |
| 170 | return |
| 171 | token, upstream_path = session_request |
| 172 | endpoint = virtual_desktop.proxy_for_token(token) |
| 173 | if not endpoint: |
| 174 | await websocket.close(code=1008) |
| 175 | return |
| 176 | |
| 177 | target = self.upstream_target(upstream_path or "/", scope.get("query_string", b"")) |
| 178 | try: |
| 179 | reader, writer, upstream, subprotocol = await self.open_websocket( |
| 180 | endpoint, |
| 181 | target, |
| 182 | tuple(scope.get("subprotocols") or ()), |
| 183 | ) |
| 184 | await websocket.accept(subprotocol=subprotocol) |
| 185 | await asyncio.gather( |
| 186 | self.browser_to_xpra(websocket, upstream, writer), |
| 187 | self.xpra_to_browser(websocket, upstream, reader, writer), |
| 188 | ) |
| 189 | except Exception: |
| 190 | try: |
| 191 | await websocket.close(code=1011) |
| 192 | except Exception: |
| 193 | pass |
| 194 | |
| 195 | async def open_websocket( |
| 196 | self, |
| 197 | endpoint: virtual_desktop.VirtualDesktopEndpoint, |
| 198 | target: str, |
| 199 | subprotocols: tuple[str, ...], |
| 200 | ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter, WSConnection, str | None]: |
| 201 | reader, writer = await asyncio.open_connection(endpoint.host, endpoint.port) |
| 202 | upstream = WSConnection(ConnectionType.CLIENT) |
| 203 | writer.write( |
| 204 | upstream.send( |
| 205 | WebSocketRequest( |
| 206 | host=f"{endpoint.host}:{endpoint.port}", |
| 207 | target=target, |
| 208 | subprotocols=list(subprotocols), |
| 209 | ), |
| 210 | ), |
| 211 | ) |
| 212 | await writer.drain() |
| 213 | |
| 214 | while True: |
| 215 | data = await asyncio.wait_for(reader.read(65536), timeout=10) |
| 216 | if not data: |
| 217 | raise ConnectionError("Xpra WebSocket handshake closed early.") |
| 218 | upstream.receive_data(data) |
| 219 | for event in upstream.events(): |
| 220 | if isinstance(event, AcceptConnection): |
| 221 | return reader, writer, upstream, event.subprotocol |
| 222 | if isinstance(event, RejectConnection): |
| 223 | raise ConnectionError(f"Xpra rejected WebSocket handshake with HTTP {event.status_code}.") |
| 224 | if isinstance(event, CloseConnection): |
| 225 | raise ConnectionError("Xpra closed WebSocket handshake.") |
| 226 | |
| 227 | async def browser_to_xpra( |
| 228 | self, |
| 229 | websocket: WebSocket, |
| 230 | upstream: WSConnection, |
| 231 | writer: asyncio.StreamWriter, |
| 232 | ) -> None: |
| 233 | try: |
| 234 | while True: |
| 235 | message = await websocket.receive() |
| 236 | if message["type"] == "websocket.disconnect": |
| 237 | writer.write(upstream.send(CloseConnection(code=1000))) |
| 238 | await writer.drain() |
| 239 | return |
| 240 | if message.get("bytes") is not None: |
| 241 | writer.write(upstream.send(BytesMessage(data=message["bytes"]))) |
| 242 | elif message.get("text") is not None: |
| 243 | writer.write(upstream.send(TextMessage(data=str(message["text"])))) |
| 244 | await writer.drain() |
| 245 | finally: |
| 246 | writer.close() |
| 247 | |
| 248 | async def xpra_to_browser( |
| 249 | self, |
| 250 | websocket: WebSocket, |
| 251 | upstream: WSConnection, |
| 252 | reader: asyncio.StreamReader, |
| 253 | writer: asyncio.StreamWriter, |
| 254 | ) -> None: |
| 255 | try: |
| 256 | while True: |
| 257 | data = await reader.read(65536) |
| 258 | if not data: |
| 259 | return |
| 260 | upstream.receive_data(data) |
| 261 | for event in upstream.events(): |
| 262 | if isinstance(event, BytesMessage): |
| 263 | await websocket.send_bytes(event.data) |
| 264 | elif isinstance(event, TextMessage): |
| 265 | await websocket.send_text(event.data) |
| 266 | elif isinstance(event, Ping): |
| 267 | writer.write(upstream.send(event.response())) |
| 268 | await writer.drain() |
| 269 | elif isinstance(event, CloseConnection): |
| 270 | writer.write(upstream.send(event.response())) |
| 271 | await writer.drain() |
| 272 | return |
| 273 | finally: |
| 274 | try: |
| 275 | await websocket.close() |
| 276 | except Exception: |
| 277 | pass |
| 278 | writer.close() |
| 279 | |
| 280 | def session_request(self, path: str) -> tuple[str, str] | None: |
| 281 | prefix = "/session/" |
| 282 | if not path.startswith(prefix): |
| 283 | return None |
| 284 | rest = path[len(prefix):] |
| 285 | token, separator, upstream_path = rest.partition("/") |
| 286 | token = unquote(token) |
| 287 | if not token: |
| 288 | return None |
| 289 | return token, f"/{upstream_path}" if separator else "/" |
| 290 | |
| 291 | def session_index_url(self, token: str) -> str: |
| 292 | quoted_token = quote(str(token), safe="") |
| 293 | base_path = f"{self.mount_path}/session/{quoted_token}/" |
| 294 | return f"{base_path}index.html?path={quote(base_path, safe='')}" |
| 295 | |
| 296 | def fetch_http( |
| 297 | self, |
| 298 | endpoint: virtual_desktop.VirtualDesktopEndpoint, |
| 299 | upstream_path: str, |
| 300 | query: str, |
| 301 | method: str, |
| 302 | headers: dict[str, str], |
| 303 | body: bytes, |
| 304 | ) -> tuple[int, dict[str, str], bytes]: |
| 305 | connection = http.client.HTTPConnection(endpoint.host, endpoint.port, timeout=60) |
| 306 | try: |
| 307 | connection.request( |
| 308 | method, |
| 309 | self.upstream_target(upstream_path, query.encode("latin-1")), |
| 310 | body=body or None, |
| 311 | headers={**headers, "Connection": "close"}, |
| 312 | ) |
| 313 | response = connection.getresponse() |
| 314 | return response.status, dict(response.getheaders()), response.read() |
| 315 | finally: |
| 316 | connection.close() |
| 317 | |
| 318 | def upstream_target(self, upstream_path: str, query_string: bytes) -> str: |
| 319 | query = query_string.decode("latin-1") |
| 320 | target = upstream_path or "/" |
| 321 | return f"{target}?{query}" if query else target |
| 322 | |
| 323 | def proxy_request_headers(self, scope: Scope) -> dict[str, str]: |
| 324 | headers: dict[str, str] = {} |
| 325 | for raw_name, raw_value in scope.get("headers", []): |
| 326 | name = raw_name.decode("latin-1") |
| 327 | lower = name.lower() |
| 328 | if lower in HOP_BY_HOP_HEADERS or lower == "origin" or lower.startswith("sec-websocket"): |
| 329 | continue |
| 330 | headers[name] = raw_value.decode("latin-1") |
| 331 | return headers |
| 332 | |
| 333 | def proxy_response_headers(self, headers: dict[str, str], token: str) -> dict[str, str]: |
| 334 | response_headers: dict[str, str] = {} |
| 335 | for name, value in dict(headers).items(): |
| 336 | lower = name.lower() |
| 337 | if lower in HOP_BY_HOP_HEADERS: |
| 338 | continue |
| 339 | if lower == "location": |
| 340 | value = self.rewrite_location(str(value), token) |
| 341 | response_headers[name] = str(value) |
| 342 | return response_headers |
| 343 | |
| 344 | def proxy_response_content(self, upstream_path: str, content: bytes) -> bytes: |
| 345 | if upstream_path.endswith("/index.html") and XPRA_WINDOW_SCRIPT in content: |
| 346 | content = content.replace(XPRA_WINDOW_SCRIPT, XPRA_WINDOW_SCRIPT_PATCH) |
| 347 | if upstream_path.endswith("/js/MenuCustom.js") and XPRA_MENU_CUSTOM_PATCH not in content: |
| 348 | return content + XPRA_MENU_CUSTOM_PATCH |
| 349 | if upstream_path.endswith("/js/Window.js") and XPRA_WINDOW_OFFSET_WARNING in content: |
| 350 | return content.replace(XPRA_WINDOW_OFFSET_WARNING, XPRA_WINDOW_OFFSET_WARNING_PATCH) |
| 351 | return content |
| 352 | |
| 353 | def rewrite_location(self, location: str, token: str) -> str: |
| 354 | quoted_token = quote(str(token), safe="") |
| 355 | prefix = f"{self.mount_path}/session/{quoted_token}" |
| 356 | parsed = urlsplit(location) |
| 357 | if parsed.scheme in {"http", "https"} and parsed.hostname in {"127.0.0.1", "localhost"}: |
| 358 | path = parsed.path or "/" |
| 359 | query = f"?{parsed.query}" if parsed.query else "" |
| 360 | return f"{prefix}{path}{query}" |
| 361 | if location.startswith("/"): |
| 362 | return f"{prefix}{location}" |
| 363 | return location |
| 364 | |
| 365 | def relative_path(self, scope: Scope) -> str: |
| 366 | raw_path = scope.get("raw_path") |
| 367 | path = raw_path.decode("latin-1") if raw_path else str(scope.get("path") or "") |
| 368 | if path.startswith(self.mount_path): |
| 369 | path = path[len(self.mount_path):] |
| 370 | return path or "/" |
| 371 | |
| 372 | def query(self, scope: Scope) -> dict[str, list[str]]: |
| 373 | return parse_qs(scope.get("query_string", b"").decode("latin-1"), keep_blank_values=True) |
| 374 | |
| 375 | def is_authorized(self, scope: Scope) -> bool: |
| 376 | credentials_hash = login.get_credentials_hash() |
| 377 | if not credentials_hash: |
| 378 | return True |
| 379 | if not self.flask_app: |
| 380 | return False |
| 381 | serializer = SecureCookieSessionInterface().get_signing_serializer(self.flask_app) |
| 382 | if not serializer: |
| 383 | return False |
| 384 | cookie_header = dict(scope.get("headers", [])).get(b"cookie", b"").decode("latin-1") |
| 385 | if not cookie_header: |
| 386 | return False |
| 387 | cookies = SimpleCookie() |
| 388 | cookies.load(cookie_header) |
| 389 | session_cookie = cookies.get(self.flask_app.config.get("SESSION_COOKIE_NAME", "session")) |
| 390 | if not session_cookie: |
| 391 | return False |
| 392 | try: |
| 393 | session_data = serializer.loads(session_cookie.value) |
| 394 | except Exception: |
| 395 | return False |
| 396 | return session_data.get("authentication") == credentials_hash |
| 397 | |
| 398 | |
| 399 | def install_route_hooks() -> None: |
| 400 | from helpers.ui_server import UiServerRuntime |
| 401 | |
| 402 | if getattr(UiServerRuntime, "_a0_virtual_desktop_route_hooks_installed", False): |
| 403 | return |
| 404 | |
| 405 | original_build_asgi_app = UiServerRuntime.build_asgi_app |
| 406 | |
| 407 | def build_asgi_app(self, startup_monitor): |
| 408 | from socketio import ASGIApp |
| 409 | from starlette.applications import Starlette |
| 410 | from starlette.routing import Mount |
| 411 | from uvicorn.middleware.wsgi import WSGIMiddleware |
| 412 | |
| 413 | from helpers import fasta2a_server, mcp_server |
| 414 | |
| 415 | with startup_monitor.stage("wsgi.middleware.create"): |
| 416 | wsgi_app = WSGIMiddleware(self.webapp) |
| 417 | |
| 418 | with startup_monitor.stage("mcp.proxy.init"): |
| 419 | mcp_app = mcp_server.DynamicMcpProxy.get_instance() |
| 420 | |
| 421 | with startup_monitor.stage("a2a.proxy.init"): |
| 422 | a2a_app = fasta2a_server.DynamicA2AProxy.get_instance() |
| 423 | |
| 424 | with startup_monitor.stage("starlette.app.create"): |
| 425 | starlette_app = Starlette( |
| 426 | routes=[ |
| 427 | Mount("/desktop", app=VirtualDesktopGateway(self.webapp, "/desktop")), |
| 428 | Mount("/mcp", app=mcp_app), |
| 429 | Mount("/a2a", app=a2a_app), |
| 430 | Mount("/", app=wsgi_app), |
| 431 | ], |
| 432 | lifespan=startup_monitor.lifespan(), |
| 433 | ) |
| 434 | |
| 435 | with startup_monitor.stage("socketio.asgi.create"): |
| 436 | return ASGIApp(self.socketio_server, other_asgi_app=starlette_app) |
| 437 | |
| 438 | UiServerRuntime.build_asgi_app = build_asgi_app |
| 439 | UiServerRuntime._a0_virtual_desktop_route_hooks_installed = True |
| 440 | UiServerRuntime._a0_virtual_desktop_original_build_asgi_app = original_build_asgi_app |
| 441 | |
| 442 | |
| 443 | def is_installed() -> bool: |
| 444 | from helpers.ui_server import UiServerRuntime |
| 445 | |
| 446 | return bool(getattr(UiServerRuntime, "_a0_virtual_desktop_route_hooks_installed", False)) |