Исправлена обработка отключения noVNC WebSocket

scripts/patch_web_console.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 from pathlib import Path
-import re
 import sys
 
 app_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path('/opt/virtuality/web/app.py')
@@ -8,8 +7,75 @@ if not app_path.exists():
     raise SystemExit(f'app.py not found: {app_path}')
 
 text = app_path.read_text()
+
+safe_proxy = '''# Virtuality noVNC console patch
+async def proxy_vnc_to_websocket(reader: asyncio.StreamReader, websocket: WebSocket) -> None:
+    while True:
+        data = await reader.read(65536)
+        if not data:
+            break
+        try:
+            await websocket.send_bytes(data)
+        except (WebSocketDisconnect, RuntimeError, ConnectionError):
+            break
+
+
+async def proxy_websocket_to_vnc(websocket: WebSocket, writer: asyncio.StreamWriter) -> None:
+    while True:
+        try:
+            message = await websocket.receive()
+        except (WebSocketDisconnect, RuntimeError, ConnectionError):
+            break
+        if message.get("type") == "websocket.disconnect":
+            break
+        if message.get("bytes") is not None:
+            writer.write(message["bytes"])
+        elif message.get("text") is not None:
+            writer.write(message["text"].encode())
+        try:
+            await writer.drain()
+        except (RuntimeError, ConnectionError, BrokenPipeError):
+            break
+'''
+
+safe_gather = '''    tasks = [
+        asyncio.create_task(proxy_vnc_to_websocket(reader, websocket)),
+        asyncio.create_task(proxy_websocket_to_vnc(websocket, writer)),
+    ]
+    try:
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+        await asyncio.gather(*done, return_exceptions=True)
+    finally:
+        writer.close()
+        try:
+            await writer.wait_closed()
+        except Exception:
+            pass
+'''
+
 if 'Virtuality noVNC console patch' in text:
-    print('noVNC console patch already applied')
+    start = text.find('# Virtuality noVNC console patch')
+    marker = '\n\n@app.get("/login", response_class=HTMLResponse)'
+    end = text.find(marker, start)
+    if start != -1 and end != -1:
+        text = text[:start] + safe_proxy + text[end:]
+
+    old_gather = '''    try:
+        await asyncio.gather(proxy_vnc_to_websocket(reader, websocket), proxy_websocket_to_vnc(websocket, writer))
+    except (WebSocketDisconnect, asyncio.CancelledError, ConnectionError):
+        pass
+    finally:
+        writer.close()
+        await writer.wait_closed()
+'''
+    if old_gather in text:
+        text = text.replace(old_gather, safe_gather, 1)
+
+    app_path.write_text(text)
+    print(f'noVNC console patch upgraded: {app_path}')
     raise SystemExit(0)
 
 text = text.replace('import crypt\n', 'import asyncio\nimport crypt\n', 1)
@@ -25,11 +91,10 @@ text = text.replace(old_get_user, new_get_user, 1)
 insert_after_vm_ip = '''def vm_vnc_display(name: str) -> str:\n    return run_cmd(["virsh", "vncdisplay", name], timeout=8)["stdout"] or "not available"\n\n\ndef vnc_display_to_port(display: str) -> int | None:\n    value = (display or "").strip()\n    if not value or value == "not available":\n        return None\n    match = re.search(r":(\\d+)$", value)\n    if not match:\n        return None\n    display_number = int(match.group(1))\n    if display_number >= 5900:\n        return display_number\n    return 5900 + display_number\n\n\ndef console_info(name: str) -> dict[str, Any]:\n    display = vm_vnc_display(name)\n    port = vnc_display_to_port(display)\n    has_novnc = bool(NOVNC_DIR and (NOVNC_DIR / "vnc.html").exists())\n    token = None\n    url = None\n    if port and has_novnc:\n        token = console_serializer.dumps({"vm": name, "port": port})\n        url = f"/novnc/vnc.html?autoconnect=1&resize=scale&path=console/ws/{token}"\n    return {"vm": name, "display": display, "port": port, "has_novnc": has_novnc, "novnc_dir": str(NOVNC_DIR) if NOVNC_DIR else "not installed", "url": url}\n\n\n'''
 text = text.replace('def vm_details(name: str) -> dict[str, Any]:\n', insert_after_vm_ip + 'def vm_details(name: str) -> dict[str, Any]:\n', 1)
 text = text.replace('"vnc": run_cmd(["virsh", "vncdisplay", name], timeout=8)["stdout"] or "not available",', '"vnc": vm_vnc_display(name),', 1)
