Обновление файла

backend/app/routers/console.py

@@ -1,19 +1,17 @@
 """WebSocket-прокси для VNC-консоли.
 
 Поток данных:
-  1. Backend получает от Proxmox одноразовый тикет через `vncproxy` (REST API,
+  1. Backend получает от Proxmox одноразовый тикет через `vncproxy` (REST API).
-     работает с API-токеном).
   2. Открывает WebSocket к Proxmox с этим тикетом и портом. Аутентификация:
      - API-токен (заголовок `Authorization: PVEAPIToken=...`) — новые PVE ≥ 7.x
      - Cookie PVEAuthCookie — если заданы PVE_USERNAME + PVE_PASSWORD
   3. Проксирует бинарный VNC-поток между браузером клиента и Proxmox.
 
-Особенности:
+Используется aiohttp вместо websockets, потому что:
-  - В Proxmox 8.x+ vncticket может содержать спецсимволы (#, /, %, :),
+  - Proxmox 8.x+ делает 302 Redirect с Location со схемой https:// вместо wss://,
-    которые ломают парсинг URL (RFC 3986). Кодируем через quote().
+    а websockets на uvloop не умеет правильно следовать за такими редиректами
-  - process_redirect в websockets не работает на uvloop (используется
+    (raise InvalidURI "scheme isn't ws or wss").
-    uvicorn[standard]), поэтому редиректы не обрабатываем — Proxmox
+  - aiohttp корректно обрабатывает HTTP→WS-редиректы с заменой scheme.
-    обычно их и не делает при корректном URL.
 """
 
 import asyncio
@@ -21,8 +19,9 @@ import logging
 from typing import Optional, Tuple
 from urllib.parse import quote
 
+import aiohttp
 import httpx
-import websockets
+from aiohttp import ClientSession, WSMsgType
 from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
 
 from ..config import settings
@@ -33,9 +32,10 @@ router = APIRouter(prefix="/console", tags=["console"])
 
 
 async def _get_pve_auth_cookie() -> Optional[Tuple[str, str]]:
-    """Аутентифицируется в Proxmox по логину/паролю и возвращает (cookie_header, csrf_token).
+    """Аутентифицируется в Proxmox по логину/паролю через REST API.
 