+text = text.replace('@app.get("/login", response_class=HTMLResponse)\n', safe_proxy + '\n@app.get("/login", response_class=HTMLResponse)\n', 1)
 
-proxy_funcs = '''\n# Virtuality noVNC console patch\nasync def proxy_vnc_to_websocket(reader: asyncio.StreamReader, websocket: WebSocket) -> None:\n    while True:\n        data = await reader.read(65536)\n        if not data:\n            break\n        await websocket.send_bytes(data)\n\n\nasync def proxy_websocket_to_vnc(websocket: WebSocket, writer: asyncio.StreamWriter) -> None:\n    while True:\n        message = await websocket.receive()\n        if message.get("type") == "websocket.disconnect":\n            break\n        if message.get("bytes") is not None:\n            writer.write(message["bytes"])\n        elif message.get("text") is not None:\n            writer.write(message["text"].encode())\n        await writer.drain()\n\n'''
+console_routes = '''\n\n@app.get("/vm/{name}/console", response_class=HTMLResponse)\ndef vm_console_page(request: Request, name: str):\n    auth_redirect = require_auth(request)\n    if auth_redirect:\n        return auth_redirect\n    if not valid_vm_name(name) or not vm_exists(name):\n        return RedirectResponse(url="/", status_code=303)\n    return templates.TemplateResponse("console.html", {"request": request, "app_name": APP_NAME, "user": AUTH_USER, "vm": vm_details(name), "console": console_info(name)})\n\n\n@app.websocket("/console/ws/{token}")\nasync def console_websocket(websocket: WebSocket, token: str):\n    if user_from_session_token(websocket.cookies.get("virtuality_session")) != AUTH_USER:\n        await websocket.close(code=1008)\n        return\n    try:\n        payload = console_serializer.loads(token)\n        vm_name = payload.get("vm")\n        target_port = int(payload.get("port"))\n    except Exception:\n        await websocket.close(code=1008)\n        return\n    if not valid_vm_name(vm_name) or not vm_exists(vm_name) or target_port < 5900 or target_port > 5999:\n        await websocket.close(code=1008)\n        return\n    await websocket.accept()\n    try:\n        reader, writer = await asyncio.open_connection("127.0.0.1", target_port)\n    except Exception:\n        await websocket.close(code=1011)\n        return\n''' + safe_gather.replace('\n', '\n')
-text = text.replace('@app.get("/login", response_class=HTMLResponse)\n', proxy_funcs + '\n@app.get("/login", response_class=HTMLResponse)\n', 1)
 
-console_routes = '''\n\n@app.get("/vm/{name}/console", response_class=HTMLResponse)\ndef vm_console_page(request: Request, name: str):\n    auth_redirect = require_auth(request)\n    if auth_redirect:\n        return auth_redirect\n    if not valid_vm_name(name) or not vm_exists(name):\n        return RedirectResponse(url="/", status_code=303)\n    return templates.TemplateResponse("console.html", {"request": request, "app_name": APP_NAME, "user": AUTH_USER, "vm": vm_details(name), "console": console_info(name)})\n\n\n@app.websocket("/console/ws/{token}")\nasync def console_websocket(websocket: WebSocket, token: str):\n    if user_from_session_token(websocket.cookies.get("virtuality_session")) != AUTH_USER:\n        await websocket.close(code=1008)\n        return\n    try:\n        payload = console_serializer.loads(token)\n        vm_name = payload.get("vm")\n        target_port = int(payload.get("port"))\n    except Exception:\n        await websocket.close(code=1008)\n        return\n    if not valid_vm_name(vm_name) or not vm_exists(vm_name) or target_port < 5900 or target_port > 5999:\n        await websocket.close(code=1008)\n        return\n    await websocket.accept()\n    try:\n        reader, writer = await asyncio.open_connection("127.0.0.1", target_port)\n    except Exception:\n        await websocket.close(code=1011)\n        return\n    try:\n        await asyncio.gather(proxy_vnc_to_websocket(reader, websocket), proxy_websocket_to_vnc(websocket, writer))\n    except (WebSocketDisconnect, asyncio.CancelledError, ConnectionError):\n        pass\n    finally:\n        writer.close()\n        await writer.wait_closed()\n'''
 text = text.replace('\n\n@app.get("/vm/{name}", response_class=HTMLResponse)\n', console_routes + '\n\n@app.get("/vm/{name}", response_class=HTMLResponse)\n', 1)
 
 app_path.write_text(text)