-    Возвращает None, если PVE_USERNAME/PVE_PASSWORD не заданы в настройках.
+    Возвращает кортеж (cookie_header, csrf_token) или None,
+    если PVE_USERNAME/PVE_PASSWORD не заданы.
     """
     if not settings.pve_username or not settings.pve_password:
         return None
@@ -49,15 +49,13 @@ async def _get_pve_auth_cookie() -> Optional[Tuple[str, str]]:
             resp = await client.post(url, data=payload)
             resp.raise_for_status()
             data = resp.json().get("data", {})
-            cookie = data.get("ticket")  # в Proxmox поле называется "ticket"
+            cookie = data.get("ticket")
             csrf = data.get("CSRFPreventionToken")
             if not cookie:
-                logger.error("auth: Proxmox не вернул PVEAuthCookie (пустой ответ)")
+                logger.error("auth: Proxmox не вернул PVEAuthCookie")
                 return None
-            # Для WebSocket-эндпоинта vncwebsocket достаточно PVEAuthCookie.
-            cookie_header = f"PVEAuthCookie={cookie}"
             logger.info("auth: PVEAuthCookie получен успешно")
-            return cookie_header, csrf or ""
+            return f"PVEAuthCookie={cookie}", csrf or ""
     except Exception as exc:
         logger.exception("auth: ошибка аутентификации в Proxmox")
         return None
@@ -77,8 +75,8 @@ async def console_ws(
 
     guest_path = "qemu" if guest_type == "vm" else "lxc"
     pve_host_only = settings.pve_host.replace("http://", "").replace("https://", "")
-    # quote() экранирует спецсимволы в тикете и порте по правилам RFC 3986.
+
-    # Без этого в Proxmox 8.x+ vncticket может содержать # / % : которые ломают URL.
+    # Экранируем спецсимволы в тикете/порте (vncticket может содержать # / % :).
     safe_ticket = quote(ticket, safe="")
     safe_port = quote(str(port), safe="")
     upstream_url = (
@@ -86,16 +84,18 @@ async def console_ws(
         f"?port={safe_port}&vncticket={safe_ticket}"
     )
 
-    # Определяем способ аутентификации.
+    # Определяем способ аутентификации: cookie приоритетнее (некоторые PVE
-    # Cookie-вариант приоритетнее, т.к. старые/строгие PVE не принимают API-токен на WS.
+    # не принимают API-токен на WebSocket).
-    headers: list = []
+    headers: dict = {}
     auth = await _get_pve_auth_cookie()
     if auth:
         cookie_header, _ = auth
-        headers.append(("Cookie", cookie_header))
+        headers["Cookie"] = cookie_header
         auth_kind = "cookie"
     else:
-        headers.append(("Authorization", f"PVEAPIToken={settings.pve_token_name}={settings.pve_token_value}"))
+        headers["Authorization"] = (
+            f"PVEAPIToken={settings.pve_token_name}={settings.pve_token_value}"
+        )
         auth_kind = "api-token"
 
     logger.info(
@@ -103,48 +103,66 @@ async def console_ws(
         node, vmid, guest_type, port, auth_kind,
     )
 
+    upstream: Optional[aiohttp.ClientWebSocketResponse] = None
+    session: Optional[ClientSession] = None
     try:
-        # extra_headers (а не additional_headers) — последний не поддерживается
+        # aiohttp сам следует за HTTP-редиректами и при этом корректно
-        # в uvloop, который использует uvicorn[standard].
+        # меняет схему с https на wss (что websockets на uvloop не делает).
-        async with websockets.connect(
+        session = ClientSession()
+        upstream = await session.ws_connect(
             upstream_url,
-            extra_headers=headers,
+            headers=headers,
-            subprotocols=["binary"],
+            ssl=False if not settings.pve_verify_ssl else None,
-            ssl=None if settings.pve_verify_ssl else False,
+            autoclose=False,
-        ) as upstream:
+            autoping=True,
-            logger.info("console_ws: upstream connected")
+            max_msg_size=8 * 1024 * 1024,
-
-            async def client_to_upstream():
-                try:
-                    while True:
-                        data = await websocket.receive_bytes()
-                        await upstream.send(data)
-                except WebSocketDisconnect:
-                    logger.info("console_ws: client disconnected")
-                except Exception as exc:
-                    logger.warning("console_ws: client_to_upstream error: %s", exc)
-
-            async def upstream_to_client():
-                try:
-                    async for message in upstream:
-                        if isinstance(message, str):
-                            message = message.encode()
-                        await websocket.send_bytes(message)
-                except websockets.ConnectionClosed as exc:
-                    logger.info("console_ws: upstream closed: code=%s", exc.code)
-                except Exception as exc:
-                    logger.warning("console_ws: upstream_to_client error: %s", exc)
-
-            await asyncio.gather(client_to_upstream(), upstream_to_client())
-    except websockets.exceptions.InvalidStatus as exc:
-        status = exc.response.status_code if exc.response else "?"
-        logger.error(
-            "console_ws: Proxmox отклонил подключение (status=%s, auth=%s). "
-            "Если 401 — задайте PVE_USERNAME/PVE_PASSWORD в backend/.env "
-            "для cookie-аутентификации.",
-            status, auth_kind,
         )
-        await websocket.close(code=1011, reason=f"Proxmox auth failed: {status}")
+        logger.info("console_ws: upstream connected")
+
+        async def client_to_upstream():
+            try:
+                while True:
+                    data = await websocket.receive_bytes()
+                    await upstream.send_bytes(data)
+            except WebSocketDisconnect:
+                logger.info("console_ws: client disconnected")
+            except Exception as exc:
+                logger.warning("console_ws: client_to_upstream error: %s", exc)
+
+        async def upstream_to_client():
+            try:
+                async for msg in upstream:
+                    if msg.type == WSMsgType.BINARY:
+                        await websocket.send_bytes(msg.data)
+                    elif msg.type == WSMsgType.TEXT:
+                        await websocket.send_bytes(msg.data.encode())
+                    elif msg.type == WSMsgType.CLOSE:
+                        logger.info("console_ws: upstream close: %s", msg.data)
+                        break
+                    elif msg.type == WSMsgType.ERROR:
+                        logger.warning("console_ws: upstream error: %s", msg.data)
+                        break
+            except Exception as exc:
+                logger.warning("console_ws: upstream_to_client error: %s", exc)
+
+        await asyncio.gather(client_to_upstream(), upstream_to_client())
+    except aiohttp.ClientResponseError as exc:
+        logger.error(
+            "console_ws: Proxmox ответил status=%s, auth=%s, message=%s",
+            exc.status, auth_kind, exc.message,
+        )
+        await websocket.close(code=1011, reason=f"Proxmox {exc.status}")
     except Exception as exc:  # noqa: BLE001
         logger.exception("console_ws: ошибка подключения к Proxmox")
         await websocket.close(code=1011, reason=str(exc)[:120])
+    finally:
+        if upstream is not None:
+            try:
+                await upstream.close()
+            except Exception:
+                pass
+        if session is not None:
+            try:
+                await session.close()
+            except Exception:
+                pass